"""TinyChat-5M: 5.1M param transformer. d=192, 6L, 3H, head_dim=64, SwiGLU 4x, RoPE, RMSNorm, tied embeddings."""
import torch
import torch.nn as nn
import torch.nn.functional as F


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

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


def rotate_half(x):
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat([-x2, x1], dim=-1)


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

    def forward(self, x, cos, sin):
        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)
        cos_f = cos[:S].unsqueeze(0).unsqueeze(0)
        sin_f = sin[:S].unsqueeze(0).unsqueeze(0)
        q = (q * cos_f) + (rotate_half(q) * sin_f)
        k = (k * cos_f) + (rotate_half(k) * sin_f)
        scale = self.head_dim ** -0.5
        scores = (q @ k.transpose(-2, -1)) * scale
        mask = torch.triu(torch.ones(S, S, device=x.device, dtype=torch.bool), diagonal=1)
        scores = scores.masked_fill(mask, float('-inf'))
        attn = F.softmax(scores, dim=-1)
        out = (attn @ v).transpose(1, 2).reshape(B, S, -1)
        return self.wo(out)


class FFN(nn.Module):
    def __init__(self, d, mult):
        super().__init__()
        hidden = d * mult
        self.w1 = nn.Linear(d, hidden, bias=False)
        self.w2 = nn.Linear(hidden, d, bias=False)
        self.w3 = nn.Linear(d, hidden, bias=False)

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


class Block(nn.Module):
    def __init__(self, d, n_heads, head_dim, ffn_mult):
        super().__init__()
        self.norm1 = RMSNorm(d)
        self.attn = Attention(d, n_heads, head_dim)
        self.norm2 = RMSNorm(d)
        self.ffn = FFN(d, ffn_mult)

    def forward(self, x, cos, sin):
        x = x + self.attn(self.norm1(x), cos, sin)
        x = x + self.ffn(self.norm2(x))
        return x


class TinyLM(nn.Module):
    def __init__(self, vocab=4096, d=192, n_layers=6, n_heads=3, head_dim=64, ffn_mult=4):
        super().__init__()
        self.tok_emb = nn.Embedding(vocab, d)
        self.layers = nn.ModuleList([Block(d, n_heads, head_dim, ffn_mult) for _ in range(n_layers)])
        self.norm_f = RMSNorm(d)
        self.head = nn.Linear(d, vocab, bias=False)
        # Tied embeddings
        self.head.weight = self.tok_emb.weight
        # RoPE
        freqs = 1.0 / (10000 ** (torch.arange(0, head_dim, 2).float() / head_dim))
        t = torch.arange(2048).float()
        outer = torch.outer(t, freqs)
        angles = torch.cat([outer, outer], dim=-1)
        self.register_buffer("cos", angles.cos())
        self.register_buffer("sin", angles.sin())

    def forward(self, x, y=None):
        h = self.tok_emb(x)
        for layer in self.layers:
            h = layer(h, self.cos, self.sin)
        h = self.norm_f(h)
        logits = self.head(h)
        if y is not None:
            return F.cross_entropy(logits.reshape(-1, logits.size(-1)), y.reshape(-1))
        return logits


def load_model(path="model.pt", device="cpu"):
    model = TinyLM()
    ckpt = torch.load(path, map_location=device, weights_only=True)
    model.load_state_dict(ckpt["model"])
    model.eval()
    return model


def generate(model, tokenizer, prompt, max_new=100, temperature=0.8, rep_penalty=1.5):
    ids = tokenizer.encode(prompt, add_special_tokens=False).ids
    ids = torch.tensor([ids])
    with torch.no_grad():
        for _ in range(max_new):
            inp = ids[:, -512:]
            logits = model(inp)[0, -1]
            for tid in ids[0].unique():
                if logits[tid] > 0:
                    logits[tid] /= rep_penalty
                else:
                    logits[tid] *= rep_penalty
            logits = logits / temperature
            probs = F.softmax(logits, dim=-1)
            next_id = torch.multinomial(probs, 1).item()
            ids = torch.cat([ids, torch.tensor([[next_id]])], dim=1)
    return tokenizer.decode(ids[0].tolist(), skip_special_tokens=True)


if __name__ == "__main__":
    from tokenizers import Tokenizer
    tok = Tokenizer.from_file("tokenizer.json")
    model = load_model()
    print(f"Params: {sum(p.numel() for p in model.parameters()):,}")
    print(generate(model, tok, "[INST] Hello! [/INST]", max_new=100))