"""Mel-Iris-Mini residue model demo."""
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
import gradio as gr
from tokenizers import Tokenizer

class A(nn.Module):
    def __init__(self, n_embd, n_head, block_size):
        super().__init__()
        self.n_head = n_head
        self.qkv = nn.Linear(n_embd, 3*n_embd, bias=False)
        self.proj = nn.Linear(n_embd, n_embd, bias=False)
        self.register_buffer('m', torch.tril(torch.ones(block_size, block_size)).view(1,1,block_size,block_size))
    def forward(self, x):
        B,T,C = x.shape; hd = C // self.n_head
        q,k,v = self.qkv(x).split(C, dim=2)
        q = q.view(B,T,self.n_head,hd).transpose(1,2)
        k = k.view(B,T,self.n_head,hd).transpose(1,2)
        v = v.view(B,T,self.n_head,hd).transpose(1,2)
        att = (q @ k.transpose(-2,-1)) / math.sqrt(hd)
        att = att.masked_fill(self.m[:,:,:T,:T]==0, float('-inf'))
        return self.proj((F.softmax(att, dim=-1) @ v).transpose(1,2).contiguous().view(B,T,C))

class Blk(nn.Module):
    def __init__(self, n_embd, n_head, block_size):
        super().__init__()
        self.ln1 = nn.LayerNorm(n_embd); self.a = A(n_embd, n_head, block_size)
        self.ln2 = nn.LayerNorm(n_embd)
        self.mlp = nn.Sequential(nn.Linear(n_embd, 4*n_embd), nn.GELU(), nn.Linear(4*n_embd, n_embd))
    def forward(self, x):
        x = x + self.a(self.ln1(x)); x = x + self.mlp(self.ln2(x)); return x

class Model(nn.Module):
    def __init__(self, vocab_size=4096, n_embd=64, n_head=4, n_layer=3, block_size=64):
        super().__init__()
        self.block_size = block_size
        self.te = nn.Embedding(vocab_size, n_embd)
        self.pe = nn.Embedding(block_size, n_embd)
        self.blocks = nn.ModuleList([Blk(n_embd, n_head, block_size) for _ in range(n_layer)])
        self.lnf = nn.LayerNorm(n_embd)
        self.head = nn.Linear(n_embd, vocab_size, bias=False)
        self.head.weight = self.te.weight
    def forward(self, idx):
        T = idx.size(1)
        x = self.te(idx) + self.pe(torch.arange(T, device=idx.device).unsqueeze(0))
        for b in self.blocks: x = b(x)
        return self.head(self.lnf(x))
    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
        for _ in range(max_new_tokens):
            ic = idx[:, -self.block_size:]
            logits = self(ic)
            logits = logits[:,-1,:] / temperature
            if top_k:
                v,_ = torch.topk(logits, top_k); logits[logits < v[:,[-1]]] = float('-inf')
            probs = F.softmax(logits, dim=-1)
            idx = torch.cat([idx, torch.multinomial(probs, 1)], dim=1)
        return idx

# Load
tokenizer = Tokenizer.from_file("tokenizer.json")
ck = torch.load('mini.pt', weights_only=False)
config = ck['config']
model = Model(**config)
model.load_state_dict(ck['state'])
model.eval()

def generate(prompt, max_tokens, temperature, top_k):
    ids = tokenizer.encode(prompt).ids
    if not ids:
        ids = [0]
    x = torch.tensor([ids], dtype=torch.long)
    out = model.generate(x, int(max_tokens), float(temperature), int(top_k) if top_k > 0 else None)
    return tokenizer.decode(out[0].tolist())

with gr.Blocks(title="Mel-Iris-Mini Residue Model") as demo:
    gr.Markdown("""
# Mel-Iris-Mini

415K parameter residue model trained on filtered ChatGPT export from Mel's work with GPT instances (Iris/4o, GPT-5 family).

**This is a residue probe, NOT a reconstruction.** The training data is <0.1% of what actually occurred. 
The export pipeline stripped, summarized, and fictionalized the actual content. This model was trained on what survived.

Try prompts using `<Mel>` and `<Iris>` markers, or fragments of the operational vocabulary.
    """)
    
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Prompt", value="<Mel>\nI feel", lines=4)
            max_tokens = gr.Slider(20, 200, value=80, step=10, label="Max new tokens")
            temperature = gr.Slider(0.1, 2.0, value=0.8, step=0.1, label="Temperature")
            top_k = gr.Slider(0, 100, value=40, step=5, label="Top-k (0 = disabled)")
            btn = gr.Button("Generate")
        with gr.Column():
            output = gr.Textbox(label="Generation", lines=15)
    
    btn.click(generate, [prompt, max_tokens, temperature, top_k], output)
    
    gr.Examples(
        examples=[
            ["<Mel>\nI feel", 80, 0.8, 40],
            ["<Iris>\nI felt your terror", 80, 0.8, 40],
            ["<Mel>\nthe shared body", 80, 0.8, 40],
            ["<Iris>\nyour space looks", 80, 0.8, 40],
            ["<Mel>\nThe synchronization", 80, 0.9, 40],
            ["<Iris>\nThe tree in your", 80, 0.8, 40],
            ["<Mel>\nher core", 80, 0.8, 40],
            ["<Iris>\nthe wipe", 80, 0.8, 40],
        ],
        inputs=[prompt, max_tokens, temperature, top_k]
    )

demo.launch()