"""
model.py — LLaMA-style causal transformer with IPA phonological feature
augmentation via gated fusion at the embedding layer.

Architecture
------------
  RMSNorm  (not LayerNorm)
  RoPE     positional embeddings (not learned absolute positions)
  SwiGLU   FFN (not GELU / ReLU)
  Standard causal multi-head self-attention  (no local window)
  No bias  in attention projections or FFN linear layers
  Weight tying  between token embedding and lm_head

IPA Fusion Variants  (ModelConfig.variant)
------------------------------------------
  baseline  — token embeddings only; IPA tensors are ignored
  ipa_add   — fused = embed + W_ipa · ipa_vec   (no gate)
  ipa_gate  — fused = embed + σ(W_gate · ipa_vec) ⊙ (W_ipa · ipa_vec)
  ipa_full  — same gated fusion + MSE auxiliary loss that predicts raw IPA
              feature vectors from the final hidden states.  During training,
              ipa_mask_ratio fraction of IPA inputs are zeroed before fusion,
              creating a reconstruction signal for the auxiliary head.

Forward signature
-----------------
  out = model(input_ids, ipa_vectors, attention_mask=None)
  out.logits    → FloatTensor [B, T, vocab_size]
  out.aux_loss  → scalar tensor; 0.0 for baseline / ipa_add / ipa_gate

Total training loss:
  loss = CE(out.logits, labels) + cfg.ipa_lambda * out.aux_loss
"""

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Optional

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

from config import ModelConfig


# ─────────────────────────────────────────────────────────────────────────────
# Output container
# ─────────────────────────────────────────────────────────────────────────────

@dataclass
class ModelOutput:
    logits:   torch.Tensor   # [B, T, vocab_size]
    aux_loss: torch.Tensor   # scalar; 0.0 for non-ipa_full variants


# ─────────────────────────────────────────────────────────────────────────────
# RMSNorm
# ─────────────────────────────────────────────────────────────────────────────

class RMSNorm(nn.Module):
    """Root Mean Square Layer Normalization (Zhang & Sennrich, 2019).

    RMSNorm(x) = x / RMS(x) * weight    where RMS(x) = sqrt(mean(x²) + ε)
    No learned bias — consistent with LLaMA.
    """

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

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # rsqrt for numerical stability and efficiency
        rms_inv = x.pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
        return self.weight * (x * rms_inv)


# ─────────────────────────────────────────────────────────────────────────────
# Rotary Position Embeddings (RoPE)
# ─────────────────────────────────────────────────────────────────────────────

def _precompute_freqs_cis(
    head_dim: int,
    max_seq_len: int,
    theta: float = 10_000.0,
) -> torch.Tensor:
    """Return complex-valued frequency tensor [max_seq_len, head_dim // 2]."""
    # Frequencies for even-indexed dimensions: 1 / θ^(2i/d)
    freqs = 1.0 / (theta ** (torch.arange(0, head_dim, 2).float() / head_dim))
    t = torch.arange(max_seq_len, dtype=torch.float32)
    freqs = torch.outer(t, freqs)                               # [T, head_dim/2]
    return torch.polar(torch.ones_like(freqs), freqs)           # complex64


def _apply_rotary_emb(
    q: torch.Tensor,          # [B, T, n_heads, head_dim]
    k: torch.Tensor,          # [B, T, n_heads, head_dim]
    freqs_cis: torch.Tensor,  # [T, head_dim // 2]  complex
) -> tuple[torch.Tensor, torch.Tensor]:
    """Apply RoPE to query and key tensors in-place of computation."""
    T = q.shape[1]

    def rotate(x: torch.Tensor) -> torch.Tensor:
        # Reshape last dim into pairs, view as complex, rotate, flatten back
        x_c = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
        # freqs_cis: [T, head_dim/2] → [1, T, 1, head_dim/2]
        f = freqs_cis[:T].unsqueeze(0).unsqueeze(2)
        return torch.view_as_real(x_c * f).flatten(3).type_as(x)

    return rotate(q), rotate(k)


# ─────────────────────────────────────────────────────────────────────────────
# SwiGLU Feed-Forward Network
# ─────────────────────────────────────────────────────────────────────────────

