"""LDT-10M — 10M-param LLaMA-style prose LM, trained from scratch.

Self-contained model definition (no transformers dependency). Load the
safetensors weights with `LDT.from_pretrained` (below) or the snippet in the
README.

Architecture (10,284,480 params, tied embeddings):
  - vocab 12288 (gollem_eval byte-level BPE)
  - d_model 320, n_layers 5, n_heads 5 (head_dim 64), SwiGLU FFN inter 896
  - RMSNorm pre-norm, RoPE (base 10000), causal attention, ctx 512
  - Standard LLaMA (no sliding window)
"""
import math
import torch
import torch.nn as nn
import torch.nn.functional as F


def precompute_rope(dim, max_pos, base=10000.0):
    freqs = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
    t = torch.arange(max_pos).float()
    angles = torch.outer(t, freqs)
    return torch.polar(torch.ones_like(angles), angles)  # complex


def apply_rope(x, freqs_cis, offset=0):
    B, nh, S, hd = x.shape
    x = x.view(B, nh, S, hd // 2, 2)
    xr = x[..., 0].float()
    xi = x[..., 1].float()
    fc = freqs_cis[offset:offset + S].to(x.device)
    fr, fi = fc.real, fc.imag
    out_r = xr * fr - xi * fi
    out_i = xr * fi + xi * fr
    out = torch.stack([out_r, out_i], dim=-1).view(B, nh, S, hd)
    return out.to(x.dtype)


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

    def forward(self, x):
        norm = x.float().pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
        return (x.float() * norm).to(x.dtype) * self.weight


class Attention(nn.Module):
    def __init__(self, d, n_heads):
        super().__init__()
        self.n_heads = n_heads
        self.head_dim = d // n_heads
        self.wq = nn.Linear(d, d, bias=False)
        self.wk = nn.Linear(d, d, bias=False)
        self.wv = nn.Linear(d, d, bias=False)
        self.wo = nn.Linear(d, d, bias=False)

    def forward(self, x, freqs_cis, offset=0):
        B, S, _ = x.shape
        q = self.wq(x).view(B, S, self.n_heads, self.head_dim).transpose(1, 2)
        k = self.wk(x).view(B, S, self.n_heads, self.head_dim).transpose(1, 2)
        v = self.wv(x).view(B, S, self.n_heads, self.head_dim).transpose(1, 2)
        q = apply_rope(q, freqs_cis, offset)
        k = apply_rope(k, freqs_cis, offset)
        out = F.scaled_dot_product_attention(q, k, v)
        out = out.transpose(1, 2).contiguous().view(B, S, -1)
        return self.wo(out)


class MLP(nn.Module):
    def __init__(self, d, ff):
        super().__init__()
        self.w1 = nn.Linear(d, ff, bias=False)  # gate
        self.w2 = nn.Linear(d, ff, bias=False)  # up
        self.w3 = nn.Linear(ff, d, bias=False)  # down

    def forward(self, x):
        return self.w3(F.silu(self.w1(x)) * self.w2(x))


class Block(nn.Module):
    def __init__(self, d, n_heads, ff):
        super().__init__()
        self.ln1 = RMSNorm(d)
        self.attn = Attention(d, n_heads)
        self.ln2 = RMSNorm(d)
        self.ffn = MLP(d, ff)

    def forward(self, x, freqs_cis, offset=0):
        x = x + self.attn(self.ln1(x), freqs_cis, offset)
        x = x + self.ffn(self.ln2(x))
        return x


class LDT(nn.Module):
    def __init__(self, vocab=12288, d=320, n_layers=5, n_heads=5, ff=896, ctx=512):
        super().__init__()
        self.vocab = vocab
        self.ctx = ctx
        self.d = d
        self.tok = nn.Embedding(vocab, d)
        self.blocks = nn.ModuleList([Block(d, n_heads, ff) for _ in range(n_layers)])
        self.ln_f = RMSNorm(d)
        self.head = nn.Linear(d, vocab, bias=False)
        self.head.weight = self.tok.weight  # tied
        self.freqs_cis = precompute_rope(d // n_heads, ctx)

    @property
    def tied_weights(self):
        return ["lm_head"]

    def forward(self, idx, targets=None):
        B, S = idx.shape
        h = self.tok(idx)
        for b in self.blocks:
            h = b(h, self.freqs_cis)
        h = self.ln_f(h)
        logits = self.head(h)
        if targets is not None:
            loss = F.cross_entropy(
                logits[:, :-1].reshape(-1, logits.size(-1)),
                targets[:, 1:].reshape(-1),
                ignore_index=-1,
            )
            return loss
        return logits

    @torch.no_grad()
    def generate(self, idx, max_new_tokens=128, temperature=0.8, top_k=40, seed=0):
        g = torch.Generator(device=idx.device).manual_seed(seed)
        for _ in range(max_new_tokens):
            ctx_in = idx[:, -self.ctx:]
            logits = self(ctx_in)[:, -1]
            if temperature and temperature > 0:
                logits = logits / temperature
            if top_k:
                v, _ = torch.topk(logits, top_k, dim=-1)
                logits[logits < v[:, -1, None]] = float("-inf")
            p = torch.softmax(logits, dim=-1)
            nxt = torch.multinomial(p, 1, generator=g)
            idx = torch.cat([idx, nxt], dim=1)
        return idx

    @classmethod
    def from_pretrained(cls, path, device="cpu"):
        """Load the safetensors weights (tied lm_head re-bound to tok.weight)."""
        from safetensors.torch import load_file

        m = cls()
        sd = load_file(path, device=device)
        # head.weight is tied to tok.weight and is NOT stored separately.
        sd.pop("head.weight", None)
        missing, unexpected = m.load_state_dict(sd, strict=False)
        assert "head.weight" in missing, f"unexpected missing keys: {missing}"
        assert not unexpected, f"unexpected keys: {unexpected}"
        m.head.weight = m.tok.weight  # restore the tie
        return m.to(device).eval()