"""
model_trace.py — Decoder-only transformer for LimbPad-CoT trace generation (T3-T6).

Architecture (Design D — LimbPad-CoT):
  - Pre-norm decoder-only transformer
  - RoPE positional encoding (no learned absolute pos)
  - Weight-tied output head (embedding weight = lm_head weight)
  - SwiGLU FFN
  - Multi-head attention (no GQA in this v1 — single group)
  - Vocab: 268 tokens (0..255 digit + 12 control tokens, v3 grammar)

Default "full" config (~15M params): d_model=384, n_layers=6, n_heads=8, ffn_mult=4
Smoke-test config: d_model=128, n_layers=2, n_heads=4

Deviations from design packet:
  - Design packet specified d_model=1024/24 layers (~310M). We target ~15M for the
    CPU inference budget. The design packet's large model is for a future GPU training
    run; the current deliverable is the 15M CPU-deployable version.
  - No Abacus limb-significance embeddings in v1 (would require knowing operand
    boundaries; we rely on RoPE + trained attention patterns).
  - No GQA: full MHA throughout (simpler, sufficient at 15M scale).
"""

import math
import sys
from dataclasses import dataclass, field

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

# Vocabulary constants (inference-only copy; must match training grammar):
# 0..255 digit tokens, 256 PAD, 257 SEP, 258 EQ, 259 ANS, 260 EOS,
# 261 MUL_ROW, 262 ADD_ROW, 263 MOD_ROW, 264 CAR,
# 265 DBL, 266 QBIT0, 267 QBIT1 (v3 grammar).
# The vocab size actually used at load time comes from the checkpoint's
# pickled ModelConfig, so older v2 checkpoints (vocab 265) still load.
VOCAB_SIZE = 268
PAD = 256


# ──────────────────────────────────────────────
# Config
# ──────────────────────────────────────────────

@dataclass
class ModelConfig:
    vocab_size: int = VOCAB_SIZE     # 268 (v3 grammar)
    d_model: int = 384
    n_layers: int = 6
    n_heads: int = 8
    ffn_mult: int = 4                # FFN hidden = d_model * ffn_mult
    max_seq_len: int = 2048
    dropout: float = 0.0             # inference model; train with small dropout if desired
    tie_weights: bool = True         # tie embedding and lm_head weights
    pad_id: int = PAD


def default_config() -> ModelConfig:
    """~15M parameter config."""
    return ModelConfig(d_model=384, n_layers=6, n_heads=8)


def smoke_config() -> ModelConfig:
    """~2M parameter config for smoke tests."""
    return ModelConfig(d_model=128, n_layers=2, n_heads=4, max_seq_len=512)


# ──────────────────────────────────────────────
# RoPE
# ──────────────────────────────────────────────

def _build_rope_cache(seq_len: int, head_dim: int, device: torch.device, dtype: torch.dtype,
                      base: float = 10000.0) -> tuple[torch.Tensor, torch.Tensor]:
    """Build cos/sin cache for RoPE. Returns (cos, sin) each [seq_len, head_dim//2]."""
    half = head_dim // 2
    inv_freq = 1.0 / (base ** (torch.arange(0, half, dtype=torch.float32, device=device) / half))
    t = torch.arange(seq_len, device=device, dtype=torch.float32)
    freqs = torch.outer(t, inv_freq)          # [seq_len, half]
    return freqs.cos().to(dtype), freqs.sin().to(dtype)


def _apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    """
    Apply RoPE to query or key tensor.
    x: [B, n_heads, T, head_dim]
    cos, sin: [T, head_dim//2]
    """
    half = x.shape[-1] // 2
    x1 = x[..., :half]
    x2 = x[..., half:]
    # rotate_half: (-x2, x1)
    cos_ = cos.unsqueeze(0).unsqueeze(0)   # [1, 1, T, half]
    sin_ = sin.unsqueeze(0).unsqueeze(0)
    return torch.cat([
        x1 * cos_ - x2 * sin_,
        x2 * cos_ + x1 * sin_,
    ], dim=-1)


# ──────────────────────────────────────────────
# Attention
# ──────────────────────────────────────────────

