# ================================================================
#   CodeSymb AI — THE GRAND AWAKENING 🔥
#   ✅ Loads 240k Base Model
#   ✅ Loads 50,000-Step SFT Chat LoRA (Just finished!)
#   ✅ Mathematically Fuses for Native Speed
#   NiftyOracleAI 🇮🇳 | codesymb.com 🌎
# ================================================================

import gradio as gr
import torch
import torch.nn as nn
import torch.nn.functional as F
import sentencepiece as spm
import spaces
import os
from huggingface_hub import hf_hub_download, login

HF_TOKEN  = os.environ.get("HF_TOKEN", "")
USERNAME  = "NiftyOracleAI"
REPO_ID   = f"{USERNAME}/NiftyOracle-Coder-30B"
TOK_REPO  = f"{USERNAME}/niftyoracle-tokenizer"

if HF_TOKEN: login(token=HF_TOKEN)

# ══════════════════════════════════════════════════════════════
# CORE ARCHITECTURE (Fixed 2048 RoPE)
# ══════════════════════════════════════════════════════════════
class RMSNorm(nn.Module):
    def __init__(self, d, eps=1e-6):
        super().__init__()
        self.w = nn.Parameter(torch.ones(d)); self.eps = eps
    def forward(self, x):
        n = torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps)
        return (x.float() * n).to(x.dtype) * self.w.to(x.dtype)

class RoPE(nn.Module):
    def __init__(self, dim, max_len=2048, base=500_000):
        super().__init__()
        inv = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        emb = torch.cat([torch.outer(torch.arange(max_len).float(), inv)] * 2, dim=-1)
        self.register_buffer("cos", emb.cos()[None, None].to(torch.float16))
        self.register_buffer("sin", emb.sin()[None, None].to(torch.float16))
    def forward(self, n): return self.cos[:, :, :n], self.sin[:, :, :n]

def rot(x):
    a, b = x[..., :x.shape[-1]//2], x[..., x.shape[-1]//2:]
    return torch.cat([-b, a], dim=-1)

def rope(q, k, cos, sin): return (q*cos)+(rot(q)*sin), (k*cos)+(rot(k)*sin)

class Attn(nn.Module):
    def __init__(self, nh=16, nkv=4, dm=2048):
        super().__init__()
        hd = dm // nh; self.nh = nh; self.nkv = nkv; self.hd = hd; self.rep = nh // nkv
        self.q = nn.Linear(dm, nh*hd, bias=False)
        self.k = nn.Linear(dm, nkv*hd, bias=False)
        self.v = nn.Linear(dm, nkv*hd, bias=False)
        self.o = nn.Linear(nh*hd, dm, bias=False)
        self.pe = RoPE(hd)
    def forward(self, x):
        B, T, _ = x.shape
        q = self.q(x).view(B,T,self.nh,self.hd).transpose(1,2)
        k = self.k(x).view(B,T,self.nkv,self.hd).transpose(1,2)
        v = self.v(x).view(B,T,self.nkv,self.hd).transpose(1,2)
        cos, sin = self.pe(T); q, k = rope(q, k, cos, sin)
        k = k.repeat_interleave(self.rep, dim=1); v = v.repeat_interleave(self.rep, dim=1)
        o = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        return self.o(o.transpose(1,2).contiguous().view(B,T,-1))

class FFN(nn.Module):
    def __init__(self, dm=2048, df=8192):
        super().__init__()
        self.g = nn.Linear(dm, df, bias=False)
        self.u = nn.Linear(dm, df, bias=False)
        self.d = nn.Linear(df, dm, bias=False)
    def forward(self, x): return self.d(F.silu(self.g(x)) * self.u(x))

class Block(nn.Module):
    def __init__(self, dm=2048, df=8192, nh=16, nkv=4):
        super().__init__()
        self.n1 = RMSNorm(dm); self.att = Attn(nh, nkv, dm)
        self.n2 = RMSNorm(dm); self.ff  = FFN(dm, df)
    def forward(self, x):
        x = x + self.att(self.n1(x)); x = x + self.ff(self.n2(x)); return x

class CodeSymb1B(nn.Module):
    def __init__(self, vocab=50000):
        super().__init__()
        DM=2048; DF=8192; NL=16; NH=16; NKV=4
        self.embed  = nn.Embedding(vocab, DM)
        self.blocks = nn.ModuleList([Block(DM,DF,NH,NKV) for _ in range(NL)])
        self.norm   = RMSNorm(DM)
        self.head   = nn.Linear(DM, vocab, bias=False)
    def forward(self, tok):
        x = self.embed(tok).to(torch.float16)
        for b in self.blocks: x = b(x)
        return self.head(self.norm(x))

# ══════════════════════════════════════════════════════════════
# LOAD & FUSE THE 50,000-STEP CHAT LORA!
# ══════════════════════════════════════════════════════════════
print("⏳ Loading Tokenizer...")
tok_path = hf_hub_download(TOK_REPO, "niftyoracle_tokenizer.model", token=HF_TOKEN or None)
sp = spm.SentencePieceProcessor(); sp.load(tok_path)

print("⏳ Loading 240k Base Model...")
model = CodeSymb1B(vocab=sp.get_piece_size())
base_ckpt = hf_hub_download(REPO_ID, "checkpoints/latest.pt", token=HF_TOKEN or None)
state = torch.load(base_ckpt, map_location="cpu", weights_only=False)

base_weights = {k: v for k, v in state.get("model", state).items() if not k.endswith((".pe.cos", ".pe.sin"))}
model.load_state_dict(base_weights, strict=False)
print("✅ Base Model Loaded.")

print("⚡ FUSING THE NEW 50,000-STEP SFT LORA...")
try:
    # 🔥 PULLING THE EXACT FILE YOU JUST UPLOADED FROM LIGHTNING AI!
    lora_ckpt = hf_hub_download(REPO_ID, "checkpoints/sft_full_latest.pt", token=HF_TOKEN or None)
    lora_state_full = torch.load(lora_ckpt, map_location="cpu", weights_only=False)
    lora_state = lora_state_full.get("model", lora_state_full)
    
    scale = 200 / 100  # Alpha(200) / Rank(100)
    with torch.no_grad():
        for i, block in enumerate(model.blocks):
            for layer_name in ["att.q", "att.k", "att.v", "att.o", "ff.g", "ff.u", "ff.d"]:
                A_key = f"blocks.{i}.{layer_name}.A"
                B_key = f"blocks.{i}.{layer_name}.B"
                
                # Check for PEFT format keys
                peft_A_key = f"base_model.model.blocks.{i}.{layer_name}.lora_A.weight"
                peft_B_key = f"base_model.model.blocks.{i}.{layer_name}.lora_B.weight"
                
                A_val = lora_state.get(peft_A_key, lora_state.get(A_key))
                B_val = lora_state.get(peft_B_key, lora_state.get(B_key))
                
                if A_val is not None and B_val is not None:
                    A = A_val.float()
                    B = B_val.float()
                    delta = (B @ A) * scale if A.shape[1] == 100 else (A @ B).T * scale
                    
                    target_layer = getattr(block, layer_name.split(".")[0])
                    sub_layer = getattr(target_layer, layer_name.split(".")[1])
                    sub_layer.weight.data += delta.to(sub_layer.weight.dtype)
                    
    print("✅ 50k Chat LoRA Fused! The AI is fully conversational! 🗣️")
except Exception as e:
    print(f"⚠️ Failed to fuse Chat LoRA: {e}")

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device).half().eval()
print(f"✅ CodeSymb is ALIVE on {device}!")

