import os
import sys

import torch
import torch.nn.functional as F
import gradio as gr
from tokenizers import Tokenizer
from huggingface_hub import hf_hub_download, snapshot_download
from transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer

# ── Repo IDs ───────────────────────────────────────────────────────────────────
REPO_V1 = "IvmeLabs/Ivme-Conversate-v1-Base"
REPO_V2 = "IvmeLabs/Ivme-Conversate-v2-Base"
REPO_CODER = "IvmeLabs/Ivme-Coder-v1"
REPO_DIFF_BASE = "IvmeLabs/ExpIvme-DiffusionConversate-v1"
REPO_DIFF_INSTRUCT = "IvmeLabs/ExpIvme-DiffusionConversate-v1-Instruct"

device = "cuda" if torch.cuda.is_available() else "cpu"


# ── Load v1 ────────────────────────────────────────────────────────────────────
def load_v1():
    tokenizer_path = hf_hub_download(repo_id=REPO_V1, filename="ivme_tokenizer.json")
    model_path = hf_hub_download(repo_id=REPO_V1, filename="ivme_base_ema.pt")
    model_py_path = hf_hub_download(repo_id=REPO_V1, filename="model.py")

    model_dir = os.path.dirname(model_py_path)
    if model_dir not in sys.path:
        sys.path.insert(0, model_dir)
    from model import IvmeConversate  # noqa: E402

    tok = Tokenizer.from_file(tokenizer_path)
    ckpt = torch.load(model_path, map_location=device, weights_only=False)
    cfg = ckpt["cfg"]
    cfg.attn_backend = "sdpa"

    model = IvmeConversate(cfg).to(device)
    model.load_state_dict(ckpt["model"])
    model.eval()

    max_ctx = (
        getattr(cfg, "block_size", None)
        or getattr(cfg, "n_ctx", None)
        or getattr(cfg, "max_seq_len", None)
        or getattr(cfg, "context_length", None)
        or 1024
    )
    eos_id = tok.token_to_id("<|eos|>")
    return {"kind": "ar-raw", "tokenizer": tok, "model": model, "max_ctx": max_ctx, "eos_id": eos_id}


# ── Load v2 ────────────────────────────────────────────────────────────────────
def load_v2():
    # v2's architecture code lives under a `model/` package in the repo, which
    # collides by name with v1's already-imported top-level `model` module, so
    # we must remove any cached `model` module before (re)importing v2's package.
    for mod_name in list(sys.modules):
        if mod_name == "model" or mod_name.startswith("model."):
            del sys.modules[mod_name]

    repo_local_dir = snapshot_download(REPO_V2, allow_patterns=["model/*"])
    if repo_local_dir not in sys.path:
        sys.path.insert(0, repo_local_dir)
    from model import IvmeConfig, IvmeConversateV2  # noqa: E402

    tokenizer_path = hf_hub_download(repo_id=REPO_V2, filename="tokenizer.json")
    ckpt_path = hf_hub_download(repo_id=REPO_V2, filename="ckpt_final.pt")

    tok = Tokenizer.from_file(tokenizer_path)

    torch.serialization.add_safe_globals([IvmeConfig])
    ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
    cfg = ckpt["config"]

    model = IvmeConversateV2(cfg)
    state_dict = ckpt["ema_state_dict"]
    state_dict = {k.removeprefix("_orig_mod."): v for k, v in state_dict.items()}
    model.load_state_dict(state_dict)
    model.to(device).eval()

    max_ctx = getattr(cfg, "context_len", 1024)
    eos_id = tok.token_to_id("<|endoftext|>")
    return {"kind": "ar-raw", "tokenizer": tok, "model": model, "max_ctx": max_ctx, "eos_id": eos_id}


# ── Load Coder-v1 (standard transformers AutoModelForCausalLM) ────────────────
def load_coder():
    tokenizer = AutoTokenizer.from_pretrained(REPO_CODER, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        REPO_CODER, trust_remote_code=True, dtype=torch.float32,
    ).to(device).eval()
    return {"kind": "ar-hf", "tokenizer": tokenizer, "model": model}


