# app.py — CodVa-2 optimized for 2 vCPU free tier (realistic version)

import os
import math
import time

# ─────────────────────────────────────────────────────────────────────────────
#  REALISTIC CPU SETTINGS FOR 2 vCPUs
# ─────────────────────────────────────────────────────────────────────────────
# On free tier you have 2 vCPUs — don't over-thread
os.environ["OMP_NUM_THREADS"] = "2"
os.environ["MKL_NUM_THREADS"] = "2"
os.environ["OPENBLAS_NUM_THREADS"] = "1"  # OpenBLAS often better with 1 thread

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

torch.set_num_threads(2)
torch.set_num_interop_threads(1)  # Don't parallelize ops — causes overhead

print(f"[cpu] 2 vCPUs | torch threads: {torch.get_num_threads()}")

# ─────────────────────────────────────────────────────────────────────────────
#  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 (simple version, no fancy optimizations that hurt on 2 vCPUs)
# ─────────────────────────────────────────────────────────────────────────────
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')))

        # Simple matmul — no fancy batching
        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 = {
    "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")

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

# ─────────────────────────────────────────────────────────────────────────────
#  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]:
    
    if not prompt or not prompt.strip():
        yield "", "⚠️ Please enter a prompt"
        return

    domain_token = DOMAIN_TOKENS.get(domain, "<|domain_code|>")
    full_prompt  = domain_token + prompt.strip()

    enc = tokenizer.encode(full_prompt)
    ids = enc.ids
    
    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)):
        x_in   = x[:, -cfg.max_len:] if x.size(1) > cfg.max_len else x
        logits = model(x_in)[0, -1, :]

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

        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')

        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)

        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)

        full_text = tokenizer.decode(ids + generated, skip_special_tokens=False)
        
        display_text = full_text
        for tok in DOMAIN_TOKENS.values():
            display_text = display_text.replace(tok, "")

        elapsed = time.time() - t0
        tps     = len(generated) / max(elapsed, 1e-3)
        stats   = (
            f"⏱ {elapsed:.1f}s  |  "
            f"🔤 {len(generated)}/{max_new}  |  "
            f"⚡ {tps:.1f} tok/s  |  "
            f"🎯 {domain}"
        )

        yield display_text, stats

        if EOS_ID >= 0 and next_tok == EOS_ID:
            break

    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 | {len(generated)} tokens | {tps:.1f} tok/s"
    )

# ─────────────────────────────────────────────────────────────────────────────
#  GRADIO UI
# ─────────────────────────────────────────────────────────────────────────────
EXAMPLES = [
    ["def fibonacci(n):\n    ", "Code", 64, 0.2, 0.95, 50],
    ["1 + 1 = ", "Math", 10, 0.1, 0.95, 50],
    ["class Stack:\n    def __init__(self):\n        ", "Code", 100, 0.3, 0.95, 50],
    ["The Roman Empire began ", "General", 80, 0.6, 0.95, 50],
]

with gr.Blocks(title="CodVa-2") as demo:
    
    gr.Markdown("""
    # 🧠 CodVa-2 — Pretrained Code LM
    **213M params** · **2 vCPUs (free tier)** · Expect ~5-8 tok/s
    
    Pretrained model — continues text. Set **low max tokens** for short answers.
    """)

    with gr.Row():
        with gr.Column(scale=1):
            prompt_box = gr.Textbox(
                label="Prompt",
                placeholder="def fibonacci(n):\n    ",
                lines=8,
            )
            
            domain_dropdown = gr.Dropdown(
                choices=list(DOMAIN_TOKENS.keys()),
                value="Code",
                label="Domain",
            )
            
            with gr.Row():
                max_new_slider = gr.Slider(8, 256, value=64, step=8, label="Max tokens")
                temp_slider    = gr.Slider(0.0, 1.5, value=0.3, step=0.05, label="Temp")
            
            with gr.Row():
                gen_btn   = gr.Button("▶ Generate", variant="primary", scale=2)
                stop_btn  = gr.Button("⏹ Stop", variant="stop", scale=1)

        with gr.Column(scale=1):
            output_box = gr.Textbox(label="Output", lines=15, interactive=False)
            stats_box = gr.Textbox(label="", lines=1, interactive=False)

    gr.Examples(
        examples=EXAMPLES,
        inputs=[prompt_box, domain_dropdown, max_new_slider, temp_slider],
        label="Examples",
    )

    gr.Markdown("""
    💡 **Free tier = slow but it works!**  
    For `1+1=` type prompts: set max tokens to 10-15 and click Stop when done.
    """)

    gen_event = gen_btn.click(
        fn=generate_stream,
        inputs=[prompt_box, domain_dropdown, max_new_slider, temp_slider, 
                gr.State(0.95), gr.State(50)],  # hidden top_p, top_k
        outputs=[output_box, stats_box],
    )
    
    stop_btn.click(fn=None, cancels=[gen_event])

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