"""echo-1 Stage 0 explainer (Gradio).

Hide a block of a real Lua program and let two from-scratch models fill it back
in: a block-diffusion model (parallel, refined over steps) and an autoregressive
baseline (left to right). Every fill runs through the real Lua interpreter against
held-out tests, so the green/red badge is proof it executes, not a guess.

Run:         .venv/bin/python -m viz.app
Public link: ECHO_SHARE=1 .venv/bin/python -m viz.app
"""

from __future__ import annotations

import html
import json
import os
import subprocess
import time

import numpy as np
import torch
import torch.nn.functional as F
import gradio as gr

from ml.ar import build_prompt
from ml.config import ModelConfig, TaskConfig
from ml.data import _ids_canvas, make_block_lua
from ml.model import Transformer
from ml.tokenizer import Tokenizer

DIFF_DIR = os.environ.get("ECHO_DIFF", "runs/ediff")
AR_DIR = os.environ.get("ECHO_AR", "runs/ear")
EVAL_PATH = os.environ.get("ECHO_EVAL", "data/easy_eval.jsonl")
VERIFIER = os.environ.get("ECHO_VERIFIER", "./target/release/echo-data")
DEVICE = "cpu"


def load_model(run_dir):
    ckpt = torch.load(f"{run_dir}/model.pt", map_location=DEVICE, weights_only=False)
    mcfg = ModelConfig(**ckpt["model_cfg"])
    model = Transformer(mcfg, causal=(ckpt["mode"] == "ar")).to(DEVICE)
    model.load_state_dict(ckpt["model"])
    model.eval()
    tok = Tokenizer.load(f"{run_dir}/tokenizer.json")
    task = TaskConfig(**ckpt["task_cfg"])
    return model, tok, task