# ── Load diffusion base + instruct (custom masked-diffusion sampler) ──────────
def load_diffusion(repo_id, instruct):
    tokenizer = AutoTokenizer.from_pretrained(repo_id, trust_remote_code=True)
    model = AutoModel.from_pretrained(
        repo_id, trust_remote_code=True,
    ).to(device).eval()
    bundle = {
        "kind": "diffusion-instruct" if instruct else "diffusion-base",
        "tokenizer": tokenizer,
        "model": model,
        "mask_token_id": model.config.mask_token_id,
    }
    if instruct:
        bundle["user_token_id"] = model.config.user_token_id
        bundle["assistant_token_id"] = model.config.assistant_token_id
        bundle["endturn_token_id"] = model.config.endturn_token_id
    return bundle


print("Loading İvme-Conversate-v1-Base...")
V1 = load_v1()
print("Loading İvme-Conversate-v2-Base...")
V2 = load_v2()
print("Loading İvme-Coder-v1...")
CODER = load_coder()
print("Loading ExpİvmeDiffusionConversate-v1 (base)...")
DIFF_BASE = load_diffusion(REPO_DIFF_BASE, instruct=False)
print("Loading ExpİvmeDiffusionConversate-v1-Instruct...")
DIFF_INSTRUCT = load_diffusion(REPO_DIFF_INSTRUCT, instruct=True)

REGISTRY = {
    "İvme-Conversate-v2-Base (recommended)": V2,
    "İvme-Conversate-v1-Base": V1,
    "İvme-Coder-v1 (Python code)": CODER,
    "Expİvme-DiffusionConversate-v1 (experimental)": DIFF_BASE,
    "Expİvme-DiffusionConversate-v1-Instruct (experimental)": DIFF_INSTRUCT,
}

BENCH = {
    # name: (v1, v2, higher_is_better)
    "WikiText-2 byte perplexity": (2.96, 2.2250, False),
    "BLiMP (macro-avg)": (61.40, 75.09, True),
    "ARC-Easy (acc_norm)": (30.85, 39.98, True),
}


# ── Generation core: raw checkpoint AR models (v1/v2) ─────────────────────────
@torch.no_grad()
def _generate_ar_raw(bundle, prompt, max_new_tokens, temperature, top_k, repetition_penalty):
    tokenizer = bundle["tokenizer"]
    model = bundle["model"]
    max_ctx = bundle["max_ctx"]
    eos_id = bundle["eos_id"]

    prompt = prompt or ""
    input_ids = tokenizer.encode(prompt).ids
    if not input_ids:
        yield prompt
        return

    generated = torch.tensor([input_ids], device=device, dtype=torch.long)
    vocab_size = None
    response_tokens: list[int] = []
    temperature = max(float(temperature), 1e-6)

    for _ in range(int(max_new_tokens)):
        window = generated[:, -max_ctx:]
        out = model(window)

        if isinstance(out, (tuple, list)):
            out = out[0]
        elif isinstance(out, dict):
            out = out.get("logits", next(iter(out.values())))

        logits = out[:, -1, :].float()

        if vocab_size is None:
            vocab_size = logits.size(-1)

        if repetition_penalty and repetition_penalty != 1.0:
            seen = torch.unique(generated[0])
            scores = logits[0, seen]
            scores = torch.where(
                scores > 0, scores / repetition_penalty, scores * repetition_penalty
            )
            logits[0, seen] = scores

        logits = logits / temperature

        k = int(top_k)
        if k > 0:
            k = min(k, vocab_size)
            topk_vals, _ = torch.topk(logits, k)
            logits[logits < topk_vals[:, -1:]] = float("-inf")

        probs = torch.softmax(logits, dim=-1)

        if not torch.isfinite(probs).all() or probs.sum() <= 0:
            next_tok = torch.argmax(logits, dim=-1, keepdim=True)
        else:
            next_tok = torch.multinomial(probs, num_samples=1)

        tok_id = next_tok.item()
        if eos_id is not None and tok_id == eos_id:
            break

        response_tokens.append(tok_id)
        generated = torch.cat([generated, next_tok], dim=1)

        yield prompt + tokenizer.decode(response_tokens)

    if not response_tokens:
        yield prompt


