# app.py — CodVa-2 HF Space Inference

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
from transformers import GPT2Tokenizer
import json

# ── Minimal inference build (copy the primitives from your training script) ────

class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 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: int, max_len: int, theta: float, device):
    freqs = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
    t     = torch.arange(max_len, device=device).float()
    freqs = torch.outer(t, freqs)
    return torch.cos(freqs), torch.sin(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_model):
        super().__init__()
        self.conv    = nn.Conv1d(d_model, d_model, kernel_size=5, padding=2, groups=d_model)
        self.to_film = nn.Linear(d_model, d_model * 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_model, d_state=32):
        super().__init__()
        self.to_delta = nn.Linear(d_model, d_state)
        self.to_out   = nn.Linear(d_state, d_model)
        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 CNNBusContrib(nn.Module):
    def __init__(self, d_model):
        super().__init__()
        self.conv_n = nn.Conv1d(d_model, d_model, 3, padding=1, groups=d_model)
        self.conv_w = nn.Conv1d(d_model, d_model, 7, padding=3, groups=d_model)
        self.proj   = nn.Linear(d_model * 2, d_model)

    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 RNNBusContrib(nn.Module):
    def __init__(self, d_model):
        super().__init__()
        self.gru  = nn.GRU(d_model, d_model // 2, num_layers=1, batch_first=True)
        self.proj = nn.Linear(d_model // 2, d_model)

    def forward(self, x):
        out, _ = self.gru(x.float())
        return self.proj(out.to(x.dtype))

class HighwayBus(nn.Module):
    def __init__(self, d_model):
        super().__init__()
        self.gate = nn.Linear(d_model * 2, d_model)

    def forward(self, bus, contribution):
        g = torch.sigmoid(self.gate(torch.cat([bus, contribution], dim=-1)))
        return bus + g * contribution

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

# ── CodVa-2 Inference Model ────

class CodVa2Inference(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        
        self.embed = nn.Embedding(cfg['vocab_size'], cfg['d_model'])
        self.film = FiLMBridge(cfg['d_model'])
        
        self.blocks = nn.ModuleList([
            TransformerBlock(cfg['d_model'], cfg['n_heads'], cfg['n_kv_heads'],
                           cfg['ffn_dim'], n_layers=cfg['n_layers'])
            for _ in range(cfg['n_layers'])
        ])
        
        self.bus_points = cfg['bus_points']
        self.highway = HighwayBus(cfg['d_model'])
        self.injectors = nn.ModuleList([BusInjector() for _ in self.bus_points])
        
        self.depth_tracker = DepthTracker(cfg['d_model'])
        self.cnn_1 = CNNBusContrib(cfg['d_model'])
        self.rnn = RNNBusContrib(cfg['d_model'])
        self.cnn_2 = CNNBusContrib(cfg['d_model'])
        
        self.norm_f = RMSNorm(cfg['d_model'])
        self.lm_head = nn.Linear(cfg['d_model'], cfg['vocab_size'], bias=False)
        self.final_injector = BusInjector()
        
        self.register_buffer("cos", torch.zeros(cfg['max_len'], cfg['d_model'] // cfg['n_heads'] // 2))
        self.register_buffer("sin", torch.zeros(cfg['max_len'], cfg['d_model'] // cfg['n_heads'] // 2))
        
    def _init_rope(self, device):
        cos, sin = build_rope(self.cfg['d_model'] // self.cfg['n_heads'],
                             self.cfg['max_len'], 500000.0, device)
        self.cos.copy_(cos); self.sin.copy_(sin)
    
    @torch.no_grad()
    def generate(self, input_ids, max_new_tokens=100, temperature=0.7, top_p=0.9):
        """Generate tokens autoregressively."""
        device = self.embed.weight.device
        
        if self.cos.sum() == 0:
            self._init_rope(device)
        
        if not isinstance(input_ids, torch.Tensor):
            input_ids = torch.tensor([input_ids], device=device)
        else:
            input_ids = input_ids.to(device)
        
        if input_ids.dim() == 1:
            input_ids = input_ids.unsqueeze(0)
        
        for _ in range(max_new_tokens):
            L = input_ids.shape[1]
            if L > self.cfg['max_len']:
                input_ids = input_ids[:, -self.cfg['max_len']:]
                L = self.cfg['max_len']
            
            # Forward pass
            h = self.embed(input_ids)
            h = self.film(h)
            bus = torch.zeros_like(h)
            
            bp = self.bus_points
            for i, block in enumerate(self.blocks):
                h = block(h, self.cos[:L], self.sin[:L])
                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 = self.lm_head(h[:, -1:, :])
            logits = 30.0 * torch.tanh(logits / 30.0)
            
            # Sample
            probs = F.softmax(logits / temperature, dim=-1)
            sorted_probs, sorted_indices = torch.sort(probs, descending=True)
            cumsum = torch.cumsum(sorted_probs, dim=-1)
            mask = cumsum <= top_p
            mask[..., 0] = True
            filtered_probs = sorted_probs * mask.float()
            filtered_probs /= filtered_probs.sum(dim=-1, keepdim=True) + 1e-9
            
            next_token = torch.multinomial(filtered_probs.squeeze(1), 1)
            input_ids = torch.cat([input_ids, next_token], dim=1)
        
        return input_ids[0].tolist()

# ── Load Model & Tokenizer ────

MODEL_REPO = "hugging-science/CodVa-2-session-001"
CHECKPOINT = "codva2_s001_step0001680.pt"

print("Loading model...")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Download config
config_path = hf_hub_download(MODEL_REPO, "config.json")
with open(config_path) as f:
    cfg = json.load(f)

# Download checkpoint
ckpt_path = hf_hub_download(MODEL_REPO, CHECKPOINT)

# Build & load model
model = CodVa2Inference(cfg).to(device)
state = torch.load(ckpt_path, map_location=device)
model.load_state_dict(state['model'], strict=False)
model.eval()

# Load GPT2 tokenizer
try:
    tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
except:
    tokenizer = GPT2Tokenizer.from_pretrained("openai-community/gpt2")

print(f"Model loaded: {sum(p.numel() for p in model.parameters())/1e6:.1f}M params")

# ── Gradio Interface ────

def generate_code(prompt: str, max_tokens: int = 100, temperature: float = 0.7, top_p: float = 0.9):
    """Generate code from prompt."""
    try:
        # Tokenize
        input_ids = tokenizer.encode(prompt, return_tensors="pt")[0]
        
        # Generate
        output_ids = model.generate(input_ids.tolist(), max_new_tokens=max_tokens,
                                   temperature=temperature, top_p=top_p)
        
        # Decode
        output = tokenizer.decode(output_ids, skip_special_tokens=True)
        
        return output
    except Exception as e:
        return f"Error: {str(e)}"

# Gradio UI
with gr.Blocks(title="CodVa-2") as demo:
    gr.Markdown("# CodVa-2 — Deep-Narrow Code Model")
    gr.Markdown("221M parameters, trained on 10B tokens of code+math")
    
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="def fibonacci(n):",
                lines=5,
                value="def hello():"
            )
            
            with gr.Row():
                max_tokens = gr.Slider(10, 300, value=100, step=10, label="Max Tokens")
                temperature = gr.Slider(0.0, 2.0, value=0.7, step=0.1, label="Temperature")
                top_p = gr.Slider(0.0, 1.0, value=0.9, step=0.05, label="Top-P")
            
            generate_btn = gr.Button("Generate", variant="primary")
        
        with gr.Column():
            output = gr.Textbox(
                label="Generated Code",
                lines=10,
                interactive=False
            )
    
    generate_btn.click(
        fn=generate_code,
        inputs=[prompt, max_tokens, temperature, top_p],
        outputs=output
    )
    
    gr.Markdown("""
    ## About CodVa-2
    - **Architecture**: 24-layer deep-narrow transformer (768d) with GQA
    - **Highway Bus**: FiLM bridge, DepthTracker, CNN experts, RNN expert
    - **Training**: 10B mixed code+math tokens (StarCoder + FineWeb)
    - **Stability**: Residual scaling, logit soft-cap, gradient clipping
    - **Status**: Early checkpoint (step 1680) — expect improvement
    """)

if __name__ == "__main__":
    demo.launch(share=True)