class CausalSelfAttention(nn.Module):
    def __init__(self, cfg: ModelConfig):
        super().__init__()
        assert cfg.d_model % cfg.n_heads == 0
        self.n_heads = cfg.n_heads
        self.head_dim = cfg.d_model // cfg.n_heads
        self.d_model = cfg.d_model
        self.dropout = cfg.dropout

        self.qkv = nn.Linear(cfg.d_model, 3 * cfg.d_model, bias=False)
        self.out_proj = nn.Linear(cfg.d_model, cfg.d_model, bias=False)

        # KV cache (filled during inference)
        self._kv_cache_k: torch.Tensor | None = None
        self._kv_cache_v: torch.Tensor | None = None
        self._cache_len: int = 0

    def forward(
        self,
        x: torch.Tensor,           # [B, T, d]
        cos: torch.Tensor,         # [T_full, head_dim//2]
        sin: torch.Tensor,
        use_cache: bool = False,
    ) -> torch.Tensor:
        B, T, D = x.shape
        qkv = self.qkv(x)          # [B, T, 3D]
        q, k, v = qkv.split(self.d_model, dim=-1)

        # Reshape to [B, n_heads, T, head_dim]
        def reshape(t):
            return t.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)

        q, k, v = reshape(q), reshape(k), reshape(v)

        # RoPE — use correct positions depending on cache
        if use_cache and self._kv_cache_k is not None:
            start = self._cache_len
        else:
            start = 0
        pos_cos = cos[start: start + T]
        pos_sin = sin[start: start + T]
        q = _apply_rope(q, pos_cos, pos_sin)
        k = _apply_rope(k, pos_cos, pos_sin)

        # KV cache update
        if use_cache:
            if self._kv_cache_k is None:
                self._kv_cache_k = k
                self._kv_cache_v = v
            else:
                self._kv_cache_k = torch.cat([self._kv_cache_k, k], dim=2)
                self._kv_cache_v = torch.cat([self._kv_cache_v, v], dim=2)
            self._cache_len += T
            k = self._kv_cache_k
            v = self._kv_cache_v

        # Scaled dot-product attention (flash/mem-efficient kernels).
        # Causal when processing a full sequence (no cache, or first prompt
        # fill); single-token cached decode attends to the whole cache.
        T_q = q.shape[2]
        is_causal = (not use_cache) or self._cache_len == T
        out = F.scaled_dot_product_attention(
            q, k, v,
            dropout_p=self.dropout if (self.dropout > 0 and self.training) else 0.0,
            is_causal=is_causal,
        )
        out = out.transpose(1, 2).contiguous().view(B, T_q, D)
        return self.out_proj(out)

    def clear_cache(self):
        self._kv_cache_k = None
        self._kv_cache_v = None
        self._cache_len = 0


# ──────────────────────────────────────────────
# SwiGLU FFN
# ──────────────────────────────────────────────

class SwiGLUFFN(nn.Module):
    def __init__(self, cfg: ModelConfig):
        super().__init__()
        hidden = cfg.d_model * cfg.ffn_mult
        self.gate_proj = nn.Linear(cfg.d_model, hidden, bias=False)
        self.up_proj   = nn.Linear(cfg.d_model, hidden, bias=False)
        self.down_proj = nn.Linear(hidden, cfg.d_model, bias=False)
        self.dropout   = cfg.dropout

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        gate = F.silu(self.gate_proj(x))
        up   = self.up_proj(x)
        h    = gate * up
        if self.dropout > 0 and self.training:
            h = F.dropout(h, p=self.dropout)
        return self.down_proj(h)


# ──────────────────────────────────────────────
# Transformer block
# ──────────────────────────────────────────────

class TransformerBlock(nn.Module):
    def __init__(self, cfg: ModelConfig):
        super().__init__()
        self.norm1 = nn.RMSNorm(cfg.d_model)
        self.attn  = CausalSelfAttention(cfg)
        self.norm2 = nn.RMSNorm(cfg.d_model)
        self.ffn   = SwiGLUFFN(cfg)

    def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor,
                use_cache: bool = False) -> torch.Tensor:
        x = x + self.attn(self.norm1(x), cos, sin, use_cache=use_cache)
        x = x + self.ffn(self.norm2(x))
        return x


# ──────────────────────────────────────────────
# Main model
# ──────────────────────────────────────────────