# ── Generation core: HF transformers AR models (Coder-v1) ─────────────────────
@torch.no_grad()
def _generate_ar_hf(bundle, prompt, max_new_tokens, temperature, top_k, repetition_penalty):
    tokenizer = bundle["tokenizer"]
    model = bundle["model"]

    prompt = prompt or ""
    inputs = tokenizer(prompt, return_tensors="pt").to(device)
    if inputs["input_ids"].shape[1] == 0:
        yield prompt
        return

    generated = inputs["input_ids"]
    response_tokens: list[int] = []
    temperature = max(float(temperature), 1e-6)
    eos_id = tokenizer.eos_token_id

    for _ in range(int(max_new_tokens)):
        out = model(generated)
        logits = out.logits[:, -1, :].float()
        vocab_size = logits.size(-1)

        if repetition_penalty and repetition_penalty != 1.0:
            seen = torch.unique(generated[0])
            scores = logits[0, seen]
            scores = torch.where(
                scores > 0, scores / repetition_penalty, scores * repetition_penalty
            )
            logits[0, seen] = scores

        logits = logits / temperature

        k = int(top_k)
        if k > 0:
            k = min(k, vocab_size)
            topk_vals, _ = torch.topk(logits, k)
            logits[logits < topk_vals[:, -1:]] = float("-inf")

        probs = torch.softmax(logits, dim=-1)
        if not torch.isfinite(probs).all() or probs.sum() <= 0:
            next_tok = torch.argmax(logits, dim=-1, keepdim=True)
        else:
            next_tok = torch.multinomial(probs, num_samples=1)

        tok_id = next_tok.item()
        if eos_id is not None and tok_id == eos_id:
            break

        response_tokens.append(tok_id)
        generated = torch.cat([generated, next_tok], dim=1)

        yield prompt + tokenizer.decode(response_tokens)

    if not response_tokens:
        yield prompt


# ── Generation core: masked-diffusion base model (unconditional/continuation) ─
@torch.no_grad()
def _generate_diffusion_base(bundle, prompt, length, steps, temperature, gumbel_temp):
    tokenizer = bundle["tokenizer"]
    model = bundle["model"]
    mask_token_id = bundle["mask_token_id"]

    length = int(length)
    steps = max(int(steps), 1)

    prefix_ids = tokenizer.encode(prompt) if prompt else []
    prefix_len = len(prefix_ids)
    total_len = prefix_len + length

    input_ids = torch.full((1, total_len), mask_token_id, dtype=torch.long, device=device)
    if prefix_len > 0:
        input_ids[0, :prefix_len] = torch.tensor(prefix_ids, dtype=torch.long, device=device)

    response_start = prefix_len

    for step in range(steps):
        logits = model(input_ids=input_ids).logits
        probs = F.softmax(logits / max(float(temperature), 1e-6), dim=-1)
        sampled = torch.multinomial(probs.view(-1, probs.size(-1)), 1).view(1, total_len)

        # Never allow the fixed prefix to be resampled.
        is_masked = input_ids == mask_token_id
        if prefix_len > 0:
            is_masked[0, :prefix_len] = False

        n_masked = is_masked.sum().item()
        if n_masked == 0:
            break

        frac_remaining = 1.0 - (step + 1) / steps
        denom = max(1 - step / steps, 1e-6)
        n_to_unmask = min(max(1, int(n_masked * (1 - frac_remaining / denom))), n_masked)

        conf = probs.gather(-1, sampled.unsqueeze(-1)).squeeze(-1)
        log_conf = torch.log(conf.clamp(min=1e-9))
        u = torch.rand_like(conf).clamp(min=1e-9, max=1 - 1e-9)
        gumbel_noise = -torch.log(-torch.log(u))
        score = (log_conf + gumbel_temp * gumbel_noise).masked_fill(~is_masked, float("-inf"))

        topk = torch.topk(score, k=n_to_unmask, dim=-1).indices
        update_mask = torch.zeros_like(is_masked).scatter_(1, topk, True)
        input_ids = torch.where(update_mask, sampled, input_ids)

        partial = input_ids[0, response_start:].tolist()
        yield (prompt or "") + tokenizer.decode(partial)

    final = input_ids[0, response_start:].tolist()
    yield (prompt or "") + tokenizer.decode(final)


