"""Transformer building blocks shared by the encoder and decoder.

Includes the TokenGS-flavored details (design §2.4): QK-norm, LayerScale
(init 1e-5), pre-norm blocks, shared K/V cross-attention (the image-token K/V is
projected once and reused across all decoder layers), and boolean attention
masks for the dynamic->static causal rule.

Boolean ``attn_mask`` follows the ``F.scaled_dot_product_attention`` convention:
``True`` = the (query, key) pair participates; ``False`` = masked out.
"""

from __future__ import annotations

import math
from typing import Optional, Tuple

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


def fourier_encode(x: torch.Tensor, n_freq: int = 10, include_input: bool = True) -> torch.Tensor:
    """Sinusoidal positional encoding gamma(x) of ``[..., D]`` -> ``[..., D*(2*n_freq[+1])]``."""
    freqs = 2.0 ** torch.arange(n_freq, device=x.device, dtype=x.dtype) * math.pi
    xb = x[..., None] * freqs  # [..., D, n_freq]
    enc = torch.cat([xb.sin(), xb.cos()], dim=-1).flatten(-2)
    if include_input:
        enc = torch.cat([x, enc], dim=-1)
    return enc


class PatchEmbed(nn.Module):
    """Conv patchifier: ``[B, in_ch, H, W]`` -> ``[B, (H/p)(W/p), dim]``."""

    def __init__(self, in_ch: int, dim: int, patch: int):
        super().__init__()
        self.patch = patch
        self.proj = nn.Conv2d(in_ch, dim, kernel_size=patch, stride=patch)

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, Tuple[int, int]]:
        x = self.proj(x)
        gh, gw = x.shape[-2], x.shape[-1]
        x = x.flatten(2).transpose(1, 2)  # [B, N, dim]
        return x, (gh, gw)


class Mlp(nn.Module):
    def __init__(self, dim: int, ratio: float = 4.0):
        super().__init__()
        hidden = int(dim * ratio)
        self.fc1 = nn.Linear(dim, hidden)
        self.act = nn.GELU()
        self.fc2 = nn.Linear(hidden, dim)

    def forward(self, x):
        return self.fc2(self.act(self.fc1(x)))


