"""The decoder's ViT trunk: RMS-sandwich blocks with 2-D RoPE."""

from __future__ import annotations

import torch
from torch import Tensor, nn
from torch.nn import functional as F

from .layers import Mlp, RMSNorm

ROPE_BASE = 10_000.0


def rope_periods(head_dim: int, device: torch.device) -> Tensor:
    """FP32 periods ``ROPE_BASE ** (2 i / (head_dim / 2))`` of the
    ``head_dim / 4`` rotation frequencies per axis."""

    exponents = (
        2.0
        * torch.arange(head_dim // 4, device=device, dtype=torch.float32)
        / (head_dim // 2)
    )
    return torch.tensor(ROPE_BASE, device=device, dtype=torch.float32) ** exponents


def rope_tables(
    height: int, width: int, *, head_dim: int, device: torch.device
) -> tuple[Tensor, Tensor]:
    """FP32 ``(sin, cos)`` ``[H * W, head_dim]`` of axial 2-D RoPE on the token
    grid: unnormalized patch-index coordinates (row, column), the
    :func:`rope_periods`, adjacent-pair layout."""

    periods = rope_periods(head_dim, device)
    rows = torch.arange(0.0, float(height), device=device, dtype=torch.float32)
    cols = torch.arange(0.0, float(width), device=device, dtype=torch.float32)
    coords = torch.stack(torch.meshgrid(rows, cols, indexing="ij"), dim=-1).flatten(
        0, 1
    )
    angles = 1.0 * coords[:, :, None] / periods[None, None, :].expand(1, 2, -1)
    angles = angles.repeat_interleave(2, dim=-1).flatten(1, 2)
    return torch.sin(angles), torch.cos(angles)


def _rotate_pairs(x: Tensor) -> Tensor:
    """Rotate consecutive channel pairs: ``(a, b) -> (-b, a)``."""

    pairs = x.reshape(*x.shape[:-1], x.shape[-1] // 2, 2)
    return torch.stack((-pairs[..., 1], pairs[..., 0]), dim=-1).reshape_as(x)


def apply_rope(x: Tensor, sin: Tensor, cos: Tensor) -> Tensor:
    """Rotate ``[B, heads, N, head_dim]`` queries or keys in the tables' dtype."""

    rotated = x.to(dtype=sin.dtype)
    sin_b = sin.view(1, 1, *sin.shape)
    cos_b = cos.view(1, 1, *cos.shape)
    return ((rotated * cos_b) + (_rotate_pairs(rotated) * sin_b)).to(dtype=x.dtype)


class Attention(nn.Module):
    """Bias-free multi-head self-attention with RMS-normalized queries and keys."""

    def __init__(self, width: int, head_dim: int) -> None:
        """Allocate the fused QKV and output projections and the q/k norms."""

        super().__init__()
        self.heads = width // head_dim
        self.head_dim = head_dim
        self.qkv = nn.Linear(width, 3 * width, bias=False)
        self.proj_out = nn.Linear(width, width, bias=False)
        self.q_norm = RMSNorm(head_dim, affine=True)
        self.k_norm = RMSNorm(head_dim, affine=True)

    def forward(self, x: Tensor, sin: Tensor, cos: Tensor) -> Tensor:
        """Attend over all tokens of ``[B, N, width]``."""

        b, n, width = x.shape
        q, k, v = self.qkv(x).chunk(3, dim=-1)
        shape = (b, n, self.heads, self.head_dim)
        q = q.view(shape).transpose(1, 2).contiguous()
        k = k.view(shape).transpose(1, 2).contiguous()
        v = v.view(shape).transpose(1, 2).contiguous()
        q = apply_rope(self.q_norm(q), sin, cos)
        k = apply_rope(self.k_norm(k), sin, cos)
        out = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0, is_causal=False)
        return self.proj_out(out.transpose(1, 2).contiguous().view(b, n, width))


class DitBlock(nn.Module):
    """ViT block: RMSNorm before and after attention and MLP."""

    def __init__(self, width: int, head_dim: int, mlp_ratio: float) -> None:
        """Allocate attention, MLP and the four norms."""

        super().__init__()
        self.attn_norm1 = RMSNorm(width, affine=True)
        self.attn_norm2 = RMSNorm(width, affine=True)
        self.mlp_norm1 = RMSNorm(width, affine=True)
        self.mlp_norm2 = RMSNorm(width, affine=True)
        self.attn = Attention(width, head_dim)
        self.mlp = Mlp(width, int(mlp_ratio * width))

    def forward(self, x: Tensor, sin: Tensor, cos: Tensor) -> Tensor:
        """Residual attention then residual MLP on ``[B, N, width]`` tokens."""

        x = x + self.attn_norm2(self.attn(self.attn_norm1(x), sin, cos))
        return x + self.mlp_norm2(self.mlp(self.mlp_norm1(x)))