# ── Generation core: masked-diffusion instruct model (chat) ───────────────────
@torch.no_grad()
def _generate_diffusion_instruct(bundle, user_message, max_response_len, steps, temperature,
                                  gumbel_temp, presence_penalty):
    tokenizer = bundle["tokenizer"]
    model = bundle["model"]
    mask_id = bundle["mask_token_id"]
    user_id = bundle["user_token_id"]
    assistant_id = bundle["assistant_token_id"]
    endturn_id = bundle["endturn_token_id"]

    max_response_len = int(max_response_len)
    steps = max(int(steps), 1)

    prefix_ids = [user_id] + tokenizer.encode(user_message or "") + [endturn_id, assistant_id]
    input_ids = torch.tensor(
        [prefix_ids + [mask_id] * max_response_len], dtype=torch.long, device=device,
    )
    prefix_len = len(prefix_ids)
    vocab_size = model.config.vocab_size

    for step in range(steps):
        logits = model(input_ids=input_ids).logits

        if presence_penalty > 0:
            response_span = input_ids[:, prefix_len:]
            visible = response_span.masked_fill(response_span == mask_id, -1)
            counts = torch.zeros(1, vocab_size, device=device)
            valid = visible[0][visible[0] >= 0]
            if len(valid) > 0:
                counts[0].scatter_add_(0, valid, torch.ones_like(valid, dtype=torch.float))
            logits = logits - presence_penalty * counts.unsqueeze(1)

        probs = F.softmax(logits / max(float(temperature), 1e-6), dim=-1)
        sampled = torch.multinomial(probs.view(-1, probs.size(-1)), 1).view(input_ids.shape)

        is_masked = input_ids == mask_id
        n_masked = is_masked.sum().item()
        if n_masked == 0:
            break

        frac_remaining = 1.0 - (step + 1) / steps
        denom = max(1 - step / steps, 1e-6)
        n_to_unmask = min(max(1, int(n_masked * (1 - frac_remaining / denom))), n_masked)

        conf = probs.gather(-1, sampled.unsqueeze(-1)).squeeze(-1)
        log_conf = torch.log(conf.clamp(min=1e-9))
        u = torch.rand_like(conf).clamp(min=1e-9, max=1 - 1e-9)
        gumbel_noise = -torch.log(-torch.log(u))
        score = (log_conf + gumbel_temp * gumbel_noise).masked_fill(~is_masked, float("-inf"))

        topk = torch.topk(score.view(1, -1), k=n_to_unmask, dim=-1).indices
        update_mask = torch.zeros_like(is_masked).view(1, -1).scatter_(1, topk, True).view(is_masked.shape)
        input_ids = torch.where(update_mask, sampled, input_ids)

        response_tokens = input_ids[0, prefix_len:].tolist()
        if endturn_id in response_tokens:
            response_tokens = response_tokens[:response_tokens.index(endturn_id)]
        yield tokenizer.decode(response_tokens)

    response_tokens = input_ids[0, prefix_len:].tolist()
    if endturn_id in response_tokens:
        response_tokens = response_tokens[:response_tokens.index(endturn_id)]
    yield tokenizer.decode(response_tokens)


