"""ProtoGPT — EXACT copy of the architecture in
PRETRAIN/stage07_tessera_proto_pilot.py so a pretrained checkpoint loads byte-exact.
Do NOT "improve" this — it must match the trainer's nn.Module key-for-key.
Plus load_base(): rebuild from a checkpoint's own config + load weights.
"""
import math
import torch
import torch.nn as nn
import torch.nn.functional as F


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

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        norm = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
        return norm * self.weight


class CausalSelfAttention(nn.Module):
    def __init__(self, d_model: int, heads: int, dropout: float):
        super().__init__()
        assert d_model % heads == 0
        self.heads = heads
        self.head_dim = d_model // heads
        self.dropout = dropout
        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
        self.proj = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        bsz, seq, _ = x.shape
        q, k, v = self.qkv(x).chunk(3, dim=-1)
        q = q.view(bsz, seq, self.heads, self.head_dim).transpose(1, 2)
        k = k.view(bsz, seq, self.heads, self.head_dim).transpose(1, 2)
        v = v.view(bsz, seq, self.heads, self.head_dim).transpose(1, 2)
        out = F.scaled_dot_product_attention(
            q, k, v,
            dropout_p=self.dropout if self.training else 0.0,
            is_causal=True,
        )
        out = out.transpose(1, 2).contiguous().view(bsz, seq, -1)
        return self.proj(out)


class Block(nn.Module):
    def __init__(self, d_model: int, heads: int, dropout: float):
        super().__init__()
        self.ln1 = RMSNorm(d_model)
        self.attn = CausalSelfAttention(d_model, heads, dropout)
        self.ln2 = RMSNorm(d_model)
        self.mlp = nn.Sequential(
            nn.Linear(d_model, 4 * d_model, bias=False),
            nn.GELU(),
            nn.Linear(4 * d_model, d_model, bias=False),
            nn.Dropout(dropout),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.attn(self.ln1(x))
        return x + self.mlp(self.ln2(x))


class ProtoGPT(nn.Module):
    def __init__(self, vocab_size, seq_len, d_model, layers, heads,
                 dropout=0.0, grad_checkpointing=False):
        super().__init__()
        self.seq_len = seq_len
        self.grad_checkpointing = grad_checkpointing
        self.tok = nn.Embedding(vocab_size, d_model)
        self.pos = nn.Embedding(seq_len, d_model)
        self.drop = nn.Dropout(dropout)
        self.blocks = nn.ModuleList([Block(d_model, heads, dropout) for _ in range(layers)])
        self.ln = RMSNorm(d_model)
        self.head = nn.Linear(d_model, vocab_size, bias=False)
        self.head.weight = self.tok.weight  # tied

    def forward(self, idx: torch.Tensor, targets: torch.Tensor | None = None,
                ignore_index: int = -100):
        _, seq = idx.shape
        if seq > self.seq_len:
            raise ValueError(f"seq {seq} exceeds seq_len={self.seq_len}")
        pos = torch.arange(seq, device=idx.device)
        x = self.drop(self.tok(idx) + self.pos(pos)[None, :, :])
        for block in self.blocks:
            x = block(x)
        logits = self.head(self.ln(x))
        loss = None
        if targets is not None:
            # masked CE: targets carry ignore_index on positions we don't train on
            loss = F.cross_entropy(
                logits.float().reshape(-1, logits.size(-1)),
                targets.reshape(-1),
                ignore_index=ignore_index,
            )
        return logits, loss


def load_base(ckpt_path: str, device: str = "cpu", dtype=torch.float32):
    """Rebuild ProtoGPT from a checkpoint's own config and load its weights."""
    ck = torch.load(ckpt_path, map_location="cpu", weights_only=False)
    c = ck["config"]
    model = ProtoGPT(
        vocab_size=c["vocab_size"], seq_len=c["seq_len"], d_model=c["d_model"],
        layers=c["layers"], heads=c["heads"], dropout=0.0, grad_checkpointing=False,
    )
    missing, unexpected = model.load_state_dict(ck["model"], strict=True)
    model.to(device=device, dtype=dtype)
    return model, c, ck.get("step")