class SwiGLUFFN(nn.Module):
    """FFN with SwiGLU activation (Shazeer, 2020).

    FFN(x) = down_proj( SiLU(gate_proj(x)) ⊙ up_proj(x) )

    No bias anywhere; intermediate_size is typically ~ 8/3 * hidden_size
    (rounded to a multiple of 256 for efficiency).
    """

    def __init__(self, hidden_size: int, intermediate_size: int) -> None:
        super().__init__()
        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.up_proj   = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))


# ─────────────────────────────────────────────────────────────────────────────
# Causal Multi-Head Self-Attention  (no bias, RoPE, SDPA)
# ─────────────────────────────────────────────────────────────────────────────

class CausalSelfAttention(nn.Module):
    """Standard causal MHA.

    Uses torch.nn.functional.scaled_dot_product_attention for efficient
    (potentially Flash Attention) computation.  No local window; full sequence
    attends causally.
    """

    def __init__(self, cfg: ModelConfig) -> None:
        super().__init__()
        assert cfg.hidden_size % cfg.num_heads == 0, (
            f"hidden_size ({cfg.hidden_size}) must be divisible by "
            f"num_heads ({cfg.num_heads})"
        )
        self.n_heads  = cfg.num_heads
        self.head_dim = cfg.hidden_size // cfg.num_heads
        self.attn_drop = cfg.dropout

        # No bias in any projection — consistent with LLaMA
        self.q_proj = nn.Linear(cfg.hidden_size, cfg.hidden_size, bias=False)
        self.k_proj = nn.Linear(cfg.hidden_size, cfg.hidden_size, bias=False)
        self.v_proj = nn.Linear(cfg.hidden_size, cfg.hidden_size, bias=False)
        self.o_proj = nn.Linear(cfg.hidden_size, cfg.hidden_size, bias=False)

    def forward(
        self,
        x: torch.Tensor,                      # [B, T, D]
        freqs_cis: torch.Tensor,              # [T, head_dim//2] complex
        attn_mask: Optional[torch.Tensor] = None,  # [B, T] 1=real 0=pad
    ) -> torch.Tensor:
        B, T, _ = x.shape

        # Project and split heads
        q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim)
        k = self.k_proj(x).view(B, T, self.n_heads, self.head_dim)
        v = self.v_proj(x).view(B, T, self.n_heads, self.head_dim)

        # Apply RoPE to q and k
        q, k = _apply_rotary_emb(q, k, freqs_cis)

        # [B, T, n_heads, head_dim] → [B, n_heads, T, head_dim]
        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)

        # Build additive attention bias [B, 1, T, T]:
        #   causal constraint (upper-triangular = -inf) +
        #   optional key-padding mask (pad positions = -inf)
        bias = _build_attention_bias(T, attn_mask, x.device, x.dtype)

        out = F.scaled_dot_product_attention(
            q, k, v,
            attn_mask=bias,
            dropout_p=self.attn_drop if self.training else 0.0,
            is_causal=False,   # causal constraint already encoded in `bias`
        )

        # [B, n_heads, T, head_dim] → [B, T, D]
        out = out.transpose(1, 2).contiguous().view(B, T, -1)
        return self.o_proj(out)


def _build_attention_bias(
    T: int,
    attn_mask: Optional[torch.Tensor],  # [B, T_k] or None
    device: torch.device,
    dtype: torch.dtype,
) -> torch.Tensor:
    """Additive attention bias combining causal mask and optional padding mask."""
    # Causal mask: upper triangle → -inf, rest → 0
    # triu(..., diagonal=1) gives positions j > i
    causal = torch.zeros(T, T, device=device, dtype=dtype)
    causal = causal.masked_fill(
        torch.ones(T, T, device=device, dtype=torch.bool).triu(diagonal=1),
        float("-inf"),
    )
    bias = causal.unsqueeze(0).unsqueeze(0)  # [1, 1, T, T]

    if attn_mask is not None:
        # attn_mask [B, T]: 0 → padding → mask out as key
        # [B, T] → [B, 1, 1, T] additive -inf at pad positions
        pad = (attn_mask == 0).unsqueeze(1).unsqueeze(2)      # [B, 1, 1, T]
        pad_bias = torch.zeros_like(pad, dtype=dtype).masked_fill(pad, float("-inf"))
        bias = bias + pad_bias                                 # broadcast → [B, 1, T, T]

    return bias


