# app.py — CodVa-2 Demo (PRETRAIN model, domain tokens only)
import os
import math
import time
import torch
import torch.nn as nn
import torch.nn.functional as F
import gradio as gr
from dataclasses import dataclass
from typing import Tuple, Generator
from tokenizers import Tokenizer
from huggingface_hub import hf_hub_download, login, HfApi

# ─────────────────────────────────────────────────────────────────────────────
#  AUTH
# ─────────────────────────────────────────────────────────────────────────────
HF_TOKEN        = os.environ.get("HF_TOKEN", "")
HF_DATASET_REPO = "Bc-AI/nova1_data"
HF_MODEL_REPO   = os.environ.get("MODEL_REPO", "hugging-science/CodVa-2-session-002")

if HF_TOKEN:
    login(token=HF_TOKEN)

# ─────────────────────────────────────────────────────────────────────────────
#  CONFIG
# ─────────────────────────────────────────────────────────────────────────────
@dataclass
class Config:
    vocab_size:       int   = 50304
    d_model:          int   = 896
    n_layers:         int   = 18
    n_heads:          int   = 14
    n_kv_heads:       int   = 2
    max_len:          int   = 2048
    rope_theta:       float = 500_000.0
    window_size:      int   = 512
    pattern_mult:     float = 2.5
    reason_mult:      float = 0.75
    reason_depth:     int   = 2
    gate_hidden:      int   = 64
    gate_init:        float = 0.0
    diff_lambda_init: float = 0.8
    tie_embeddings:   bool  = True

    @property
    def head_dim(self):
        return self.d_model // self.n_heads

    @property
    def pattern_dim(self):
        return ((int(self.d_model * self.pattern_mult) + 255) // 256) * 256

    @property
    def reason_dim(self):
        return ((int(self.d_model * self.reason_mult) + 63) // 64) * 64

# ─────────────────────────────────────────────────────────────────────────────
#  MODEL
# ─────────────────────────────────────────────────────────────────────────────
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):
        x32 = x.float()
        return (x32 * torch.rsqrt(
            x32.pow(2).mean(-1, keepdim=True) + self.eps
        ) * self.w).to(x.dtype)


def precompute_rope(head_dim, max_len, theta, device):
    inv_freq = 1.0 / (theta ** (
        torch.arange(0, head_dim, 2, device=device, dtype=torch.float32) / head_dim
    ))
    pos   = torch.arange(max_len, device=device, dtype=torch.float32)
    freqs = torch.outer(pos, inv_freq)
    return freqs.cos(), freqs.sin()


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


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


class DifferentialAttention(nn.Module):
    def __init__(self, cfg, layer_idx, local=False):
        super().__init__()
        self.n_pairs    = cfg.n_heads // 2
        self.n_kv_pairs = max(1, cfg.n_kv_heads // 2)
        self.n_rep      = self.n_pairs // self.n_kv_pairs
        self.head_dim   = cfg.head_dim
        self.local      = local
        self.window     = cfg.window_size
        d               = cfg.d_model

        self.wq       = nn.Linear(d, 2 * self.n_pairs    * self.head_dim, bias=False)
        self.wk       = nn.Linear(d, 2 * self.n_kv_pairs * self.head_dim, bias=False)
        self.wv       = nn.Linear(d,     self.n_kv_pairs * self.head_dim, bias=False)
        self.wo       = nn.Linear(    self.n_pairs * self.head_dim, d,    bias=False)
        self.q_norm   = RMSNorm(self.head_dim)
        self.k_norm   = RMSNorm(self.head_dim)
        self.lambda1  = nn.Parameter(torch.tensor(0.0))
        self.lambda2  = nn.Parameter(torch.tensor(0.0))
        self.out_norm = RMSNorm(self.head_dim)

    def _window_mask(self, L, device):
        idx  = torch.arange(L, device=device)
        dist = idx.unsqueeze(0) - idx.unsqueeze(1)
        mask = (dist > 0) | (dist < -self.window)
        return mask.float().masked_fill(mask, float('-inf'))

    def forward(self, x, cos, sin):
        B, L, _ = x.shape
        hd = self.head_dim

        q_all  = self.wq(x).view(B, L, self.n_pairs,    2, hd).transpose(1, 2)
        k_all  = self.wk(x).view(B, L, self.n_kv_pairs, 2, hd).transpose(1, 2)
        v      = self.wv(x).view(B, L, self.n_kv_pairs,    hd).transpose(1, 2)
        q1, q2 = q_all[..., 0, :], q_all[..., 1, :]
        k1, k2 = k_all[..., 0, :], k_all[..., 1, :]

        q1 = self.q_norm(q1); q2 = self.q_norm(q2)
        k1 = self.k_norm(k1); k2 = self.k_norm(k2)
        q1 = apply_rope(q1, cos, sin); q2 = apply_rope(q2, cos, sin)
        k1 = apply_rope(k1, cos, sin); k2 = apply_rope(k2, cos, sin)
        k1 = repeat_kv(k1, self.n_rep); k2 = repeat_kv(k2, self.n_rep)
        v  = repeat_kv(v,  self.n_rep)

        scale = 1.0 / math.sqrt(hd)
        mask  = (self._window_mask(L, x.device) if self.local
                 else torch.zeros(L, L, device=x.device).masked_fill(
                     ~torch.ones(L, L, device=x.device, dtype=torch.bool).tril(),
                     float('-inf')))

        a1   = F.softmax(torch.matmul(q1, k1.transpose(-2, -1)) * scale + mask, dim=-1)
        a2   = F.softmax(torch.matmul(q2, k2.transpose(-2, -1)) * scale + mask, dim=-1)
        lam  = torch.exp(self.lambda1) - torch.exp(self.lambda2) + 0.5
        out  = torch.matmul(a1 - lam * a2, v)
        out  = self.out_norm(out).transpose(1, 2).contiguous().view(B, L, -1)
        return self.wo(out)


class DualPathFFN(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        d, pd, rd     = cfg.d_model, cfg.pattern_dim, cfg.reason_dim
        self.pat_gate = nn.Linear(d,  pd, bias=False)
        self.pat_up   = nn.Linear(d,  pd, bias=False)
        self.pat_down = nn.Linear(pd, d,  bias=False)
        layers = [nn.Linear(d, rd, bias=False), nn.SiLU()]
        for _ in range(cfg.reason_depth - 1):
            layers += [nn.Linear(rd, rd, bias=False), nn.SiLU()]
        layers.append(nn.Linear(rd, d, bias=False))
        self.reason = nn.Sequential(*layers)
        self.merge  = nn.Parameter(torch.zeros(d))

    def forward(self, x):
        pat = self.pat_down(F.silu(self.pat_gate(x)) * self.pat_up(x))
        w   = torch.sigmoid(self.merge)
        return w * pat + (1.0 - w) * self.reason(x)


class TokenImportanceGate(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(cfg.d_model, cfg.gate_hidden, bias=True),
            nn.SiLU(),
            nn.Linear(cfg.gate_hidden, 1, bias=True),
        )
    def forward(self, x):
        return x * (0.5 + torch.sigmoid(self.net(x)))


class Block(nn.Module):
    def __init__(self, cfg, layer_idx):
        super().__init__()
        self.norm1 = RMSNorm(cfg.d_model)
        self.norm2 = RMSNorm(cfg.d_model)
        self.attn  = DifferentialAttention(cfg, layer_idx, local=(layer_idx % 2 == 0))
        self.ffn   = DualPathFFN(cfg)

    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 CodVa2(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg        = cfg
        self.embed      = nn.Embedding(cfg.vocab_size, cfg.d_model)
        self.importance = TokenImportanceGate(cfg)
        self.blocks     = nn.ModuleList([Block(cfg, i) for i in range(cfg.n_layers)])
        self.final_norm = RMSNorm(cfg.d_model)
        self.register_buffer("rope_cos", torch.zeros(cfg.max_len, cfg.head_dim // 2))
        self.register_buffer("rope_sin", torch.zeros(cfg.max_len, cfg.head_dim // 2))
        self._rope_ready = False

    def _init_rope(self, device):
        c, s = precompute_rope(
            self.cfg.head_dim, self.cfg.max_len, self.cfg.rope_theta, device
        )
        self.rope_cos.copy_(c)
        self.rope_sin.copy_(s)
        self._rope_ready = True

    def forward(self, tokens):
        B, L   = tokens.shape
        device = tokens.device
        if not self._rope_ready:
            self._init_rope(device)
        cos = self.rope_cos[:L]
        sin = self.rope_sin[:L]
        x   = self.importance(self.embed(tokens))
        for block in self.blocks:
            x = block(x, cos, sin)
        return F.linear(self.final_norm(x), self.embed.weight)

# ─────────────────────────────────────────────────────────────────────────────
#  LOAD TOKENIZER + MODEL
# ─────────────────────────────────────────────────────────────────────────────
print("[init] loading tokenizer...")
tok_path  = hf_hub_download(
    repo_id=HF_DATASET_REPO,
    filename="nova_tokenizer.json",
    repo_type="dataset",
    token=HF_TOKEN or None,
)
tokenizer = Tokenizer.from_file(tok_path)

# Domain tokens
DOMAIN_TOKENS = {
    "Code":      "<|domain_code|>",
    "Math":      "<|domain_math|>",
    "General":   "<|domain_general|>",
    "Reasoning": "<|domain_reasoning|>",
}

EOS_ID = tokenizer.token_to_id("<|endoftext|>") or tokenizer.token_to_id("</s>") or -1
print(f"[init] tokenizer | vocab={tokenizer.get_vocab_size()} | eos={EOS_ID}")

print("[init] loading model...")
cfg            = Config()
cfg.vocab_size = (tokenizer.get_vocab_size() + 63) // 64 * 64
model          = CodVa2(cfg)

api    = HfApi()
files  = list(api.list_repo_files(
    repo_id=HF_MODEL_REPO, repo_type="model", token=HF_TOKEN or None
))
finals = sorted([f for f in files if "final" in f and f.endswith(".safetensors")])
ckpts  = sorted(
    [f for f in files if "step" in f and f.endswith(".pt")],
    key=lambda x: int(x.split("step")[-1].split(".")[0])
)

if finals:
    print(f"[init] loading final: {finals[-1]}")
    from safetensors.torch import load_file
    wpath = hf_hub_download(HF_MODEL_REPO, finals[-1],
                            repo_type="model", token=HF_TOKEN or None)
    model.load_state_dict(load_file(wpath, device="cpu"), strict=True)

elif ckpts:
    print(f"[init] loading checkpoint: {ckpts[-1]}")
    wpath = hf_hub_download(HF_MODEL_REPO, ckpts[-1],
                            repo_type="model", token=HF_TOKEN or None)
    ckpt  = torch.load(wpath, map_location="cpu", weights_only=False)
    model.load_state_dict(ckpt["model"], strict=True)

else:
    print("[init] WARNING: no weights found — random init")

model.eval()
n_params = sum(p.numel() for p in model.parameters())
print(f"[init] ready | {n_params/1e6:.1f}M params | CPU inference")

# ─────────────────────────────────────────────────────────────────────────────
#  STREAMING GENERATION
# ─────────────────────────────────────────────────────────────────────────────
@torch.no_grad()
def generate_stream(
    prompt:      str,
    domain:      str,
    max_new:     int   = 256,
    temperature: float = 0.8,
    top_p:       float = 0.95,
    top_k:       int   = 50,
) -> Generator[Tuple[str, str], None, None]:
    """
    Pretrain-style generation with domain token prepending.
    Model sees: <|domain_X|>{prompt}
    Continues from there.
    """
    if not prompt or not prompt.strip():
        yield "", "⚠️ Please enter a prompt"
        return

    # Prepend domain token (matches training data format)
    domain_token = DOMAIN_TOKENS.get(domain, "<|domain_code|>")
    full_prompt  = domain_token + prompt.strip()

    enc = tokenizer.encode(full_prompt)
    ids = enc.ids
    
    # Truncate if too long
    if len(ids) > cfg.max_len - max_new:
        ids = ids[-(cfg.max_len - max_new):]

    x         = torch.tensor([ids], dtype=torch.long)
    generated = []
    t0        = time.time()

    for step in range(int(max_new)):
        # Truncate context to max_len
        x_in   = x[:, -cfg.max_len:] if x.size(1) > cfg.max_len else x
        logits = model(x_in)[0, -1, :].float()

        # Temperature scaling
        logits = logits / max(float(temperature), 1e-5)

        # Top-k filtering
        if top_k > 0:
            k           = min(int(top_k), logits.size(-1))
            topk_vals,_ = torch.topk(logits, k)
            logits[logits < topk_vals[-1]] = float('-inf')

        # Top-p (nucleus) filtering
        probs                = F.softmax(logits, dim=-1)
        sorted_p, sorted_idx = probs.sort(descending=True)
        cumsum_p             = sorted_p.cumsum(0)
        sorted_p[(cumsum_p - sorted_p) > float(top_p)] = 0.0
        sorted_p             = sorted_p / sorted_p.sum().clamp(min=1e-9)

        # Sample
        next_tok = sorted_idx[torch.multinomial(sorted_p, num_samples=1)].item()

        generated.append(next_tok)
        x = torch.cat([x, torch.tensor([[next_tok]])], dim=1)

        # Decode what we have so far (strip domain token from display)
        full_text = tokenizer.decode(ids + generated, skip_special_tokens=False)
        
        # Remove domain token from display
        display_text = full_text
        for tok in DOMAIN_TOKENS.values():
            display_text = display_text.replace(tok, "")

        # Build stats
        elapsed = time.time() - t0
        tps     = len(generated) / max(elapsed, 1e-3)
        stats   = (
            f"⏱ {elapsed:.1f}s  |  "
            f"🔤 {len(generated)} / {max_new} tokens  |  "
            f"⚡ {tps:.1f} tok/s  |  "
            f"🎯 {domain}  |  "
            f"🌡 {temperature}  top-p {top_p}  top-k {int(top_k)}"
        )

        yield display_text, stats

        # Stop on EOS
        if EOS_ID >= 0 and next_tok == EOS_ID:
            break

    # Final yield
    full_text = tokenizer.decode(ids + generated, skip_special_tokens=False)
    for tok in DOMAIN_TOKENS.values():
        full_text = full_text.replace(tok, "")
    
    elapsed = time.time() - t0
    tps     = len(generated) / max(elapsed, 1e-3)
    yield (
        full_text,
        f"✅ Done  |  ⏱ {elapsed:.1f}s  |  "
        f"🔤 {len(generated)} tokens  |  ⚡ {tps:.1f} tok/s"
    )

# ─────────────────────────────────────────────────────────────────────────────
#  GRADIO UI
# ─────────────────────────────────────────────────────────────────────────────
EXAMPLES = [
    ["def fibonacci(n):\n    ", "Code", 128, 0.2, 0.95, 50],
    ["class BinaryTree:\n    def __init__(self):\n        ", "Code", 256, 0.3, 0.95, 50],
    ["import torch\nimport torch.nn as nn\n\n", "Code", 200, 0.4, 0.95, 50],
    ["SELECT users.name, orders.total FROM ", "Code", 100, 0.3, 0.90, 40],
    ["# Quicksort implementation\ndef quicksort(arr):\n    ", "Code", 200, 0.2, 0.95, 50],
    ["Problem: Find the derivative of f(x) = x^3 + 2x^2 - 5x + 1\n\nSolution: ", "Math", 150, 0.4, 0.95, 50],
    ["Theorem: The sum of angles in a triangle equals 180 degrees.\n\nProof: ", "Math", 200, 0.5, 0.95, 50],
    ["Let $f(x) = \\int_0^x t^2 dt$. Then ", "Math", 128, 0.3, 0.95, 50],
    ["The history of the Roman Empire began ", "General", 200, 0.7, 0.95, 50],
    ["Photosynthesis is the process by which ", "General", 150, 0.5, 0.95, 50],
]

CSS = """
.container { max-width: 1100px; margin: auto; }
.code-box { 
    font-family: 'JetBrains Mono', 'Fira Code', 'Courier New', monospace !important;
    font-size: 13px !important; 
    line-height: 1.5 !important; 
}
"""

with gr.Blocks(title="CodVa-2 Pretrain Demo") as demo:

    gr.HTML(f"<style>{CSS}</style>")
    
    gr.Markdown("""
    # 🧠 CodVa-2 — Pretrained Code LM
    **213M parameters** · Differential Attention · Trained on code/math/general corpus
    
    This is a **pretrained** model (not instruction-tuned). It continues text in the style of its training domain.
    Use the domain selector to control what kind of continuation you get.
    """)

    with gr.Row():
        # ── Left: inputs ──────────────────────────────────────────────────────
        with gr.Column(scale=1):
            prompt_box = gr.Textbox(
                label="Prompt (raw text, model will continue)",
                placeholder="def fibonacci(n):\n    ",
                lines=10,
                elem_classes=["code-box"],
            )
            
            domain_dropdown = gr.Dropdown(
                choices=list(DOMAIN_TOKENS.keys()),
                value="Code",
                label="Domain (prepends domain token)",
                info="Code, Math, General, or Reasoning — tells the model what style to use"
            )
            
            with gr.Row():
                max_new_slider = gr.Slider(16, 512, value=256, step=16,
                                           label="Max new tokens")
                temp_slider    = gr.Slider(0.0, 2.0, value=0.8, step=0.05,
                                           label="Temperature")
            with gr.Row():
                topp_slider = gr.Slider(0.1, 1.0, value=0.95, step=0.05,
                                        label="Top-p")
                topk_slider = gr.Slider(1, 200, value=50, step=1,
                                        label="Top-k")
            with gr.Row():
                gen_btn   = gr.Button("▶ Generate", variant="primary", scale=3)
                stop_btn  = gr.Button("⏹ Stop",    variant="stop",    scale=1)
                clear_btn = gr.Button("🗑 Clear",                      scale=1)

        # ── Right: output ─────────────────────────────────────────────────────
        with gr.Column(scale=1):
            output_box = gr.Textbox(
                label="Generated continuation (streaming)",
                lines=20,
                interactive=False,
                elem_classes=["code-box"],
            )
            stats_box = gr.Textbox(
                label="",
                lines=1,
                interactive=False,
            )

    gr.Examples(
        examples=EXAMPLES,
        inputs=[prompt_box, domain_dropdown, max_new_slider, temp_slider, topp_slider, topk_slider],
        label="📋 Example prompts — click to load",
        examples_per_page=10,
    )

    gr.Markdown("""
    ---
    💡 **Tips:**  
    - **Domain matters:** Code domain → code syntax, Math → equations, General → prose  
    - **Lower temp (0.1-0.3)** = deterministic, predictable (good for code)  
    - **Higher temp (0.7-1.2)** = creative, varied (good for text)  
    - This model has seen **~2B tokens** (20% trained). Expect coherent syntax but sometimes wrong logic.  
    - By 10B tokens it should be much stronger.
    """)

    # ── Wire up events ────────────────────────────────────────────────────────
    gen_event = gen_btn.click(
        fn=generate_stream,
        inputs=[prompt_box, domain_dropdown, max_new_slider, temp_slider, topp_slider, topk_slider],
        outputs=[output_box, stats_box],
    )
    
    prompt_box.submit(
        fn=generate_stream,
        inputs=[prompt_box, domain_dropdown, max_new_slider, temp_slider, topp_slider, topk_slider],
        outputs=[output_box, stats_box],
    )
    
    stop_btn.click(fn=None, cancels=[gen_event])

    clear_btn.click(
        fn=lambda: ("", "", ""),
        outputs=[prompt_box, output_box, stats_box],
    )

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True,
    )