# ── Unified dispatcher used by the Playground tab ──────────────────────────────
def continue_text(model_choice, prompt, max_new_tokens, temperature, top_k, repetition_penalty,
                   diff_steps, gumbel_temp, presence_penalty):
    bundle = REGISTRY[model_choice]
    kind = bundle["kind"]

    if kind == "ar-raw":
        yield from _generate_ar_raw(bundle, prompt, max_new_tokens, temperature, top_k, repetition_penalty)
    elif kind == "ar-hf":
        yield from _generate_ar_hf(bundle, prompt, max_new_tokens, temperature, top_k, repetition_penalty)
    elif kind == "diffusion-base":
        yield from _generate_diffusion_base(
            bundle, prompt, max_new_tokens, diff_steps, temperature, gumbel_temp
        )
    elif kind == "diffusion-instruct":
        yield from _generate_diffusion_instruct(
            bundle, prompt, max_new_tokens, diff_steps, temperature, gumbel_temp, presence_penalty
        )
    else:
        yield prompt


def compare_generate(prompt, max_new_tokens, temperature, top_k, repetition_penalty):
    """Run v1 and v2 on the same prompt/settings, streaming both in parallel steps."""
    gen_v1 = _generate_ar_raw(V1, prompt, max_new_tokens, temperature, top_k, repetition_penalty)
    gen_v2 = _generate_ar_raw(V2, prompt, max_new_tokens, temperature, top_k, repetition_penalty)

    last_v1, last_v2 = prompt, prompt
    done_v1 = done_v2 = False

    while not (done_v1 and done_v2):
        if not done_v1:
            try:
                last_v1 = next(gen_v1)
            except StopIteration:
                done_v1 = True
        if not done_v2:
            try:
                last_v2 = next(gen_v2)
            except StopIteration:
                done_v2 = True
        yield last_v1, last_v2


def benchmark_table():
    rows = []
    for name, (v1, v2, higher_better) in BENCH.items():
        delta = (v2 - v1) if higher_better else (v1 - v2)
        pct = (delta / abs(v1)) * 100 if v1 else 0
        arrow = "↑" if higher_better else "↓"
        rows.append([name + f" {arrow}", f"{v1:.2f}", f"{v2:.2f}", f"{'+' if delta >= 0 else ''}{delta:.2f} ({pct:+.0f}%)"])
    return rows


# ── UI ─────────────────────────────────────────────────────────────────────────
CSS = """
body, .gradio-container { font-family: 'Inter', system-ui, sans-serif; }
#component-0 { max-width: 900px; margin: 0 auto; padding: 16px; }
footer { display: none !important; }
.ivme-output textarea { font-size: 1.02rem; line-height: 1.6; }
"""

EXAMPLES = [
    "The theory of relativity states that",
    "In the beginning, the universe was",
    "def fibonacci(n):",
    "The most important thing to remember about cooking is",
    "Once upon a time, in a small village by the sea,",
    "Python is a programming language that",
]

CODE_EXAMPLES = [
    "def fibonacci(n):",
    "class BinaryTree:",
    "import numpy as np\n\ndef normalize(",
    "# Sort a list using quicksort\ndef quicksort(arr):",
]

CHAT_EXAMPLES = [
    "Hi there, how are you?",
    "What's your favorite color?",
    "Can you help me plan my day?",
    "Tell me something interesting.",
]

DIFFUSION_KEYS = {
    "Expİvme-DiffusionConversate-v1 (experimental)",
    "Expİvme-DiffusionConversate-v1-Instruct (experimental)",
}
INSTRUCT_KEY = "Expİvme-DiffusionConversate-v1-Instruct (experimental)"
CODER_KEY = "İvme-Coder-v1 (Python code)"