class Attention(nn.Module):
    """Multi-head attention supporting self- and cross-attention with QK-norm.

    For cross-attention with shared K/V, pass precomputed ``kv`` (a tuple of
    ``[B, n_heads, M, head_dim]`` key/value tensors) and the q is taken from x.
    """

    def __init__(self, dim: int, n_heads: int, qk_norm: bool = True, kv_dim: Optional[int] = None):
        super().__init__()
        assert dim % n_heads == 0
        self.n_heads = n_heads
        self.head_dim = dim // n_heads
        self.scale = self.head_dim ** -0.5
        kv_dim = kv_dim or dim
        self.q_proj = nn.Linear(dim, dim)
        self.k_proj = nn.Linear(kv_dim, dim)
        self.v_proj = nn.Linear(kv_dim, dim)
        self.out_proj = nn.Linear(dim, dim)
        self.q_norm = nn.LayerNorm(self.head_dim) if qk_norm else nn.Identity()
        self.k_norm = nn.LayerNorm(self.head_dim) if qk_norm else nn.Identity()

    def _split(self, x: torch.Tensor) -> torch.Tensor:
        B, N, _ = x.shape
        return x.view(B, N, self.n_heads, self.head_dim).transpose(1, 2)

    def project_kv(self, context: torch.Tensor):
        """Project external context -> (k, v) heads, for caching across layers."""
        k = self.k_norm(self._split(self.k_proj(context)))
        v = self._split(self.v_proj(context))
        return k, v

    def forward(
        self,
        x: torch.Tensor,
        context: Optional[torch.Tensor] = None,
        kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
        attn_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        B, N, _ = x.shape
        q = self.q_norm(self._split(self.q_proj(x)))
        if kv is not None:
            k, v = kv
        else:
            src = context if context is not None else x
            k = self.k_norm(self._split(self.k_proj(src)))
            v = self._split(self.v_proj(src))
        out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
        out = out.transpose(1, 2).reshape(B, N, -1)
        return self.out_proj(out)


class SharedImageKV(nn.Module):
    """Project image tokens to (K, V) heads **once**, reused by every decoder
    cross-attention layer (TokenGS optimization: O(N_I*D)->O(N_I))."""

    def __init__(self, dim: int, n_heads: int, ctx_dim: int, qk_norm: bool = True):
        super().__init__()
        self.n_heads = n_heads
        self.head_dim = dim // n_heads
        self.k_proj = nn.Linear(ctx_dim, dim)
        self.v_proj = nn.Linear(ctx_dim, dim)
        self.k_norm = nn.LayerNorm(self.head_dim) if qk_norm else nn.Identity()

    def forward(self, context: torch.Tensor):
        B, M, _ = context.shape
        k = self.k_proj(context).view(B, M, self.n_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(context).view(B, M, self.n_heads, self.head_dim).transpose(1, 2)
        return self.k_norm(k), v


class CrossAttention(nn.Module):
    """Query-only cross-attention; consumes externally projected (K, V)."""

    def __init__(self, dim: int, n_heads: int, qk_norm: bool = True):
        super().__init__()
        self.n_heads = n_heads
        self.head_dim = dim // n_heads
        self.q_proj = nn.Linear(dim, dim)
        self.q_norm = nn.LayerNorm(self.head_dim) if qk_norm else nn.Identity()
        self.out_proj = nn.Linear(dim, dim)

    def forward(self, x, kv):
        B, N, _ = x.shape
        q = self.q_proj(x).view(B, N, self.n_heads, self.head_dim).transpose(1, 2)
        q = self.q_norm(q)
        k, v = kv
        out = F.scaled_dot_product_attention(q, k, v)
        out = out.transpose(1, 2).reshape(B, N, -1)
        return self.out_proj(out)


class LayerScale(nn.Module):
    def __init__(self, dim: int, init: float = 1e-5):
        super().__init__()
        self.gamma = nn.Parameter(init * torch.ones(dim))

    def forward(self, x):
        return x * self.gamma


class EncoderBlock(nn.Module):
    """Pre-norm self-attention block with QK-norm + LayerScale."""

    def __init__(self, dim, n_heads, mlp_ratio=4.0, qk_norm=True, layerscale_init=1e-5):
        super().__init__()
        self.norm1 = nn.LayerNorm(dim)
        self.attn = Attention(dim, n_heads, qk_norm=qk_norm)
        self.ls1 = LayerScale(dim, layerscale_init)
        self.norm2 = nn.LayerNorm(dim)
        self.mlp = Mlp(dim, mlp_ratio)
        self.ls2 = LayerScale(dim, layerscale_init)

    def forward(self, x, attn_mask=None):
        x = x + self.ls1(self.attn(self.norm1(x), attn_mask=attn_mask))
        x = x + self.ls2(self.mlp(self.norm2(x)))
        return x


class DecoderBlock(nn.Module):
    """Pre-norm DETR block: (cross-attn to image tokens with shared K/V) ->
    (self-attn among GS tokens with optional dynamic->static causal mask) -> MLP."""

    def __init__(self, dim, n_heads, mlp_ratio=4.0, qk_norm=True, layerscale_init=1e-5):
        super().__init__()
        self.norm_cross = nn.LayerNorm(dim)
        self.cross_attn = CrossAttention(dim, n_heads, qk_norm=qk_norm)
        self.ls_cross = LayerScale(dim, layerscale_init)
        self.norm_self = nn.LayerNorm(dim)
        self.self_attn = Attention(dim, n_heads, qk_norm=qk_norm)
        self.ls_self = LayerScale(dim, layerscale_init)
        self.norm_mlp = nn.LayerNorm(dim)
        self.mlp = Mlp(dim, mlp_ratio)
        self.ls_mlp = LayerScale(dim, layerscale_init)

    def forward(self, x, kv, self_mask=None):
        x = x + self.ls_cross(self.cross_attn(self.norm_cross(x), kv=kv))
        x = x + self.ls_self(self.self_attn(self.norm_self(x), attn_mask=self_mask))
        x = x + self.ls_mlp(self.mlp(self.norm_mlp(x)))
        return x