# ══════════════════════════════════════════════════════════════
# INFERENCE ENGINE 
# ══════════════════════════════════════════════════════════════
@spaces.GPU(duration=60)
def chat_with_codesymb(message: str, history: list) -> str:
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)

    # EXACT SAME SYSTEM PROMPT FROM TRAINING!
    system_prompt = (
        "You are CodeSymb AI — expert coding assistant built by Deepakkumar Malkar "
        "at NiftyOracle India. Write clean code, explain simply, be friendly. "
        "Nothing is impossible. codesymb.com"
    )
    
    full_prompt = f"<|system|>\n{system_prompt}\n"
    for msg in history[-3:]:
        if isinstance(msg, dict):
            role = "<|user|>" if msg.get("role") == "user" else "<|assistant|>"
            full_prompt += f"{role}\n{msg.get('content', '')}\n"
        elif hasattr(msg, "role"):
            role = "<|user|>" if msg.role == "user" else "<|assistant|>"
            full_prompt += f"{role}\n{msg.content}\n"
        else:
            full_prompt += f"<|user|>\n{msg[0]}\n<|assistant|>\n{msg[1]}\n"
            
    full_prompt += f"<|user|>\n{message}\n<|assistant|>\n"
    
    prompt_ids = sp.encode(full_prompt)
    if len(prompt_ids) > 1024: prompt_ids = prompt_ids[-1024:]
        
    inp = torch.tensor([prompt_ids], dtype=torch.long).to(device)
    generated = []
    END_SEQ = sp.encode("<|endoftext|>")
    
    with torch.no_grad():
        for _ in range(512):
            with torch.autocast("cuda", dtype=torch.float16):
                logits = model(inp)[0, -1, :].float()

            for gid in set(generated[-50:]):
                logits[gid] -= 2.0

            logits = logits / 0.5  
            
            top_k_values, _ = torch.topk(logits, 40)
            logits[logits < top_k_values[-1]] = -float('inf')

            sorted_logits, sorted_idx = torch.sort(logits, descending=True)
            cum_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
            remove = cum_probs > 0.90
            remove[1:] = remove[:-1].clone(); remove[0] = False
            sorted_logits[remove] = -float("inf")
            
            logits_filtered = torch.zeros_like(logits)
            logits_filtered.scatter_(0, sorted_idx, sorted_logits)
            probs = F.softmax(logits_filtered, dim=-1)
            
            next_t = torch.multinomial(probs, 1).item()
            
            if len(generated) >= 3 and generated[-1] == next_t and generated[-2] == next_t:
                break

            generated.append(next_t)
            inp = torch.cat([inp, torch.tensor([[next_t]], device=device)], dim=1)

            decoded = sp.decode(generated)
            if "<|endoftext|>" in decoded or "<|user|>" in decoded:
                break

    response = sp.decode(generated).strip()
    for tag in ["<|user|>", "<|system|>", "<|assistant|>", "<|endoftext|>"]:
        response = response.split(tag)[0].strip()

    return response if response else "I am CodeSymb! How can I help you? 🔥"

with gr.Blocks(theme=gr.themes.Base(primary_hue="cyan", secondary_hue="blue"), css="body { background: #060D1F !important; color: white !important; }") as demo:
    gr.HTML("""
    <h1 style='text-align:center; color:#00D4FF;'>⚡ CodeSymb AI</h1>
    <p style='text-align:center; color:#A0AEC0;'>
        Built by <strong>Deepakkumar Malkar</strong> at NiftyOracle 🇮🇳<br>
    </p>
    """)
    gr.ChatInterface(
        fn=chat_with_codesymb,
        examples=[
            ["Hi! Who are you?"],
            ["Write a Python function to reverse a string."],
            ["What is an API?"],
        ]
    )

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