#!/usr/bin/env python3
"""
app.py — HuggingFace Spaces demo for Competitive Docking Memory (CDM) Language Model
DuoNeural — Archon + Jesse Caldwell + Aura — 2026
Gradio 5 compatible. Routing gate visualization (not Logit Lens).
"""

import json
import torch
import gradio as gr
from huggingface_hub import hf_hub_download

model = None
tokenizer = None
DEVICE = "cpu"
MODEL_REPO = "DuoNeural/CDM-V2-TinyStories-37M"

SLOT_COLORS = [
    "#FF6B6B","#4ECDC4","#45B7D1","#96CEB4","#FECA57","#FF9FF3","#54A0FF","#5F27CD",
    "#00D2D3","#FF9F43","#C8D6E5","#8395A7","#EE5A24","#009432","#C4E538","#A3CB38",
]

EXAMPLES = [
    "Once upon a time there was a little girl named Lily.",
    "Tom loved trains more than anything else in the world.",
    "The rabbit hopped through the sunny meadow looking for",
    "In a small village there lived a clever fox named",
    "She wanted to play outside but the weather was",
]

SLOT_ROLES = {
    11: "PUNCT",   # 0-indexed — slot 11 specializes on punctuation (step5000: 71%, step30k: 88%)
    9:  "AGENCY",  # character agency verbs
    0:  "IDENTITY",# character names / pronouns
    15: "VERBS",   # action verbs
    5:  "ARTICLES",# articles & modifiers
}


def load_model():
    global model, tokenizer
    if model is not None:
        return

    from transformers import GPT2TokenizerFast
    from cdm_model_v2 import CDMConfigV2, CDMLanguageModelV2

    tokenizer = GPT2TokenizerFast.from_pretrained("gpt2")
    tokenizer.pad_token = tokenizer.eos_token

    model_pt_path = hf_hub_download(repo_id=MODEL_REPO, filename="model.pt")
    config_path   = hf_hub_download(repo_id=MODEL_REPO, filename="config.json")

    with open(config_path) as f:
        cfg_dict = json.load(f)

    cfg = CDMConfigV2(
        vocab_size=cfg_dict.get("vocab_size", 50257),
        d_model=cfg_dict.get("d_model", 384),
        n_layers=cfg_dict.get("n_layers", 8),
        n_heads=cfg_dict.get("n_heads", 8),
        n_kv_heads=cfg_dict.get("n_kv_heads", 4),
        d_ff=cfg_dict.get("d_ff", 1024),
        K=cfg_dict.get("K", 16),
        max_len=cfg_dict.get("max_len", 512),
    )

    ckpt = torch.load(model_pt_path, map_location=DEVICE, weights_only=False)
    model = CDMLanguageModelV2(cfg)
    model.load_state_dict(ckpt.get("model_state", ckpt))
    model = model.to(DEVICE)
    model.eval()


print("Loading CDM model on startup...")
try:
    load_model()
    print("Model loaded OK")
except Exception as e:
    print(f"Startup load failed (will retry on first request): {e}")


def gate_to_css_color(base_hex: str, intensity: float) -> str:
    """Blend base color with dark background by intensity (0-1)."""
    r = int(base_hex[1:3], 16)
    g = int(base_hex[3:5], 16)
    b = int(base_hex[5:7], 16)
    bg = 10
    r2 = int(bg + (r - bg) * intensity)
    g2 = int(bg + (g - bg) * intensity)
    b2 = int(bg + (b - bg) * intensity)
    return f"#{r2:02x}{g2:02x}{b2:02x}"


