"""Small shared neural-network layers."""

from __future__ import annotations

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


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

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Accumulating the variance in fp32 prevents low-precision callers from
        # underflowing; the normalized activation is cast back before scaling.
        normalized = x.float() * torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + self.eps)
        return normalized.to(dtype=x.dtype) * self.weight


class SwiGLU(nn.Module):
    def __init__(self, width: int, hidden_width: int, dropout: float, bias: bool = False) -> None:
        super().__init__()
        self.gate = nn.Linear(width, hidden_width, bias=bias)
        self.up = nn.Linear(width, hidden_width, bias=bias)
        self.down = nn.Linear(hidden_width, width, bias=bias)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.dropout(self.down(F.silu(self.gate(x)) * self.up(x)))


def apply_rope(q: torch.Tensor, k: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """Apply rotary embeddings to tensors shaped ``[B, H, T, D]``."""

    _, _, seq_len, head_dim = q.shape
    if head_dim % 2:
        raise ValueError("RoPE requires an even head dimension")
    # RoPE pairs adjacent channels. Frequencies decrease geometrically from the
    # fastest pair to the slowest pair, then each pair is rotated by its token
    # position. There are intentionally no learned position parameters.
    positions = torch.arange(seq_len, device=q.device, dtype=torch.float32)
    frequencies = 1.0 / (
        10000.0
        ** (torch.arange(0, head_dim, 2, device=q.device, dtype=torch.float32) / head_dim)
    )
    angles = torch.outer(positions, frequencies)
    cos = angles.cos()[None, None, :, :].to(dtype=q.dtype)
    sin = angles.sin()[None, None, :, :].to(dtype=q.dtype)

    def rotate(x: torch.Tensor) -> torch.Tensor:
        # [x_even, x_odd] is multiplied by the 2-D rotation matrix for every
        # position/frequency pair, then interleaved back into the head axis.
        even, odd = x[..., 0::2], x[..., 1::2]
        return torch.stack((even * cos - odd * sin, even * sin + odd * cos), dim=-1).flatten(-2)

    return rotate(q), rotate(k)