"""Attention layers used by the correctness-first scaffold."""

from __future__ import annotations

import math

import torch
from torch import nn

from .padding import masked_hidden


class RMSNorm(nn.Module):
    def __init__(self, hidden_size: int, eps: float = 1e-5):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        variance = x.pow(2).mean(dim=-1, keepdim=True)
        return x * torch.rsqrt(variance + self.eps) * self.weight


class FeedForward(nn.Module):
    def __init__(self, hidden_size: int, intermediate_size: int, dropout: float, activation: str = "gelu"):
        super().__init__()
        self.activation_name = activation
        if activation == "geglu":
            self.in_proj = nn.Linear(hidden_size, intermediate_size * 2)
            self.out_proj = nn.Linear(intermediate_size, hidden_size)
        else:
            self.in_proj = nn.Linear(hidden_size, intermediate_size)
            self.out_proj = nn.Linear(intermediate_size, hidden_size)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.in_proj(x)
        if self.activation_name == "geglu":
            value, gate = x.chunk(2, dim=-1)
            x = value * torch.nn.functional.gelu(gate)
        else:
            x = torch.nn.functional.gelu(x)
        return self.out_proj(self.dropout(x))


class StrataBertAttentionLayer(nn.Module):
    def __init__(self, config, layer_type: str):
        super().__init__()
        if config.hidden_size % config.num_attention_heads != 0:
            raise ValueError("hidden_size must be divisible by num_attention_heads")
        self.layer_type = layer_type
        self.num_heads = config.num_attention_heads
        self.head_dim = config.hidden_size // config.num_attention_heads
        self.window = config.local_attention_window
        self.norm = RMSNorm(config.hidden_size, config.norm_eps)
        self.qkv = nn.Linear(config.hidden_size, config.hidden_size * 3)
        self.out_proj = nn.Linear(config.hidden_size, config.hidden_size)
        self.dropout = nn.Dropout(config.attention_dropout)
        self.ffn_norm = RMSNorm(config.hidden_size, config.norm_eps)
        self.ffn = FeedForward(
            config.hidden_size,
            config.intermediate_size,
            config.hidden_dropout,
            config.hidden_activation,
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor,
        segment_ids: torch.Tensor | None = None,
    ) -> torch.Tensor:
        residual = hidden_states
        x = self.norm(hidden_states)
        batch, length, hidden = x.shape
        qkv = self.qkv(x).view(batch, length, 3, self.num_heads, self.head_dim)
        q, k, v = qkv.unbind(dim=2)
        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)

        scores = torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(self.head_dim)
        key_mask = attention_mask[:, None, None, :]
        query_mask = attention_mask[:, None, :, None]
        scores = scores.masked_fill(~key_mask, -1.0e4)
        if segment_ids is not None:
            same_segment = segment_ids[:, None, :, None] == segment_ids[:, None, None, :]
            scores = scores.masked_fill(~same_segment, -1.0e4)
        if self.layer_type == "local_attention":
            idx = torch.arange(length, device=x.device)
            local = (idx[None, :] - idx[:, None]).abs() <= self.window
            scores = scores.masked_fill(~local[None, None, :, :], -1.0e4)
        probs = torch.softmax(scores, dim=-1)
        probs = self.dropout(probs) * query_mask.to(probs.dtype)
        attended = torch.matmul(probs, v).transpose(1, 2).contiguous().view(batch, length, hidden)
        hidden_states = residual + self.out_proj(attended)
        hidden_states = masked_hidden(hidden_states, attention_mask)

        hidden_states = hidden_states + self.ffn(self.ffn_norm(hidden_states))
        return masked_hidden(hidden_states, attention_mask)