# ─────────────────────────────────────────────────────────
# EGen Labs | ErebusTN
# AXIOM-1 — EGen-Core/AXIOM-1
# File: axiom1/model/components.py
# ─────────────────────────────────────────────────────────

"""
Core components for AXIOM-1.

All components match the canonical specification in basis/AXIOM1_Context_File.md.
Locked choices — no substitutions allowed:
  - RMSNorm (not LayerNorm)
  - RoPE (not learned absolute)
  - GQA 8q/2kv (not MHA/MQA)
  - SwiGLU with no bias (not ReLU/GELU)
  - DropPath (not standard dropout)
"""

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


class RMSNorm(nn.Module):
    """
    Root Mean Square Layer Normalization.

    Used everywhere in AXIOM-1 instead of LayerNorm.
    Reason: no mean-centering step → no activation drift across loop iterations.
    """

    def __init__(self, d: int, eps: float = 1e-6) -> None:
        super().__init__()
        self.weight = nn.Parameter(torch.ones(d))
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        rms = x.pow(2).mean(-1, keepdim=True).add(self.eps).sqrt()
        return (x / rms) * self.weight


def apply_rope(
    q: torch.Tensor,
    k: torch.Tensor,
    seq_len: int,
    head_dim: int,
    base: int = 10000
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Rotary Position Embedding applied to q and k.

    q, k shape: (batch, n_heads, seq_len, head_dim)

    RoPE is applied INSIDE the attention, not to the embedding.
    This means position info is refreshed at every loop iteration — correct behavior.
    """
    device = q.device
    dtype = q.dtype

    theta = 1.0 / (
        base ** (torch.arange(0, head_dim, 2, dtype=torch.float32, device=device) / head_dim)
    )
    positions = torch.arange(seq_len, dtype=torch.float32, device=device)
    freqs = torch.outer(positions, theta)           # (T, head_dim/2)
    cos = freqs.cos()[None, None, :, :]             # (1, 1, T, head_dim/2)
    sin = freqs.sin()[None, None, :, :]

    def rotate_half(x: torch.Tensor) -> torch.Tensor:
        x_even = x[..., ::2]
        x_odd = x[..., 1::2]
        return torch.stack([-x_odd, x_even], dim=-1).flatten(-2)

    # Broadcast cos/sin across even and odd positions
    cos_full = torch.stack([cos, cos], dim=-1).flatten(-2)
    sin_full = torch.stack([sin, sin], dim=-1).flatten(-2)

    q_rot = (q.float() * cos_full + rotate_half(q.float()) * sin_full).to(dtype)
    k_rot = (k.float() * cos_full + rotate_half(k.float()) * sin_full).to(dtype)
    return q_rot, k_rot


class GroupedQueryAttention(nn.Module):
    """
    GQA with 8 query heads and 2 KV heads (4× repetition).

    Uses Flash Attention (F.scaled_dot_product_attention) for efficiency.
    RoPE is applied inside this module.
    No dropout — DropPath handles regularization at the block level.
    """

    def __init__(self, d_model: int, n_q_heads: int, n_kv_heads: int) -> None:
        super().__init__()
        assert n_q_heads % n_kv_heads == 0, \
            f"n_q_heads ({n_q_heads}) must be divisible by n_kv_heads ({n_kv_heads})"

        self.n_q = n_q_heads
        self.n_kv = n_kv_heads
        self.n_rep = n_q_heads // n_kv_heads
        self.hd = d_model // n_q_heads

        self.q_proj = nn.Linear(d_model, n_q_heads * self.hd, bias=False)
        self.k_proj = nn.Linear(d_model, n_kv_heads * self.hd, bias=False)
        self.v_proj = nn.Linear(d_model, n_kv_heads * self.hd, bias=False)
        self.o_proj = nn.Linear(n_q_heads * self.hd, d_model, bias=False)

    def forward(
        self,
        x: torch.Tensor,
        mask: torch.Tensor | None = None,
        kv_cache: object | None = None,
        step_idx: int | None = None
    ) -> torch.Tensor:
        B, T, _ = x.shape

        q = self.q_proj(x).view(B, T, self.n_q, self.hd).transpose(1, 2)
        k = self.k_proj(x).view(B, T, self.n_kv, self.hd).transpose(1, 2)
        v = self.v_proj(x).view(B, T, self.n_kv, self.hd).transpose(1, 2)

        # Apply RoPE
        q, k = apply_rope(q, k, T, self.hd)

        # KV cache update (inference only)
        if kv_cache is not None and step_idx is not None:
            k, v = kv_cache.update(step_idx, k, v)

        # Expand KV to match Q heads
        k = k.repeat_interleave(self.n_rep, dim=1)
        v = v.repeat_interleave(self.n_rep, dim=1)

        # Flash Attention — causal by default during training
        if mask is None:
            out = F.scaled_dot_product_attention(q, k, v, is_causal=True, dropout_p=0.0)
        else:
            out = F.scaled_dot_product_attention(q, k, v, attn_mask=~mask, dropout_p=0.0)

        out = out.transpose(1, 2).contiguous().view(B, T, -1)
        return self.o_proj(out)


class SwiGLU_FFN(nn.Module):
    """
    SwiGLU feed-forward network. No bias anywhere.

    d_ff = int(d_model × 8/3) to match 4× param count of standard FFN.
    No bias: prevents activation mean drift across loop iterations.
    """

    def __init__(self, d_model: int, d_ff: int) -> None:
        super().__init__()
        d_inner = int(d_ff * 2 / 3)    # SwiGLU splits d_ff into two parallel paths
        self.gate = nn.Linear(d_model, d_inner, bias=False)
        self.up = nn.Linear(d_model, d_inner, bias=False)
        self.down = nn.Linear(d_inner, d_model, bias=False)

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


class DropPath(nn.Module):
    """
    Stochastic depth: randomly skip entire sub-layer outputs during training.

    Applied at block level, not token level.
    Rate 0.10 for AXIOM-1.
    This is the ONLY dropout-style regularization in AXIOM-1.
    Standard dropout is NOT used inside the loop.
    """

    def __init__(self, p: float = 0.10) -> None:
        super().__init__()
        self.p = p

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if not self.training or self.p == 0.0:
            return x
        keep_prob = 1.0 - self.p
        shape = (x.shape[0],) + (1,) * (x.ndim - 1)
        mask = torch.bernoulli(torch.full(shape, keep_prob, device=x.device))
        return x * mask / keep_prob
