import torch
import torch.nn as nn
import torch.nn.functional as F
import math
import gradio as gr
from huggingface_hub import hf_hub_download, HfApi
from transformers import GPT2Tokenizer
import json

# ── Hardcoded config ──────────────────────────────────────────────────────────
CFG = {
    "vocab_size": 50304,
    "d_model":    768,
    "n_layers":   24,
    "n_heads":    12,
    "n_kv_heads": 3,
    "ffn_dim":    1792,
    "max_len":    2048,
    "bus_points": [4, 9, 15, 21],
}

MODEL_REPO = "hugging-science/CodVa-2-session-001"

# ── Primitives ────────────────────────────────────────────────────────────────

class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.w = nn.Parameter(torch.ones(dim))
    def forward(self, x):
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.w

def build_rope(head_dim, max_len, theta, device):
    freqs = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
    t = torch.arange(max_len, device=device).float()
    return torch.cos(torch.outer(t, freqs)), torch.sin(torch.outer(t, freqs))

def apply_rope(x, cos, sin):
    B, H, L, D = x.shape
    x1, x2 = x[..., :D//2], x[..., D//2:]
    c = cos[:L].unsqueeze(0).unsqueeze(0)
    s = sin[:L].unsqueeze(0).unsqueeze(0)
    return torch.cat([x1*c - x2*s, x2*c + x1*s], dim=-1)

def repeat_kv(x, n_rep):
    if n_rep == 1: return x
    B, H, L, D = x.shape
    return x[:, :, None, :, :].expand(B, H, n_rep, L, D).reshape(B, H*n_rep, L, D)

class SwiGLU(nn.Module):
    def __init__(self, dim, hidden):
        super().__init__()
        self.gate = nn.Linear(dim, hidden, bias=False)
        self.up   = nn.Linear(dim, hidden, bias=False)
        self.down = nn.Linear(hidden, dim, bias=False)
    def forward(self, x):
        return self.down(F.silu(self.gate(x)) * self.up(x))

class GQA(nn.Module):
    def __init__(self, dim, n_heads, n_kv):
        super().__init__()
        self.n_heads = n_heads
        self.n_kv    = n_kv
        self.n_rep   = n_heads // n_kv
        self.hd      = dim // n_heads
        self.wq = nn.Linear(dim, n_heads * self.hd, bias=False)
        self.wk = nn.Linear(dim, n_kv    * self.hd, bias=False)
        self.wv = nn.Linear(dim, n_kv    * self.hd, bias=False)
        self.wo = nn.Linear(n_heads * self.hd, dim, bias=False)
        self.q_norm = RMSNorm(self.hd)
        self.k_norm = RMSNorm(self.hd)
    def forward(self, x, cos, sin):
        B, L, _ = x.shape
        q = self.wq(x).view(B, L, self.n_heads, self.hd).transpose(1, 2)
        k = self.wk(x).view(B, L, self.n_kv,    self.hd).transpose(1, 2)
        v = self.wv(x).view(B, L, self.n_kv,    self.hd).transpose(1, 2)
        q = self.q_norm(q); k = self.k_norm(k)
        q = apply_rope(q, cos, sin); k = apply_rope(k, cos, sin)
        k = repeat_kv(k, self.n_rep); v = repeat_kv(v, self.n_rep)
        out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        return self.wo(out.transpose(1, 2).contiguous().view(B, L, -1))

class TransformerBlock(nn.Module):
    def __init__(self, dim, n_heads, n_kv, ffn_dim, n_layers):
        super().__init__()
        self.norm1 = RMSNorm(dim)
        self.norm2 = RMSNorm(dim)
        self.attn  = GQA(dim, n_heads, n_kv)
        self.ffn   = SwiGLU(dim, ffn_dim)
        self.res_scale = 1.0 / math.sqrt(2.0 * max(n_layers, 1))
    def forward(self, x, cos, sin):
        x = x + self.res_scale * self.attn(self.norm1(x), cos, sin)
        x = x + self.res_scale * self.ffn(self.norm2(x))
        return x

class FiLMBridge(nn.Module):
    def __init__(self, d):
        super().__init__()
        self.conv    = nn.Conv1d(d, d, 5, padding=2, groups=d)
        self.to_film = nn.Linear(d, d * 2)
    def forward(self, h):
        local = self.conv(h.transpose(1,2)).transpose(1,2)
        gamma, beta = self.to_film(local).chunk(2, dim=-1)
        return h * (1 + gamma) + beta

class DepthTracker(nn.Module):
    def __init__(self, d, d_state=32):
        super().__init__()
        self.to_delta = nn.Linear(d, d_state)
        self.to_out   = nn.Linear(d_state, d)
        self.decay    = nn.Parameter(torch.ones(d_state) * 0.9)
    def forward(self, x):
        B, L, _ = x.shape
        delta = torch.tanh(self.to_delta(x))
        decay = torch.sigmoid(self.decay)
        depth = torch.zeros(B, delta.size(-1), device=x.device, dtype=x.dtype)
        outs = []
        for t in range(L):
            depth = depth * decay + delta[:, t]
            outs.append(depth)
        return self.to_out(torch.stack(outs, dim=1))

class CNNBus(nn.Module):
    def __init__(self, d):
        super().__init__()
        self.conv_n = nn.Conv1d(d, d, 3, padding=1, groups=d)
        self.conv_w = nn.Conv1d(d, d, 7, padding=3, groups=d)
        self.proj   = nn.Linear(d * 2, d)
    def forward(self, x):
        t = x.transpose(1, 2)
        n = F.silu(self.conv_n(t)).transpose(1, 2)
        w = F.silu(self.conv_w(t)).transpose(1, 2)
        return self.proj(torch.cat([n, w], dim=-1))

class RNNBus(nn.Module):
    def __init__(self, d):
        super().__init__()
        self.gru  = nn.GRU(d, d // 2, batch_first=True)
        self.proj = nn.Linear(d // 2, d)
    def forward(self, x):
        out, _ = self.gru(x.float())
        return self.proj(out.to(x.dtype))

class HighwayBus(nn.Module):
    def __init__(self, d):
        super().__init__()
        self.gate = nn.Linear(d * 2, d)
    def forward(self, bus, contrib):
        g = torch.sigmoid(self.gate(torch.cat([bus, contrib], dim=-1)))
        return bus + g * contrib

class BusInjector(nn.Module):
    def __init__(self):
        super().__init__()
        self.gate = nn.Parameter(torch.zeros(1))
    def forward(self, h, bus):
        return h + torch.sigmoid(self.gate) * bus

# ── Model ─────────────────────────────────────────────────────────────────────

class CodVa2(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        D = cfg['d_model']
        self.embed          = nn.Embedding(cfg['vocab_size'], D)
        self.film           = FiLMBridge(D)
        self.blocks         = nn.ModuleList([
            TransformerBlock(D, cfg['n_heads'], cfg['n_kv_heads'],
                             cfg['ffn_dim'], cfg['n_layers'])
            for _ in range(cfg['n_layers'])
        ])
        self.bus_points     = cfg['bus_points']
        self.highway        = HighwayBus(D)
        self.injectors      = nn.ModuleList([BusInjector() for _ in self.bus_points])
        self.depth_tracker  = DepthTracker(D)
        self.cnn_1          = CNNBus(D)
        self.rnn            = RNNBus(D)
        self.cnn_2          = CNNBus(D)
        self.final_injector = BusInjector()
        self.norm_f         = RMSNorm(D)
        self.lm_head        = nn.Linear(D, cfg['vocab_size'], bias=False)
        hd = D // cfg['n_heads']
        self.register_buffer("cos", torch.zeros(cfg['max_len'], hd // 2))
        self.register_buffer("sin", torch.zeros(cfg['max_len'], hd // 2))

    def _init_rope(self):
        cos, sin = build_rope(
            self.cfg['d_model'] // self.cfg['n_heads'],
            self.cfg['max_len'], 500000.0,
            self.embed.weight.device)
        self.cos.copy_(cos); self.sin.copy_(sin)

    @torch.no_grad()
    def generate(self, input_ids, max_new_tokens=128, temperature=0.8, top_p=0.92):
        if self.cos.sum() == 0: self._init_rope()
        ids = torch.tensor([input_ids]) if not isinstance(input_ids, torch.Tensor) \
              else input_ids.unsqueeze(0)
        bp = self.bus_points
        for _ in range(max_new_tokens):
            L = ids.shape[1]
            if L > self.cfg['max_len']:
                ids = ids[:, -self.cfg['max_len']:]
                L = ids.shape[1]
            h   = self.embed(ids)
            h   = self.film(h)
            bus = torch.zeros_like(h)
            cos = self.cos[:L]; sin = self.sin[:L]
            for i, block in enumerate(self.blocks):
                h = block(h, cos, sin)
                if   i == bp[0]: bus = self.highway(bus, self.depth_tracker(h)); h = self.injectors[0](h, bus)
                elif i == bp[1]: bus = self.highway(bus, self.cnn_1(h));         h = self.injectors[1](h, bus)
                elif i == bp[2]: h = self.injectors[2](h, bus)
                elif i == bp[3]: bus = self.highway(bus, self.rnn(h));           h = self.injectors[3](h, bus)
            bus = self.highway(bus, self.cnn_2(h))
            h   = self.final_injector(h, bus)
            h   = self.norm_f(h)
            logits = 30.0 * torch.tanh(self.lm_head(h[:, -1, :]) / 30.0)
            probs  = F.softmax(logits / temperature, dim=-1)
            sp, si = torch.sort(probs, descending=True)
            mask   = (torch.cumsum(sp, dim=-1) - sp) < top_p
            sp[~mask] = 0.0
            sp /= sp.sum() + 1e-9
            next_tok = si[torch.multinomial(sp, 1)]
            ids = torch.cat([ids, next_tok.unsqueeze(0)], dim=1)
            if next_tok.item() == 50256: break  # GPT2 EOS
        return ids[0].tolist()

# ── Load ──────────────────────────────────────────────────────────────────────

print("Loading tokenizer...")
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")

print("Finding latest checkpoint...")
api    = HfApi()
files  = list(api.list_repo_files(MODEL_REPO))
ckpts  = sorted([f for f in files if f.endswith(".pt") and "step" in f],
                key=lambda x: int(x.split("step")[-1].split(".")[0]))
latest = ckpts[-1]
print(f"Loading {latest}...")
ckpt   = torch.load(hf_hub_download(MODEL_REPO, latest), map_location="cpu")

model = CodVa2(CFG)
model.load_state_dict(ckpt['model'], strict=False)
model.eval()
print(f"Ready! {sum(p.numel() for p in model.parameters())/1e6:.1f}M params")

# ── Gradio ────────────────────────────────────────────────────────────────────

def generate(prompt, max_tokens, temperature, top_p):
    if not prompt.strip(): return "Enter a prompt."
    try:
        ids = tokenizer.encode(prompt)[-256:]
        out = model.generate(ids, int(max_tokens), float(temperature), float(top_p))
        return tokenizer.decode(out, skip_special_tokens=True)
    except Exception as e:
        return f"Error: {e}"

with gr.Blocks(title="CodVa-2", theme=gr.themes.Monochrome()) as demo:
    gr.Markdown("""
    # 🍰 CodVa-2 — Code Model
    **221M params** | 24L×768d | Highway Bus | GQA  
    ⚠️ *CPU only — ~30-60s per generation. Early checkpoint, still training.*
    """)
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Prompt", lines=5,
                                value="def fibonacci(n):\n    ")
            with gr.Row():
                max_tok = gr.Slider(10, 200, 80,  step=10,   label="Max Tokens")
                temp    = gr.Slider(0.1, 2.0, 0.8, step=0.1, label="Temperature")
                topp    = gr.Slider(0.1, 1.0, 0.92,step=0.05,label="Top-P")
            btn = gr.Button("Generate 🚀", variant="primary")
        with gr.Column():
            out = gr.Textbox(label="Output", lines=10, interactive=False)

    btn.click(generate, [prompt, max_tok, temp, topp], out)

    gr.Examples([
        ["def fibonacci(n):\n    ", 100, 0.8, 0.92],
        ["class Stack:\n    def __init__(self):\n        ", 120, 0.8, 0.92],
        ["# binary search\ndef search(arr, target):\n    ", 100, 0.7, 0.9],
    ], [prompt, max_tok, temp, topp])

demo.launch()