MODEL_NOTES = {
    "İvme-Conversate-v2-Base (recommended)": (
        "Autoregressive base model, general text. Not instruction-tuned — continues text, doesn't chat."
    ),
    "İvme-Conversate-v1-Base": (
        "Autoregressive base model, general text (earlier version). Not instruction-tuned."
    ),
    CODER_KEY: (
        "Autoregressive base model trained only on Python source. Writes code-*shaped* text reliably; "
        "does not reliably write *correct* code. Not instruction-tuned — give it a code prefix to continue."
    ),
    "Expİvme-DiffusionConversate-v1 (experimental)": (
        "🧪 Experimental masked-diffusion model (not autoregressive). Generates a fixed-length span via "
        "iterative denoising instead of left-to-right decoding. Not instruction-tuned, no chat behavior. "
        "Weak general capability (near-chance on ARC-Easy) — expect local fluency, not coherent long-form text."
    ),
    INSTRUCT_KEY: (
        "🧪 Experimental masked-diffusion model, SFT'd for basic chat. Enter a single user message (not a "
        "free-form prompt). Known limitation per the model card: output is not reliably grammatical — "
        "locally plausible words that often don't compose into coherent sentences."
    ),
}


def on_model_change(model_choice):
    """Toggle which settings are relevant/visible and swap in the right examples + notes."""
    is_diffusion = model_choice in DIFFUSION_KEYS
    is_instruct = model_choice == INSTRUCT_KEY
    is_coder = model_choice == CODER_KEY

    if is_instruct:
        examples = CHAT_EXAMPLES
        prompt_label = "User message"
        prompt_placeholder = "Hi there, how are you?"
    elif is_coder:
        examples = CODE_EXAMPLES
        prompt_label = "Prompt (Python)"
        prompt_placeholder = "def fibonacci(n):"
    else:
        examples = EXAMPLES
        prompt_label = "Prompt"
        prompt_placeholder = "The theory of relativity states that…"

    return (
        gr.update(visible=not is_diffusion),  # AR-only settings group
        gr.update(visible=is_diffusion),  # diffusion-only settings group
        gr.update(visible=is_instruct),  # presence penalty (instruct diffusion only)
        gr.update(label=prompt_label, placeholder=prompt_placeholder),
        gr.Dataset(samples=[[e] for e in examples]),
        gr.update(value=MODEL_NOTES.get(model_choice, "")),
    )


