"""Building blocks: RMSNorm, RoPE, grouped-query causal attention with optional KV cache, SwiGLU."""
from __future__ import annotations

import math
from typing import Optional

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


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

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        dtype = x.dtype
        xf = x.float()
        out = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps)
        return (out * self.weight.float()).to(dtype)


def rope_cos_sin(positions: torch.Tensor, head_dim: int, theta: float, device, dtype=torch.float32):
    """positions: [..., T] integer tensor -> cos, sin of shape [..., T, head_dim]."""
    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device, dtype=torch.float32) / head_dim))
    freqs = positions.to(torch.float32).unsqueeze(-1) * inv_freq  # [..., T, hd/2]
    emb = torch.cat([freqs, freqs], dim=-1)
    return emb.cos().to(dtype), emb.sin().to(dtype)


def rotate_half(x: torch.Tensor) -> torch.Tensor:
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat([-x2, x1], dim=-1)


def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    """x: [B, H, T, hd]; cos/sin: [B or 1, T, hd] -> broadcast over heads."""
    cos = cos.unsqueeze(1)
    sin = sin.unsqueeze(1)
    return x * cos + rotate_half(x) * sin


class SwiGLU(nn.Module):
    def __init__(self, dim: int, hidden: int):
        super().__init__()
        self.w_gate = nn.Linear(dim, hidden, bias=False)
        self.w_up = nn.Linear(dim, hidden, bias=False)
        self.w_down = nn.Linear(hidden, dim, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))


class LayerKVCache:
    """Per-layer KV cache. Tensors are [B, Hkv, T, hd]. One instance per request batch."""

    def __init__(self):
        self.k: Optional[torch.Tensor] = None
        self.v: Optional[torch.Tensor] = None

    @property
    def length(self) -> int:
        return 0 if self.k is None else self.k.shape[2]

    def append(self, k: torch.Tensor, v: torch.Tensor):
        if self.k is None:
            self.k, self.v = k, v
        else:
            self.k = torch.cat([self.k, k], dim=2)
            self.v = torch.cat([self.v, v], dim=2)
        return self.k, self.v

    def reset(self):
        self.k = None
        self.v = None


class CausalSelfAttention(nn.Module):
    """GQA causal attention with RoPE. Supports full-sequence (training) and cached incremental decoding."""

    def __init__(self, d_model: int, n_heads: int, n_kv_heads: int, rope_theta: float):
        super().__init__()
        self.n_heads = n_heads
        self.n_kv_heads = n_kv_heads
        self.head_dim = d_model // n_heads
        self.rope_theta = rope_theta
        self.wq = nn.Linear(d_model, n_heads * self.head_dim, bias=False)
        self.wk = nn.Linear(d_model, n_kv_heads * self.head_dim, bias=False)
        self.wv = nn.Linear(d_model, n_kv_heads * self.head_dim, bias=False)
        self.wo = nn.Linear(n_heads * self.head_dim, d_model, bias=False)

    def forward(self, x: torch.Tensor, positions: torch.Tensor, cache: Optional[LayerKVCache] = None,
                attn_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
        """x: [B, T, d]; positions: [B, T] absolute positions (for RoPE and causal masking).

        If cache is given, k/v are appended and attention runs over the full cached length.
        attn_mask: optional boolean [B, 1, T, T_total] (True = attend). If None and cache is None,
        a standard causal mask is used.
        """
        B, T, _ = x.shape
        q = self.wq(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        k = self.wk(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
        v = self.wv(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
        cos, sin = rope_cos_sin(positions, self.head_dim, self.rope_theta, x.device, x.dtype)
        q = apply_rope(q, cos, sin)
        k = apply_rope(k, cos, sin)
        if cache is not None:
            k, v = cache.append(k, v)
        rep = self.n_heads // self.n_kv_heads
        if rep > 1:
            k = k.repeat_interleave(rep, dim=1)
            v = v.repeat_interleave(rep, dim=1)
        if attn_mask is None and cache is None:
            out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        else:
            if attn_mask is None:
                # incremental decoding: query positions attend to all cached keys at or before them.
                T_total = k.shape[2]
                key_pos = torch.arange(T_total, device=x.device)
                # cached keys are stored in insertion order == absolute position order
                attn_mask = (key_pos[None, None, None, :] <= positions[:, None, :, None])
            out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
        out = out.transpose(1, 2).reshape(B, T, self.n_heads * self.head_dim)
        return self.wo(out)