def generate_and_visualize(prompt: str, max_new_tokens: int, temperature: float, top_k: int):
    if model is None:
        try:
            load_model()
        except Exception as e:
            return f"[Model load failed: {e}]", "<p>Model not loaded.</p>", ""

    ids = tokenizer.encode(prompt, return_tensors="pt").to(DEVICE)
    generated_text, snapshots = model.generate_with_slots(
        ids, max_new=int(max_new_tokens), tokenizer=tokenizer,
        temperature=float(temperature), top_k=int(top_k),
    )

    K = model.cfg.K

    # ── Routing gate heatmap table ────────────────────────────────────────────
    rows = []
    slot_token_counts = [[] for _ in range(K)]  # track which tokens route to each slot

    for tok_str, all_layer_gates, winner in snapshots:
        tok_display = tok_str.replace("<","&lt;").replace(">","&gt;").replace(" ","·").strip() or "⏎"
        last_gates = all_layer_gates[-1]  # final layer gate distribution

        # Track routing for summary
        slot_token_counts[winner].append(tok_str.strip())

        # Winner badge
        win_color = SLOT_COLORS[winner % len(SLOT_COLORS)]
        winner_label = SLOT_ROLES.get(winner, f"S{winner+1}")
        tok_cell = (
            f'<td style="background:#1a1a2e;color:#e0e0e0;padding:3px 8px;'
            f'font-size:13px;font-weight:bold;border-right:2px solid #333;white-space:nowrap;">'
            f'{tok_display} '
            f'<span style="color:{win_color};font-size:10px;font-weight:normal;">→{winner_label}</span>'
            f'</td>'
        )

        gate_cells = []
        for k, gv in enumerate(last_gates):
            intensity = min(1.0, gv * K)  # scale up since each slot avg ~1/K
            bg = gate_to_css_color(SLOT_COLORS[k % len(SLOT_COLORS)], intensity)
            pct = f"{gv*100:.0f}%" if gv > 0.08 else ""
            border = f"2px solid {SLOT_COLORS[k%len(SLOT_COLORS)]}" if k == winner else "1px solid #111"
            gate_cells.append(
                f'<td style="background:{bg};color:#fff;padding:2px 5px;'
                f'font-size:10px;font-family:monospace;border:{border};'
                f'text-align:center;min-width:28px;">{pct}</td>'
            )
        rows.append(f'<tr style="border-bottom:1px solid #111;">{tok_cell}{"".join(gate_cells)}</tr>')

    # Header
    slot_headers = "".join(
        f'<th style="background:{SLOT_COLORS[k%len(SLOT_COLORS)]};color:#000;'
        f'padding:4px 3px;font-size:10px;font-weight:bold;text-align:center;min-width:28px;">'
        f'{"★" if k in SLOT_ROLES else f"S{k+1}"}</th>'
        for k in range(K)
    )
    table_html = f"""<div style="overflow-x:auto;background:#050510;border-radius:8px;padding:8px;">
<p style="color:#888;font-size:11px;margin:0 0 6px 0;">
Cell brightness = routing affinity. ★ = slots with known specialization (probe-confirmed).
→ label = winning slot for each token. <b>Watch S12 light up for punctuation!</b> (slot 11, 0-indexed)
</p>
<table style="border-collapse:collapse;width:100%;font-family:monospace;">
<thead><tr>
<th style="background:#1a1a2e;color:#aaa;padding:4px 7px;font-size:11px;text-align:left;border-right:2px solid #333;white-space:nowrap;">Token → Winner</th>
{slot_headers}</tr></thead>
<tbody>{"".join(rows)}</tbody>
</table></div>"""

    # ── Routing summary ───────────────────────────────────────────────────────
    summary_lines = ["Slot routing summary (final layer, generated tokens):", ""]
    for k in range(K):
        tokens = slot_token_counts[k]
        role = SLOT_ROLES.get(k, "")
        role_str = f" [{role}]" if role else ""
        if tokens:
            top_toks = ", ".join(f'"{t}"' for t in tokens[:6] if t)
            summary_lines.append(f"S{k+1:2d}{role_str}: {len(tokens)} wins — {top_toks}")
        else:
            summary_lines.append(f"S{k+1:2d}{role_str}: 0 wins")
    summary = "\n".join(summary_lines)

    return generated_text, table_html, summary


with gr.Blocks(title="CDM V2 — Competitive Docking Memory | DuoNeural") as demo:
    gr.Markdown("""
# 🧠 Competitive Docking Memory (CDM) — Routing Gate Visualization

**[DuoNeural](https://huggingface.co/DuoNeural) 2026 | Archon · Jesse Caldwell · Aura**

This 37M-parameter model maintains **K=16 persistent memory slots** per layer that compete to process each token.
Slots **spontaneously specialize** without supervision: slot 11 (S12) dominates punctuation, slot 0 (S1) tracks character names, etc.

The heatmap shows **routing gate affinity** — how strongly each slot claims each token.
Bright cell = that slot "owns" the token. The `→ label` shows the winning slot.

> *K_eff ≈ 15.9/16 · 99.8% Shannon Capacity Saturation · val CE 1.5934 vs 1.6516 vanilla GPT baseline*
""")

    with gr.Row():
        with gr.Column(scale=3):
            prompt_in = gr.Textbox(label="Story Prompt", placeholder="Once upon a time...",
                                   lines=2, value="Once upon a time there was a little girl named Lily.")
            with gr.Row():
                max_tokens  = gr.Slider(10, 80, value=50, step=5,    label="Max new tokens")
                temperature = gr.Slider(0.1, 1.4, value=0.85, step=0.05, label="Temperature")
                top_k       = gr.Slider(5, 80,  value=40, step=5,    label="Top-k")
            btn = gr.Button("⚡ Generate + Show Routing Gates", variant="primary")
        with gr.Column(scale=2):
            routing_summary = gr.Textbox(label="Routing Summary — which tokens each slot claimed",
                                         lines=12, interactive=False)

    generated_out = gr.Textbox(label="Generated Story", lines=4, interactive=False)
    gr.Markdown("""### Routing Gate Heatmap
*Each row = one generated token. Columns S1–S16 = memory slots (★ = probe-confirmed specializations).
Cell brightness = gate affinity. `→ label` = winning slot. Try a prompt with punctuation to watch S12!*""")
    slot_table = gr.HTML()

    gr.Examples(examples=EXAMPLES, inputs=prompt_in, label="Try these prompts")

    btn.click(fn=generate_and_visualize,
              inputs=[prompt_in, max_tokens, temperature, top_k],
              outputs=[generated_out, slot_table, routing_summary])

    gr.Markdown("""---
**Known slot specializations** (from routing probe, step 30000, last layer):
S12 (slot 11) = PUNCT · S10 (slot 9) = CHARACTER AGENCY · S1 (slot 0) = CHARACTER IDENTITY ·
S16 (slot 15) = ACTION VERBS · S6 (slot 5) = ARTICLES

**Architecture:** 37.1M params · d=384 · 8 layers · K=16 competitive memory slots · GQA · SwiGLU FFN  
**Model:** [DuoNeural/CDM-V2-TinyStories-37M](https://huggingface.co/DuoNeural/CDM-V2-TinyStories-37M)  
**Paper:** *Competitive Docking Memory: Emergent Slot Specialization in Language Models* — DuoNeural 2026""")

demo.launch()