"""Attention for Ethoric: RMSNorm, RoPE, grouped-query attention with QK-norm,
a per-head output gate, and an optional local window on selected layers.

Two backends compute the same thing: flex (FlexAttention with a compiled block
mask, CUDA) and dense (an explicit boolean mask handed to SDPA, used on CPU and
MPS). They agree to 1e-5 in fp32; if you change one, change both.

Works on Transformers 4.x and 5.x; the v5 differences are marked inline.
"""
from __future__ import annotations

import math
from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F

# ----------------------------------------------------------------------------- rope

# A plain module-level dict, not a buffer. Transformers 5 builds models on the
# meta device and fills only tensors the weights file contains, so a
# non-persistent buffer would come back as uninitialised memory. Keyed by
# device and dtype, so .to()/.cuda() need not move anything.
_ROPE_CACHE: dict[tuple, tuple[torch.Tensor, torch.Tensor]] = {}


def rope_tables(seq_len: int, head_dim: int, theta: float, device, dtype):
    """cos/sin of shape [seq_len, head_dim/2], built once per (shape, device)."""
    key = (int(seq_len), int(head_dim), float(theta), str(device), str(dtype))
    hit = _ROPE_CACHE.get(key)
    if hit is not None:
        return hit
    half = head_dim // 2
    freqs = 1.0 / (theta ** (torch.arange(0, half, device=device, dtype=torch.float32) / half))
    t = torch.arange(seq_len, device=device, dtype=torch.float32)
    angles = torch.outer(t, freqs)                     # [T, half]
    cos, sin = angles.cos().to(dtype), angles.sin().to(dtype)
    _ROPE_CACHE[key] = (cos, sin)
    return cos, sin


def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    """Half-split rotation (first half against second), not the interleaved
    variant. The two are not interchangeable -- swapping them degrades output
    without raising anything.

    x is [B, H, T, hd]. cos/sin are [T, hd/2] when every row of the batch shares
    the same positions, or [B, T, hd/2] when they do not (left-padded batched
    generation, where Transformers derives position_ids from the attention
    mask). Both broadcast against the head axis.
    """
    x1, x2 = x.chunk(2, dim=-1)
    if cos.dim() == 2:
        cos = cos.view(1, 1, cos.shape[0], cos.shape[1])
        sin = sin.view(1, 1, sin.shape[0], sin.shape[1])
    else:
        cos = cos.unsqueeze(1)
        sin = sin.unsqueeze(1)
    return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1)


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

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Always reduce in fp32: in fp16 the mean of squares underflows on the
        # long tail and the norm comes back slightly wrong.
        dtype = x.dtype
        x32 = x.float()
        out = x32 * torch.rsqrt(x32.pow(2).mean(-1, keepdim=True) + self.eps)
        return (out * self.weight.float()).to(dtype)


# ----------------------------------------------------------------------------- masks

def dense_mask(
    q_pos: torch.Tensor,            # [B, Tq]
    k_pos: torch.Tensor,            # [B, Tk]
    q_doc: Optional[torch.Tensor],  # [B, Tq] or None
    k_doc: Optional[torch.Tensor],  # [B, Tk] or None
    window: Optional[int] = None,
) -> torch.Tensor:
    """Boolean [B, 1, Tq, Tk]; True = attend.

    Three predicates, ANDed: causal (key no later than query), same document
    when doc ids are supplied, and within `window` positions when the layer is
    windowed. Positions, not indices -- with a KV cache the key at index 0 is
    not at position 0.
    """
    allowed = k_pos[:, None, :] <= q_pos[:, :, None]
    if q_doc is not None and k_doc is not None:
        allowed &= q_doc[:, :, None] == k_doc[:, None, :]
    if window is not None:
        allowed &= (q_pos[:, :, None] - k_pos[:, None, :]) < int(window)
    return allowed.unsqueeze(1)


def build_flex_block_mask(
    q_pos: torch.Tensor,
    k_pos: torch.Tensor,
    q_doc: Optional[torch.Tensor],
    k_doc: Optional[torch.Tensor],
    window: Optional[int],
    n_head: int,
):
    """Compile a BlockMask for the same predicate as `dense_mask`.

    Rebuilt per batch because document ids change; the compiled kernels are
    reused across calls by torch.compile's cache. Keep this in step with
    `dense_mask` -- the two backends are expected to agree numerically, and a
    mismatch shows up as a quiet quality drop rather than an error.
    """
    from torch.nn.attention.flex_attention import create_block_mask

    B, Tq = q_pos.shape
    Tk = k_pos.shape[1]

    def mask_mod(b, h, q_idx, kv_idx):
        allowed = k_pos[b, kv_idx] <= q_pos[b, q_idx]
        if q_doc is not None and k_doc is not None:
            allowed = allowed & (q_doc[b, q_idx] == k_doc[b, kv_idx])
        if window is not None:
            allowed = allowed & ((q_pos[b, q_idx] - k_pos[b, kv_idx]) < window)
        return allowed

    return create_block_mask(mask_mod, B, n_head, Tq, Tk, device=q_pos.device)


# ----------------------------------------------------------------------------- attention