DMODEL, TOK, TASK = load_model(DIFF_DIR)
AMODEL, ATOK, ATASK = load_model(AR_DIR)
if TASK.tile_size >= TASK.block_len:
    TASK.tile_size = max(2, TASK.block_len // 4)


def curate_examples(path, n=10):
    """A small, varied, readable set: short programs spread across features."""
    rows = []
    try:
        rows = [json.loads(l) for l in open(path)]
    except FileNotFoundError:
        pass
    picked, seen = [], set()
    # one pass preferring variety by feature, short and readable
    def score(r):
        return len(r["source"].splitlines())
    rows = [r for r in rows if 5 <= len(r["source"].splitlines()) <= 12]
    want = ["recursion", "table_build", "loops", "closure", None]
    for feat in want * 3:
        for r in sorted(rows, key=score):
            key = r["source"]
            if key in seen:
                continue
            f = r.get("features", {})
            if feat is None or f.get(feat):
                picked.append(r)
                seen.add(key)
                break
        if len(picked) >= n:
            break
    if not picked:
        picked = rows[:n]
    return picked


RECORDS = curate_examples(EVAL_PATH)
def label(i, r):
    f = r.get("features", {})
    tags = [k for k in ("recursion", "loops", "table_build", "closure") if f.get(k)]
    return f"#{i+1}  ({', '.join(tags) or 'simple'}, {len(r['source'].splitlines())} lines)"
CHOICES = [(label(i, r), i) for i, r in enumerate(RECORDS)]


def tests_for(idx):
    return RECORDS[idx].get("tests", []) if 0 <= idx < len(RECORDS) else []


# ---- model passes ----

@torch.no_grad()
def diffuse(pre, blk_ids, suf, n_inner):
    enc = _ids_canvas(TOK, pre, blk_ids, suf, TASK, ar=False)
    if enc is None:
        return None, None, 0.0
    ids, region, _bid, attn = (torch.from_numpy(x).to(DEVICE) for x in enc)
    pos = torch.arange(ids.size(0), device=DEVICE)
    ctx_idx = pos[attn & (~region)]
    region_idx = pos[region]
    R = region_idx.numel()
    tile = TASK.tile_size
    target = list(blk_ids)

    t0 = time.perf_counter()
    caches = DMODEL.encode_context(ids[ctx_idx].unsqueeze(0), ctx_idx.unsqueeze(0))
    cur = [TOK.mask_id] * R
    masked_g = [True] * R
    frames = [{"ids": list(cur), "masked": list(masked_g), "label": "input (everything hidden)"}]

    for bi, b0 in enumerate(range(0, R, TASK.block_len)):
        bp = region_idx[b0:b0 + TASK.block_len]
        Lb = bp.numel()
        blk = torch.full((Lb,), TOK.mask_id, dtype=torch.long, device=DEVICE)
        masked = torch.ones(Lb, dtype=torch.bool, device=DEVICE)
        for inner in range(n_inner):
            logits, _ = DMODEL.decode_block(blk.unsqueeze(0), bp.unsqueeze(0), caches)
            conf, pred = F.softmax(logits[0], -1).max(-1)
            blk = torch.where(masked, pred, blk)
            masked = torch.zeros(Lb, dtype=torch.bool, device=DEVICE)
            for k in range(Lb):
                cur[b0 + k] = int(blk[k]); masked_g[b0 + k] = False
            frames.append({"ids": list(cur), "masked": list(masked_g),
                           "label": f"block {bi+1}, refine step {inner+1} of {n_inner}"})
            if inner == n_inner - 1:
                break
            fm = 1.0 - (inner + 1) / n_inner
            nt = (Lb + tile - 1) // tile
            keep = round(nt * fm)
            if keep <= 0:
                continue
            tc = torch.stack([conf[t*tile:(t+1)*tile].mean() for t in range(nt)])
            for t in torch.argsort(tc)[:keep].tolist():
                lo, hi = t*tile, min((t+1)*tile, Lb)
                blk[lo:hi] = TOK.mask_id; masked[lo:hi] = True
                for k in range(lo, hi):
                    masked_g[b0 + k] = True
        _, caches = DMODEL.decode_block(blk.unsqueeze(0), bp.unsqueeze(0), caches)
    lat = (time.perf_counter() - t0) * 1000
    frames.append({"ids": list(cur), "masked": [False]*R, "label": "final"})
    recon = TOK.decode(list(pre) + cur + list(suf))
    return frames, {"pre": pre, "suf": suf, "target": target, "recon": recon}, lat


@torch.no_grad()
def ar_fill(pre, suf):
    head = build_prompt(ATOK, pre, suf, ATASK)
    if head is None:
        return None, 0.0
    t0 = time.perf_counter()
    out = AMODEL.generate(torch.tensor(head, device=DEVICE), max_new=ATASK.max_decode, eos_id=ATOK.eos_id)
    lat = (time.perf_counter() - t0) * 1000
    return out, lat


def verify(source, tests):
    if not tests or not os.path.exists(VERIFIER):
        return None
    try:
        p = subprocess.run([VERIFIER, "verify-batch"],
                           input=json.dumps({"source": source, "tests": tests}),
                           capture_output=True, text=True, timeout=20)
        for line in p.stdout.splitlines():
            if line.strip():
                return json.loads(line)["pass"]
    except Exception:
        return None
    return None


# ---- readable layout: lay Lua tokens onto indented lines, colour the region ----

PRE = ("font:13px/1.7 ui-monospace,monospace;white-space:pre;background:#1e1e1e;"
       "color:#9aa;padding:14px;border-radius:8px;overflow-x:auto")
STARTERS = {"local", "return", "for", "while", "if", "function", "repeat"}
CLOSERS = {"end", "else", "elseif", "until"}


def _chip(txt, status):
    e = html.escape(txt)
    if status == "ctx":
        return e
    if status == "mask":
        return "<span style='background:#3a3a3a;color:#777;border-radius:3px;padding:0 3px'>__</span>"
    bg = "#2f9e57" if status == "ok" else "#c0392b"
    return f"<span style='background:{bg};color:#fff;border-radius:3px;padding:0 2px'>{e}</span>"


def layout(items):
    """items: list of (text, status). status in ctx/ok/wrong/mask (mask has text None).
    Returns HTML with newlines and indentation."""
    lines, cur, indent = [], [], 0

    def flush():
        if cur:
            lines.append("  " * indent + " ".join(cur)); cur.clear()

    for text, status in items:
        if status == "mask":
            cur.append(_chip("", "mask")); continue
        if text in CLOSERS:
            flush()
            if text in ("end", "until"):
                indent = max(0, indent - 1)
        elif text in STARTERS:
            flush()
        cur.append(_chip(text, status))
        if text in ("do", "then"):
            flush(); indent += 1
        elif text == "function":
            indent += 1
        elif text == "else":
            flush()
    flush()
    body = "\n".join(lines).strip("\n")
    return f"<div style='{PRE}'>{body}</div>"


def _tok(t):
    return TOK.itos.get(int(t), "?")


def diff_items(state, frame):
    items = [(_tok(t), "ctx") for t in state["pre"]]
    ok = tot = 0
    for tid, tgt, m in zip(frame["ids"], state["target"], frame["masked"]):
        if m:
            items.append((None, "mask"))
        else:
            tot += 1; good = tid == tgt; ok += good
            items.append((_tok(tid), "ok" if good else "wrong"))
    items += [(_tok(t), "ctx") for t in state["suf"]]
    return items, ok, tot


def ar_items(pre, pred, target, suf):
    items = [(_tok(t), "ctx") for t in pre]
    ok = 0
    for k, tid in enumerate(pred):
        good = k < len(target) and tid == target[k]; ok += good
        items.append((_tok(tid), "ok" if good else "wrong"))
    items += [(_tok(t), "ctx") for t in suf]
    return items, ok, len(target)


def header(ok, tot, lat, label_txt):
    pct = 100 * ok / tot if tot else 0
    return (f"<div style='color:#bbb;font:13px ui-monospace;margin-bottom:5px'>"
            f"{label_txt} &nbsp;·&nbsp; tokens right: {ok}/{tot} ({pct:.0f}%) &nbsp;·&nbsp; {lat:.0f} ms</div>")


def badge(ok):
    if ok is None:
        t, c = "verifier off", "#777"
    elif ok:
        t, c = "runs and matches the tests", "#2f9e57"
    else:
        t, c = "wrong output when executed", "#c0392b"
    return f"<span style='background:{c};color:#fff;padding:3px 9px;border-radius:5px;font:13px ui-monospace'>{t}</span>"


# ---- callbacks ----

def show_program(idx):
    if idx is None or not (0 <= int(idx) < len(RECORDS)):
        return ""
    src = RECORDS[int(idx)]["source"]
    return f"<div style='{PRE}'>{html.escape(src)}</div>"


def run(idx, frac, n_inner, seed):
    idx = int(idx)
    src = RECORDS[idx]["source"]
    blk = make_block_lua(src, float(frac), np.random.RandomState(int(seed) + 1), TOK)
    if blk is None:
        return None, "program too small to hide a block", "", "", "", "", gr.update(maximum=1, value=0)
    pre, block, suf = blk
    tests = tests_for(idx)

    frames, dstate, dlat = diffuse(pre, block, suf, int(n_inner))
    if frames is None:
        return None, "this program is longer than the model's window", "", "", "", "", gr.update(maximum=1, value=0)
    dok = verify(dstate["recon"], tests)
    apred, alat = ar_fill(pre, suf)
    arecon = TOK.decode(list(pre) + (apred or []) + list(suf))
    aok = verify(arecon, tests)

    st = {"pre": pre, "suf": suf, "target": block, "frames": frames, "dlat": dlat}
    last = len(frames) - 1
    items, ok, tot = diff_items(st, frames[-1])
    dview = header(ok, tot, dlat, frames[-1]["label"]) + layout(items)
    aitems, aok_n, atot = ar_items(pre, apred or [], block, suf)
    aview = header(aok_n, atot, alat, "left to right") + layout(aitems)
    return (st, dview, badge(dok), aview, badge(aok),
            gr.update(maximum=last, value=last, label=f"diffusion step (drag to replay), 0 to {last}"))


def scrub(state, step):
    if not state:
        return ""
    fr = state["frames"][max(0, min(int(step), len(state["frames"]) - 1))]
    items, ok, tot = diff_items(state, fr)
    return header(ok, tot, state["dlat"], fr["label"]) + layout(items)


INTRO = """
## echo-1: diffusion vs autoregressive, on pure Lua

Two tiny models, trained from scratch only on pure Lua (no libraries, no English).
We **hide a block** of a real program; each model writes it back.

- The **diffusion** model fills the whole block at once, then **refines it over a few
  steps** (you can replay the steps).
- The **autoregressive** model writes it one token at a time, left to right.

Then the **real Lua interpreter runs each result** against held-out input/output
tests. The badge is proof it actually works, not a guess.

Colours: <span style='background:#2f9e57;color:#fff;padding:0 4px;border-radius:3px'>green</span> = same token as the original,
<span style='background:#c0392b;color:#fff;padding:0 4px;border-radius:3px'>red</span> = different,
<span style='background:#3a3a3a;color:#999;padding:0 4px;border-radius:3px'>__</span> = still hidden.

Honest note: at this tiny scale the diffusion refinement visibly helps, but the
autoregressive model is more accurate. Diffusion's edge here is speed and the fact
that it revises in parallel. The point is the method, not raw skill.
"""

CONTROLS = """
**how to read the controls**
- **program**: pick one. Tags show what it uses (recursion, loops, tables).
- **hide fraction**: how much of the program body we erase for the models to rebuild. Bigger means harder.
- **diffusion steps**: how many times the diffusion model revises the block. 1 means one shot, higher means more refinement.
- **seed**: changes which block gets hidden.
"""


def build():
    with gr.Blocks(title="echo-1 Stage 0") as demo:
        gr.Markdown(INTRO)
        with gr.Row():
            pick = gr.Dropdown(CHOICES, value=0, label="program", scale=3)
            frac = gr.Slider(0.15, 0.6, value=0.3, step=0.05, label="hide fraction",
                             info="how much of the body to erase", scale=1)
            ninner = gr.Slider(1, 8, value=4, step=1, label="diffusion steps",
                               info="how many refinement passes", scale=1)
            seed = gr.Number(value=1, label="seed", info="which block to hide", precision=0, scale=1)
        gr.Markdown("### the program")
        prog = gr.HTML()
        go = gr.Button("Hide a block and let both models fill it", variant="primary")
        with gr.Row():
            with gr.Column():
                gr.Markdown("### diffusion (parallel, refined)")
                dbadge = gr.HTML()
                dview = gr.HTML()
                step = gr.Slider(0, 1, value=0, step=1, label="diffusion step")
            with gr.Column():
                gr.Markdown("### autoregressive (left to right)")
                abadge = gr.HTML()
                aview = gr.HTML()
        gr.Markdown(CONTROLS)
        st = gr.State()
        demo.load(show_program, pick, prog)
        pick.change(show_program, pick, prog)
        go.click(run, [pick, frac, ninner, seed], [st, dview, dbadge, aview, abadge, step])
        step.change(scrub, [st, step], dview)
    return demo


if __name__ == "__main__":
    build().launch(
        server_name="0.0.0.0",
        server_port=int(os.environ.get("PORT", 7860)),
        share=bool(os.environ.get("ECHO_SHARE")),
    )