class TraceTransformer(nn.Module):
    """
    Decoder-only transformer for scratchpad trace generation.
    Input: token id sequence. Output: next-token logits.
    """

    def __init__(self, cfg: ModelConfig | None = None):
        super().__init__()
        if cfg is None:
            cfg = default_config()
        self.cfg = cfg

        self.embedding = nn.Embedding(cfg.vocab_size, cfg.d_model, padding_idx=cfg.pad_id)
        self.blocks = nn.ModuleList([TransformerBlock(cfg) for _ in range(cfg.n_layers)])
        self.norm_out = nn.RMSNorm(cfg.d_model)
        self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)

        if cfg.tie_weights:
            self.lm_head.weight = self.embedding.weight

        # RoPE cache — built lazily
        self._rope_cos: torch.Tensor | None = None
        self._rope_sin: torch.Tensor | None = None

        self._init_weights()

    def _init_weights(self):
        std = 0.02
        nn.init.normal_(self.embedding.weight, std=std)
        for block in self.blocks:
            nn.init.normal_(block.attn.qkv.weight, std=std)
            nn.init.normal_(block.attn.out_proj.weight, std=std / math.sqrt(2 * self.cfg.n_layers))
            nn.init.normal_(block.ffn.gate_proj.weight, std=std)
            nn.init.normal_(block.ffn.up_proj.weight, std=std)
            nn.init.normal_(block.ffn.down_proj.weight, std=std / math.sqrt(2 * self.cfg.n_layers))
        if not self.cfg.tie_weights:
            nn.init.normal_(self.lm_head.weight, std=std)

    def _get_rope(self, seq_len: int, device: torch.device, dtype: torch.dtype):
        max_len = max(seq_len, self.cfg.max_seq_len)
        if (self._rope_cos is None
                or self._rope_cos.shape[0] < max_len
                or self._rope_cos.device != device):
            cos, sin = _build_rope_cache(max_len, self.cfg.d_model // self.cfg.n_heads,
                                         device=device, dtype=dtype)
            self._rope_cos = cos
            self._rope_sin = sin
        return self._rope_cos[:seq_len], self._rope_sin[:seq_len]

    def forward(
        self,
        input_ids: torch.Tensor,    # [B, T]
        use_cache: bool = False,
    ) -> torch.Tensor:              # [B, T, vocab_size]
        B, T = input_ids.shape
        device = input_ids.device
        x = self.embedding(input_ids)     # [B, T, d_model]
        dtype = x.dtype

        # For cache mode, we only need the positions of the current tokens
        if use_cache and self.blocks[0].attn._kv_cache_k is not None:
            start = self.blocks[0].attn._cache_len
            cos, sin = self._get_rope(start + T, device, dtype)
        else:
            cos, sin = self._get_rope(T, device, dtype)

        for block in self.blocks:
            x = block(x, cos, sin, use_cache=use_cache)

        x = self.norm_out(x)
        return self.lm_head(x)            # [B, T, vocab_size]

    def clear_cache(self):
        for block in self.blocks:
            block.attn.clear_cache()

    def num_parameters(self) -> int:
        return sum(p.numel() for p in self.parameters())

    def num_trainable_parameters(self) -> int:
        return sum(p.numel() for p in self.parameters() if p.requires_grad)


# ──────────────────────────────────────────────
# Parameter count estimate
# ──────────────────────────────────────────────

def estimate_params(cfg: ModelConfig) -> int:
    d = cfg.d_model
    V = cfg.vocab_size
    n = cfg.n_layers
    ffn_h = d * cfg.ffn_mult

    embedding = V * d                         # also lm_head if tied
    attn_per_layer = 3 * d * d + d * d        # qkv + out
    ffn_per_layer = 2 * d * ffn_h + ffn_h * d # gate + up + down
    norms_per_layer = 2 * d                   # 2 RMSNorm
    lm_head = 0 if cfg.tie_weights else V * d
    final_norm = d

    total = (embedding
             + n * (attn_per_layer + ffn_per_layer + norms_per_layer)
             + lm_head + final_norm)
    return total


# ──────────────────────────────────────────────
# CLI
# ──────────────────────────────────────────────

if __name__ == "__main__":
    for name, cfg in [("default (~15M)", default_config()), ("smoke (~2M)", smoke_config())]:
        model = TraceTransformer(cfg)
        est = estimate_params(cfg)
        actual = model.num_parameters()
        print(f"{name}: estimated={est:,}  actual={actual:,}")
        # Quick forward pass
        x = torch.randint(0, 256, (2, 32))
        logits = model(x)
        print(f"  forward ok: input {tuple(x.shape)} -> logits {tuple(logits.shape)}")
        del model