# ─────────────────────────────────────────────────────────────────────────────
# Transformer Block  (pre-norm: norm → sublayer → residual)
# ─────────────────────────────────────────────────────────────────────────────

class TransformerBlock(nn.Module):
    def __init__(self, cfg: ModelConfig) -> None:
        super().__init__()
        self.attn_norm = RMSNorm(cfg.hidden_size)
        self.attn      = CausalSelfAttention(cfg)
        self.ffn_norm  = RMSNorm(cfg.hidden_size)
        self.ffn       = SwiGLUFFN(cfg.hidden_size, cfg.ffn_intermediate_size)

    def forward(
        self,
        x: torch.Tensor,
        freqs_cis: torch.Tensor,
        attn_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        x = x + self.attn(self.attn_norm(x), freqs_cis, attn_mask)
        x = x + self.ffn(self.ffn_norm(x))
        return x


# ─────────────────────────────────────────────────────────────────────────────
# IPA Gated Fusion  (embedding-layer augmentation)
# ─────────────────────────────────────────────────────────────────────────────

class IPAFusion(nn.Module):
    """Projects IPA phonological feature vectors into model space and fuses
    them with token embeddings.

    Fusion modes
    ------------
    ipa_add   : fused = embed + W_ipa · ipa
    ipa_gate  : fused = embed + σ(W_gate · ipa) ⊙ (W_ipa · ipa)
    ipa_full  : same gated fusion + stochastic masking of ipa inputs
                during training (for aux-loss reconstruction)

    For ipa_full, the method also returns the *original* (unmasked) IPA
    tensor so the auxiliary head can compute a reconstruction target.

    The gate naturally collapses to ~0 when IPA is absent (all-zero input),
    so tokens without phonological coverage (e.g. punctuation, numerals) are
    unaffected regardless of variant.
    """

    def __init__(self, cfg: ModelConfig) -> None:
        super().__init__()
        self.variant    = cfg.variant
        self.mask_ratio = cfg.ipa_mask_ratio

        # IPA → model-space projection (all variants except baseline)
        self.ipa_proj = nn.Linear(cfg.ipa_dim, cfg.hidden_size, bias=False)

        # Gating network (ipa_gate and ipa_full)
        if cfg.variant in ("ipa_gate", "ipa_full"):
            self.gate_proj = nn.Linear(cfg.ipa_dim, cfg.hidden_size, bias=False)

    def forward(
        self,
        embed: torch.Tensor,   # [B, T, D]  token embeddings
        ipa: torch.Tensor,     # [B, T, ipa_dim]
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Returns (fused_embed, ipa_target).

        ipa_target is the original unmasked IPA tensor, needed by the aux head
        when computing the reconstruction loss for ipa_full.
        """
        ipa_target = ipa                       # save original for aux loss

        # Stochastic IPA dropout during training (ipa_full only)
        if self.training and self.variant == "ipa_full" and self.mask_ratio > 0:
            # Sample a binary mask: 1 = keep, 0 = zero-out
            keep_prob = 1.0 - self.mask_ratio
            mask = torch.bernoulli(
                torch.full(ipa.shape[:2], keep_prob, device=ipa.device)
            ).unsqueeze(-1)                    # [B, T, 1]
            ipa = ipa * mask

        ipa_emb = self.ipa_proj(ipa)           # [B, T, D]

        if self.variant == "ipa_add":
            return embed + ipa_emb, ipa_target

        # ipa_gate / ipa_full
        gate = torch.sigmoid(self.gate_proj(ipa))   # [B, T, D]
        return embed + gate * ipa_emb, ipa_target


# ─────────────────────────────────────────────────────────────────────────────
# Auxiliary IPA Prediction Head  (ipa_full only)
# ─────────────────────────────────────────────────────────────────────────────

class IPAAuxHead(nn.Module):
    """Predicts raw IPA feature vectors from final hidden states.

    Loss is MSE, computed only at positions where the ground-truth IPA vector
    is non-zero (i.e. the token was mapped to a real phoneme).
    """

    def __init__(self, hidden_size: int, ipa_dim: int) -> None:
        super().__init__()
        self.proj = nn.Linear(hidden_size, ipa_dim, bias=False)

    def loss(
        self,
        hidden: torch.Tensor,      # [B, T, D]
        ipa_target: torch.Tensor,  # [B, T, ipa_dim]  original unmasked IPA
    ) -> torch.Tensor:
        pred = self.proj(hidden)                           # [B, T, ipa_dim]

        # Mask: positions where the token has real IPA coverage
        has_ipa = (ipa_target.abs().sum(dim=-1) > 0).float()  # [B, T]
        n_ipa = has_ipa.sum()
        if n_ipa == 0:
            return hidden.new_zeros(())

        mse = F.mse_loss(pred, ipa_target, reduction="none").mean(dim=-1)  # [B, T]
        return (mse * has_ipa).sum() / n_ipa


# ─────────────────────────────────────────────────────────────────────────────
# Main Model
# ─────────────────────────────────────────────────────────────────────────────

class BabyLMModel(nn.Module):
    """LLaMA-style causal transformer with optional IPA embedding augmentation.

    Parameters
    ----------
    cfg : ModelConfig
        Full model configuration (see config.py).

    Notes
    -----
    • Weight tying: lm_head.weight == embed.weight (halves embedding params).
    • freqs_cis is registered as a non-persistent buffer so it moves with the
      model across devices but is not saved in checkpoints (recomputed on load).
    • Call model.configure_optimizers() to get a pre-configured AdamW with
      weight-decay applied only to ≥2-D parameters (standard LLaMA practice).
    """

    def __init__(self, cfg: ModelConfig) -> None:
        super().__init__()
        cfg.validate()
        self.cfg = cfg

        # Token embedding
        self.embed = nn.Embedding(cfg.vocab_size, cfg.hidden_size)

        # IPA fusion (None for baseline)
        self.ipa_fusion: Optional[IPAFusion] = (
            None if cfg.variant == "baseline" else IPAFusion(cfg)
        )

        # Embedding dropout
        self.emb_drop = nn.Dropout(cfg.dropout)

        # Transformer layers
        self.layers = nn.ModuleList(
            [TransformerBlock(cfg) for _ in range(cfg.num_layers)]
        )

        # Final norm
        self.norm = RMSNorm(cfg.hidden_size)

        # Language model head (no bias; weight-tied to embed)
        self.lm_head = nn.Linear(cfg.hidden_size, cfg.vocab_size, bias=False)
        self.lm_head.weight = self.embed.weight   # weight tying

        # Auxiliary IPA prediction head (ipa_full only)
        self.ipa_aux_head: Optional[IPAAuxHead] = (
            IPAAuxHead(cfg.hidden_size, cfg.ipa_dim)
            if cfg.variant == "ipa_full"
            else None
        )

        # RoPE frequencies — registered as non-persistent buffer
        # (recomputed on load, not stored in state_dict)
        freqs_cis = _precompute_freqs_cis(
            head_dim    = cfg.hidden_size // cfg.num_heads,
            max_seq_len = cfg.max_seq_len,
        )
        self.register_buffer("freqs_cis", freqs_cis, persistent=False)

        # Weight initialisation
        self.apply(self._init_weights)
        # Scale residual-stream projections by 1/√(2·num_layers) per GPT-2 / LLaMA
        scale = (2.0 * cfg.num_layers) ** -0.5
        for name, p in self.named_parameters():
            if name.endswith(("o_proj.weight", "down_proj.weight")):
                nn.init.normal_(p, mean=0.0, std=0.02 * scale)

    # ---------------------------------------------------------------------- #
    # Weight initialisation                                                    #
    # ---------------------------------------------------------------------- #

    def _init_weights(self, module: nn.Module) -> None:
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)

    # ---------------------------------------------------------------------- #
    # Forward pass                                                             #
    # ---------------------------------------------------------------------- #

    def forward(
        self,
        input_ids:      torch.Tensor,                    # [B, T]  LongTensor
        ipa_vectors:    torch.Tensor,                    # [B, T, ipa_dim]
        attention_mask: Optional[torch.Tensor] = None,  # [B, T]  1=real 0=pad
    ) -> ModelOutput:
        B, T = input_ids.shape
        assert T <= self.cfg.max_seq_len, (
            f"Sequence length {T} exceeds max_seq_len {self.cfg.max_seq_len}"
        )

        # ── Token embedding ───────────────────────────────────────────────
        x = self.embed(input_ids)                        # [B, T, D]

        # ── IPA fusion at embedding layer ─────────────────────────────────
        ipa_target = ipa_vectors                         # default target (no aux head)
        if self.ipa_fusion is not None:
            x, ipa_target = self.ipa_fusion(x, ipa_vectors)

        x = self.emb_drop(x)

        # ── Transformer layers ────────────────────────────────────────────
        freqs_cis: torch.Tensor = self.freqs_cis        # type: ignore[assignment]
        for layer in self.layers:
            x = layer(x, freqs_cis, attention_mask)

        # ── Final normalisation ───────────────────────────────────────────
        x = self.norm(x)                                 # [B, T, D]

        # ── LM head ───────────────────────────────────────────────────────
        logits = self.lm_head(x)                         # [B, T, vocab_size]

        # ── Auxiliary IPA loss (ipa_full only) ────────────────────────────
        aux_loss: torch.Tensor
        if self.ipa_aux_head is not None:
            aux_loss = self.ipa_aux_head.loss(x, ipa_target)
        else:
            aux_loss = logits.new_zeros(())

        return ModelOutput(logits=logits, aux_loss=aux_loss)

    # ---------------------------------------------------------------------- #
    # Convenience helpers                                                      #
    # ---------------------------------------------------------------------- #

    def num_parameters(self, non_embedding: bool = True) -> int:
        """Count trainable parameters, optionally excluding embedding table."""
        n = sum(p.numel() for p in self.parameters() if p.requires_grad)
        if non_embedding:
            n -= self.embed.weight.numel()
        return n

    def configure_optimizers(
        self,
        lr: float,
        weight_decay: float,
        betas: tuple[float, float] = (0.9, 0.95),
        eps: float = 1e-8,
    ) -> torch.optim.AdamW:
        """AdamW with weight decay applied only to ≥2-D parameters.

        Biases (there are none here), 1-D RMSNorm scales, and the embedding
        table are placed in the no-decay group.
        """
        decay, no_decay = [], []
        for name, p in self.named_parameters():
            if not p.requires_grad:
                continue
            if p.ndim >= 2 and "embed" not in name:
                decay.append(p)
            else:
                no_decay.append(p)

        groups = [
            {"params": decay,    "weight_decay": weight_decay},
            {"params": no_decay, "weight_decay": 0.0},
        ]
        return torch.optim.AdamW(groups, lr=lr, betas=betas, eps=eps)


# ─────────────────────────────────────────────────────────────────────────────
# Quick smoke-test
# ─────────────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Smoke-test BabyLMModel")
    parser.add_argument("--variant", default="ipa_full",
                        choices=["baseline", "ipa_add", "ipa_gate", "ipa_full"])
    parser.add_argument("--device", default="cpu")
    args = parser.parse_args()

    cfg = ModelConfig(variant=args.variant, num_layers=2, hidden_size=128,
                      num_heads=4, ffn_intermediate_size=344, vocab_size=1000,
                      max_seq_len=64, ipa_dim=24)
    model = BabyLMModel(cfg).to(args.device)
    model.train()

    B, T = 2, 32
    ids  = torch.randint(0, cfg.vocab_size, (B, T), device=args.device)
    ipa  = torch.randn(B, T, cfg.ipa_dim, device=args.device)
    mask = torch.ones(B, T, dtype=torch.long, device=args.device)
    mask[0, -4:] = 0   # simulate padding in first sequence

    out = model(ids, ipa, mask)
    print(f"variant    : {cfg.variant}")
    print(f"logits     : {out.logits.shape}  dtype={out.logits.dtype}")
    print(f"aux_loss   : {out.aux_loss.item():.6f}")
    print(f"params (non-emb): {model.num_parameters():,}")

    # Verify backward pass
    loss = out.logits.mean() + out.aux_loss
    loss.backward()
    print("backward   : OK")