class EthoricAttention(nn.Module):
    """Grouped-query attention: n_head queries over n_kv_head key/value heads.

    Two details are easy to trip over. QK-norm is applied before RoPE, not
    after. And the output gate is a sigmoid initialised near-open rather than at
    zero, so a freshly built model already passes attention through -- see
    `_apply_residual_init` in modeling_ethoric.py.
    """

    def __init__(self, config, layer_idx: int):
        super().__init__()
        self.layer_idx = int(layer_idx)
        self.n_head = int(config.n_head)
        self.n_kv_head = int(config.n_kv_head)
        self.head_dim = int(config.head_dim)
        self.d_model = int(config.d_model)
        self.rope_theta = float(config.rope_theta)
        self.max_seq_len = int(config.max_seq_len)
        self.scale = 1.0 / math.sqrt(self.head_dim)
        self.attn_gate = str(config.attn_gate)
        self.backend = "dense"                    # set by the model

        self.q_proj = nn.Linear(self.d_model, self.n_head * self.head_dim, bias=False)
        self.k_proj = nn.Linear(self.d_model, self.n_kv_head * self.head_dim, bias=False)
        self.v_proj = nn.Linear(self.d_model, self.n_kv_head * self.head_dim, bias=False)
        self.o_proj = nn.Linear(self.n_head * self.head_dim, self.d_model, bias=False)

        if str(config.qk_norm) == "per_head":
            self.q_norm = RMSNorm(self.head_dim, config.norm_eps)
            self.k_norm = RMSNorm(self.head_dim, config.norm_eps)
        else:
            self.q_norm = self.k_norm = None

        if self.attn_gate == "headwise":
            self.gate_proj = nn.Linear(self.d_model, self.n_head, bias=True)
        elif self.attn_gate == "elementwise":
            self.gate_proj = nn.Linear(self.d_model, self.n_head * self.head_dim, bias=True)
        else:
            self.gate_proj = None

        # Layers listed in `attn_window_layers` attend only to the most recent
        # `attn_window` positions; every other layer sees the full context. This
        # is part of the trained architecture, not a tuning knob: dropping the
        # window from a model trained with it degrades output badly, and nothing
        # about the weights file will tell you it is missing.
        window_layers = list(getattr(config, "attn_window_layers", None) or [])
        self.attn_window = (
            int(getattr(config, "attn_window", 0) or 0)
            if self.layer_idx in window_layers
            else None
        )

    def forward(
        self,
        h: torch.Tensor,                       # normalized input [B, T, d]
        *,
        positions: torch.Tensor,               # [B, T]
        doc_ids: Optional[torch.Tensor],       # [B, T] or None
        layer_cache: Optional[dict] = None,
        block_mask=None,
    ) -> torch.Tensor:
        B, T, _ = h.shape
        H, hd = self.n_head, self.head_dim

        q = self.q_proj(h).view(B, T, H, hd).transpose(1, 2)              # [B,H,T,hd]
        k = self.k_proj(h).view(B, T, self.n_kv_head, hd).transpose(1, 2)  # [B,Hkv,T,hd]
        v = self.v_proj(h).view(B, T, self.n_kv_head, hd).transpose(1, 2)

        if self.q_norm is not None:
            q, k = self.q_norm(q), self.k_norm(k)

        # Build the table at max_seq_len and index into it. Deriving the length
        # from positions.max().item() would force a device sync in every layer
        # of every step, which dominates the step time at this model size.
        cos, sin = rope_tables(self.max_seq_len, hd, self.rope_theta, h.device, q.dtype)
        # Gather per batch row. Rows are identical whenever the model builds
        # positions itself, but `positions` is a public forward kwarg: under
        # left-padded batched generation the rows differ, and indexing with
        # positions[0] would silently rotate every row by row 0's offsets.
        # Checking for identical rows would cost a device sync per layer, which
        # is more than the gather.
        cos_t = cos[positions] if positions.shape[0] else cos[:T]
        sin_t = sin[positions] if positions.shape[0] else sin[:T]
        q = apply_rope(q, cos_t, sin_t)
        k = apply_rope(k, cos_t, sin_t)

        # The cache holds positions and doc ids alongside k/v because the mask
        # predicate needs them for the whole prefix, not just the new tokens.
        k_pos, k_doc, k_all, v_all = positions, doc_ids, k, v
        if layer_cache is not None:
            if layer_cache.get("k") is not None:
                k_all = torch.cat([layer_cache["k"], k], dim=2)
                v_all = torch.cat([layer_cache["v"], v], dim=2)
                k_pos = torch.cat([layer_cache["positions"], positions], dim=1)
                if doc_ids is not None and layer_cache.get("doc_ids") is not None:
                    k_doc = torch.cat([layer_cache["doc_ids"], doc_ids], dim=1)
            layer_cache["k"], layer_cache["v"] = k_all, v_all
            layer_cache["positions"] = k_pos
            layer_cache["doc_ids"] = k_doc

        if self.n_kv_head != H:
            rep = H // self.n_kv_head
            k_all = k_all.repeat_interleave(rep, dim=1)
            v_all = v_all.repeat_interleave(rep, dim=1)

        window = self.attn_window

        if self.backend == "flex" and block_mask is None:
            raise RuntimeError("flex backend requires a precomputed block mask")

        if self.backend == "flex":
            from torch.nn.attention.flex_attention import flex_attention

            out = flex_attention(q, k_all, v_all, block_mask=block_mask, scale=self.scale)
        else:
            mask = dense_mask(positions, k_pos, doc_ids, k_doc, window)
            out = F.scaled_dot_product_attention(
                q, k_all, v_all, attn_mask=mask, scale=self.scale
            )

        # Output gate, applied after attention and before o_proj: one scalar
        # per head (headwise) or per channel (elementwise).
        if self.gate_proj is not None:
            gate = torch.sigmoid(self.gate_proj(h))
            if self.attn_gate == "headwise":
                gate = gate.permute(0, 2, 1).unsqueeze(-1)                 # [B,H,T,1]
            else:
                gate = gate.view(B, T, H, hd).permute(0, 2, 1, 3)
            out = out * gate.to(out.dtype)

        out = out.transpose(1, 2).reshape(B, T, H * hd)
        return self.o_proj(out)