with gr.Blocks(css=CSS, title="İvme-Conversate") as demo:
    gr.Markdown(
        "## İvme-Conversate — Tiny Language Models\n"
        "A family of sub-130M-parameter language models from IvmeLabs: autoregressive base models, "
        "a Python-only coder model, and experimental masked-diffusion models."
    )

    with gr.Tabs():
        # ── Tab 1: single-model playground with picker ──────────────────────
        with gr.Tab("Playground"):
            model_picker = gr.Dropdown(
                choices=list(REGISTRY.keys()),
                value="İvme-Conversate-v2-Base (recommended)",
                label="Model",
            )

            model_note = gr.Markdown(MODEL_NOTES["İvme-Conversate-v2-Base (recommended)"])

            prompt_box = gr.Textbox(
                label="Prompt",
                placeholder="The theory of relativity states that…",
                lines=3,
                value="The theory of relativity states that",
            )

            with gr.Row():
                gen_btn = gr.Button("Generate", variant="primary", scale=3)
                clear_btn = gr.Button("Clear", scale=1)

            output_box = gr.Textbox(
                label="Output",
                lines=12,
                show_copy_button=True,
                elem_classes="ivme-output",
                interactive=False,
            )

            example_set = gr.Examples(examples=[[e] for e in EXAMPLES], inputs=prompt_box, label="Try a prompt")

            with gr.Accordion("Settings", open=False):
                with gr.Group(visible=True) as ar_settings:
                    with gr.Row():
                        max_tokens = gr.Slider(16, 512, value=200, step=8, label="Max new tokens")
                        temperature = gr.Slider(0.1, 2.0, value=0.7, step=0.05, label="Temperature")
                    with gr.Row():
                        top_k = gr.Slider(0, 200, value=40, step=1, label="Top-k (0 = disabled)")
                        rep_penalty = gr.Slider(1.0, 2.0, value=1.15, step=0.05, label="Repetition penalty")

                with gr.Group(visible=False) as diff_settings:
                    gr.Markdown(
                        "Masked-diffusion sampling: the model denoises a fully-masked span over a fixed "
                        "number of steps rather than decoding left-to-right."
                    )
                    with gr.Row():
                        diff_length = gr.Slider(16, 256, value=96, step=8, label="Response length (tokens)")
                        diff_steps = gr.Slider(4, 64, value=32, step=2, label="Diffusion steps")
                    with gr.Row():
                        diff_temperature = gr.Slider(0.1, 2.0, value=1.0, step=0.05, label="Temperature")
                        diff_gumbel_temp = gr.Slider(0.0, 2.0, value=1.0, step=0.1, label="Gumbel temp (unmask noise)")
                    diff_presence_penalty = gr.Slider(
                        0.0, 3.0, value=1.2, step=0.1,
                        label="Presence penalty (Instruct only — suppresses repetition)",
                        visible=False,
                    )

            gen_inputs = [
                model_picker, prompt_box, max_tokens, temperature, top_k, rep_penalty,
                diff_steps, diff_gumbel_temp, diff_presence_penalty,
            ]
            # Note: for diffusion models, `max_tokens` slider doubles as response length via diff_length
            # binding below; wire diff_length into the same "max_new_tokens" slot dynamically:

            def route_generate(model_choice, prompt, max_new_tokens, temperature, top_k, repetition_penalty,
                                length, steps, d_temperature, gumbel_temp, presence_penalty):
                bundle = REGISTRY[model_choice]
                kind = bundle["kind"]
                if kind == "diffusion-base":
                    yield from _generate_diffusion_base(bundle, prompt, length, steps, d_temperature, gumbel_temp)
                elif kind == "diffusion-instruct":
                    yield from _generate_diffusion_instruct(
                        bundle, prompt, length, steps, d_temperature, gumbel_temp, presence_penalty
                    )
                elif kind == "ar-hf":
                    yield from _generate_ar_hf(bundle, prompt, max_new_tokens, temperature, top_k, repetition_penalty)
                else:
                    yield from _generate_ar_raw(bundle, prompt, max_new_tokens, temperature, top_k, repetition_penalty)

            full_inputs = [
                model_picker, prompt_box, max_tokens, temperature, top_k, rep_penalty,
                diff_length, diff_steps, diff_temperature, diff_gumbel_temp, diff_presence_penalty,
            ]
            gen_btn.click(route_generate, full_inputs, output_box)
            prompt_box.submit(route_generate, full_inputs, output_box)
            clear_btn.click(lambda: ("", ""), None, [prompt_box, output_box], queue=False)

            model_picker.change(
                on_model_change,
                inputs=model_picker,
                outputs=[ar_settings, diff_settings, diff_presence_penalty, prompt_box, example_set.dataset, model_note],
            )

        # ── Tab 2: side-by-side compare ──────────────────────────────────────
        with gr.Tab("Compare v1 vs v2"):
            gr.Markdown(
                "Run the **same prompt and settings** through both autoregressive base models at once "
                "to see the difference training data made, plus the benchmark deltas below. "
                "(Coder-v1 and the diffusion models aren't included here since they use different "
                "generation mechanics — try them individually in the Playground tab.)"
            )

            cmp_prompt = gr.Textbox(
                label="Prompt",
                lines=3,
                value="Once upon a time, there was a",
            )

            with gr.Row():
                cmp_gen_btn = gr.Button("Generate both", variant="primary", scale=3)
                cmp_clear_btn = gr.Button("Clear", scale=1)

            with gr.Row():
                cmp_out_v1 = gr.Textbox(
                    label="v1-Base",
                    lines=10,
                    show_copy_button=True,
                    elem_classes="ivme-output",
                    interactive=False,
                )
                cmp_out_v2 = gr.Textbox(
                    label="v2-Base",
                    lines=10,
                    show_copy_button=True,
                    elem_classes="ivme-output",
                    interactive=False,
                )

            gr.Examples(examples=[[e] for e in EXAMPLES], inputs=cmp_prompt, label="Try a prompt")

            with gr.Accordion("Settings", open=False):
                with gr.Row():
                    cmp_max_tokens = gr.Slider(16, 512, value=150, step=8, label="Max new tokens")
                    cmp_temperature = gr.Slider(0.1, 2.0, value=0.8, step=0.05, label="Temperature")
                with gr.Row():
                    cmp_top_k = gr.Slider(0, 200, value=50, step=1, label="Top-k (0 = disabled)")
                    cmp_rep_penalty = gr.Slider(1.0, 2.0, value=1.0, step=0.05, label="Repetition penalty")

            gr.Markdown("### Benchmark improvement, v1 → v2")
            gr.Dataframe(
                headers=["Benchmark", "v1", "v2", "Δ (v1 → v2)"],
                value=benchmark_table(),
                interactive=False,
                row_count=(len(BENCH), "fixed"),
            )

            cmp_inputs = [cmp_prompt, cmp_max_tokens, cmp_temperature, cmp_top_k, cmp_rep_penalty]
            cmp_gen_btn.click(compare_generate, cmp_inputs, [cmp_out_v1, cmp_out_v2])
            cmp_prompt.submit(compare_generate, cmp_inputs, [cmp_out_v1, cmp_out_v2])
            cmp_clear_btn.click(lambda: ("", "", ""), None, [cmp_prompt, cmp_out_v1, cmp_out_v2], queue=False)

        # ── Tab 3: diffusion chat (Instruct model, dedicated chat-style UI) ──
        with gr.Tab("Diffusion Chat (experimental)"):
            gr.Markdown(
                "### Expİvme-DiffusionConversate-v1-Instruct\n"
                "🧪 **Experimental.** A 130M-parameter masked-diffusion model, SFT'd for basic chat. "
                "Per the model card: output is **not reliably grammatical** — expect locally plausible "
                "word choice that often doesn't compose into coherent sentences. Included here in the "
                "spirit of the model card's own honesty about its limitations, not as a working assistant."
            )

            chat_input = gr.Textbox(
                label="Your message",
                placeholder="Hi there, how are you?",
                lines=2,
            )

            with gr.Row():
                chat_btn = gr.Button("Send", variant="primary", scale=3)
                chat_clear_btn = gr.Button("Clear", scale=1)

            chat_output = gr.Textbox(
                label="Assistant (diffusion-sampled)",
                lines=6,
                show_copy_button=True,
                elem_classes="ivme-output",
                interactive=False,
            )

            gr.Examples(examples=[[e] for e in CHAT_EXAMPLES], inputs=chat_input, label="Try a message")

            with gr.Accordion("Settings", open=False):
                with gr.Row():
                    chat_len = gr.Slider(16, 128, value=64, step=8, label="Max response length")
                    chat_steps = gr.Slider(4, 64, value=32, step=2, label="Diffusion steps")
                with gr.Row():
                    chat_temperature = gr.Slider(0.1, 2.0, value=1.0, step=0.05, label="Temperature")
                    chat_gumbel = gr.Slider(0.0, 2.0, value=1.0, step=0.1, label="Gumbel temp")
                chat_presence = gr.Slider(0.0, 3.0, value=1.2, step=0.1, label="Presence penalty")

            def diffusion_chat(user_message, length, steps, temperature, gumbel_temp, presence_penalty):
                yield from _generate_diffusion_instruct(
                    DIFF_INSTRUCT, user_message, length, steps, temperature, gumbel_temp, presence_penalty
                )

            chat_inputs = [chat_input, chat_len, chat_steps, chat_temperature, chat_gumbel, chat_presence]
            chat_btn.click(diffusion_chat, chat_inputs, chat_output)
            chat_input.submit(diffusion_chat, chat_inputs, chat_output)
            chat_clear_btn.click(lambda: ("", ""), None, [chat_input, chat_output], queue=False)

demo.queue().launch()