"""
Gradiographer — live MRI viewer for Gemma-3-1B (and siblings).

Watches tracked words' ranks through each layer/head cell as tokens generate,
with threshold-based suppress/stop interventions. Two projection methods:
  - lm     : standard logit-lens (final_norm + lm_head); residual-only by default.
  - interp : projection-free readout (no lm_head); every head + Layer by default.

Presentation follows small-multiples discipline: every tracked word × method is a
panel on ONE shared color scale, so target-vs-control and lm-vs-interp are honest
visual comparisons rather than separately-autoscaled images.

Local-first, Spaces-compatible. Defaults to google/gemma-3-1b-it.
"""

from __future__ import annotations

import os
import sys
from dataclasses import dataclass

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import gradio as gr
import numpy as np
import plotly.graph_objects as go
from plotly.subplots import make_subplots

from projection_map import load_model, build_col_specs
from mri_live import (
    TrackedWord,
    InterventionConfig,
    MethodTrace,
    stream_trace,
    tokenize_word,
    word_preview,
    make_rank_threshold_trigger,
)

DEFAULT_MODEL = os.environ.get("GRADIOGRAPHER_MODEL", "google/gemma-3-1b-it")

_METHOD_COLOR = {"lm": "#4cc9f0", "interp": "#ffb703"}  # vivid on dark

# Discrete rank bands — finer near the surface, coarse for the buried tail.
# Color: Magma sampled bright(top) -> dark; 'buried' sits near the dark canvas.
BAND_LABELS = ["top", "top-10", "top-100", "top-1k", "top-2.5k", "buried"]
BAND_COLORS = ["#fcfdbf", "#fe9f6d", "#de4968", "#8c2981", "#3b0f70", "#1a1a22"]
BAND_EDGES = [1, 10, 100, 1000, 2500]  # upper bound of bands 0..4; >1000 = buried


def _bands(rank2d):
    """Map a 2D rank array to band indices 0(top)..5(buried)."""
    b = np.full(rank2d.shape, 5.0)
    b[rank2d <= 2500] = 4
    b[rank2d <= 1000] = 3
    b[rank2d <= 100] = 2
    b[rank2d <= 10] = 1
    b[rank2d <= 1] = 0
    return b


def _discrete_scale(colors):
    """Stepped colorscale so each integer band index renders a flat color."""
    n = len(colors)
    cs = []
    for i, c in enumerate(colors):
        cs.append([i / n, c])
        cs.append([(i + 1) / n, c])
    return cs


# ---------------------------------------------------------------------------
# Lazy model state
# ---------------------------------------------------------------------------

@dataclass
class ModelBundle:
    model: object
    tokenizer: object
    device: str
    num_layers: int
    num_heads: int
    num_vocab: int
    col_labels_default: list[str]


_bundle: ModelBundle | None = None


def get_bundle() -> ModelBundle:
    global _bundle
    if _bundle is None:
        model, tokenizer, device, num_layers, num_heads = load_model(DEFAULT_MODEL)
        col_specs, _, _, _ = build_col_specs(num_heads)
        default_cols = [c[2] for c in col_specs if c[0] in ("single", "layer")]
        num_vocab = int(model.lm_head.weight.shape[0])
        _bundle = ModelBundle(model, tokenizer, device, num_layers, num_heads,
                              num_vocab, default_cols)
    return _bundle


# ---------------------------------------------------------------------------
# UI helpers
# ---------------------------------------------------------------------------

def render_word_preview(words_raw: str) -> str:
    if not words_raw.strip():
        return "_(enter one or more comma-separated words — first is the focus, others read as controls)_"
    b = get_bundle()
    tok = b.tokenizer
    lines = []
    for w in [x.strip() for x in words_raw.split(",") if x.strip()]:
        prev = word_preview(tok, w)
        ws = " | ".join(f"`{repr(s)}`(id {tid})" for tid, s in prev["with_leading_space"])
        ns = " | ".join(f"`{repr(s)}`(id {tid})" for tid, s in prev["no_leading_space"])
        lines.append(f"**{w}**")
        lines.append(f"- with leading space: {ws}")
        lines.append(f"- no leading space: {ns}")
    return "\n".join(lines)


def _empty_fig(title: str = "") -> go.Figure:
    f = go.Figure()
    f.update_layout(template="plotly_dark", height=320, title=title,
                    margin=dict(l=40, r=20, t=40, b=40))
    return f


def build_heatmap_grid(hist_by_wm, words, methods, layer_axis) -> go.Figure:
    """Small multiples (rows = tracked words, cols = methods) on ONE shared
    discrete rank-band scale. Bands are absolute categories (top / top-5 / ... /
    buried), so a single legend is both honest and comparable across methods —
    no log decoding, no per-panel autoscale. Crisp tiles on a dark field: a
    surfaced concept glows like structure on a scan. Each panel collapses cols
    to min-rank per layer (best visibility of the tracked word); hover = exact rank."""
    rows, cols = len(words), len(methods)
    titles = [f"{w.label} · {m}" for w in words for m in methods]
    fig = make_subplots(rows=rows, cols=cols, subplot_titles=titles,
                        shared_xaxes=True, shared_yaxes=True,
                        horizontal_spacing=0.09, vertical_spacing=0.14)
    ylabels = [f"L{l:02d}" for l in layer_axis]
    any_data = False
    for r, _w in enumerate(words):
        for c, m in enumerate(methods):
            hist = hist_by_wm.get((r, m))
            if not hist:
                continue
            rankmin = np.stack([rm.min(axis=1) for rm in hist], axis=0).T  # (layers, steps)
            fig.add_trace(go.Heatmap(
                z=_bands(rankmin), customdata=rankmin,
                y=ylabels, x=[f"t{i}" for i in range(len(hist))],
                coloraxis="coloraxis", xgap=1.5, ygap=1.5, zsmooth=False,
                hovertemplate="L%{y} · %{x}<br>rank %{customdata}<extra></extra>",
            ), row=r + 1, col=c + 1)
            any_data = True
    if not any_data:
        return _empty_fig("rank heatmaps")
    fig.update_layout(
        template="plotly_dark",
        coloraxis=dict(
            colorscale=_discrete_scale(BAND_COLORS), cmin=-0.5, cmax=5.5,
            colorbar=dict(title="rank<br>band", tickvals=list(range(6)),
                          ticktext=BAND_LABELS, thickness=16, len=0.92, outlinewidth=0),
        ),
        height=max(320, 250 * rows),
        margin=dict(l=55, r=95, t=56, b=40),
        title="MRI — where each word surfaces across layers (x = generation step) · bright = nearer the surface",
    )
    fig.update_annotations(font_size=12)
    return fig


def build_trace_grid(step_methods, words, methods, layer_axis) -> go.Figure:
    """Rows = tracked words; lm vs interp overlaid as two lines per word so the
    divergence is read within-panel. Log-rank y, REVERSED so the surface (rank 1)
    is at the top; faint dotted band-edges echo the heatmap legend."""
    rows = len(words)
    titles = [w.label for w in words]
    fig = make_subplots(rows=rows, cols=1, subplot_titles=titles, shared_xaxes=True,
                        vertical_spacing=0.13)
    for r, _w in enumerate(words):
        for m in methods:
            mt = step_methods.get(m)
            if mt is None:
                continue
            per_layer = mt.rank_matrix(r).min(axis=1)
            fig.add_trace(go.Scatter(
                x=mt.layers_used, y=per_layer, mode="lines+markers",
                name=m, legendgroup=m, showlegend=(r == 0),
                line=dict(color=_METHOD_COLOR.get(m), width=2), marker=dict(size=5),
            ), row=r + 1, col=1)
        for e in BAND_EDGES:
            fig.add_hline(y=e, line=dict(color="rgba(255,255,255,0.10)", width=1, dash="dot"),
                          row=r + 1, col=1)
        fig.update_yaxes(type="log", autorange="reversed", title_text="rank",
                         row=r + 1, col=1)
    fig.update_xaxes(title_text="layer", row=rows, col=1)
    fig.update_layout(template="plotly_dark", height=max(320, 240 * rows),
                      margin=dict(l=55, r=20, t=50, b=40),
                      title="per-layer rank at this step — lm vs interp (higher = nearer the surface)")
    fig.update_annotations(font_size=12)
    return fig


def _fmt_rank(r: int) -> str:
    return f"{round(r/1000)}k" if r >= 1000 else str(r)


def _five_closest(mt: MethodTrace, tracked_idx: int = 0) -> str:
    if not mt or not mt.cells:
        return ""
    cells = sorted(mt.cells, key=lambda c: c.ranks[tracked_idx])[:5]
    parts = []
    for c in cells:
        loc = f"L{c.layer:02d}·R" if c.col == "Layer" else f"L{c.layer:02d}{c.col}"
        parts.append(f"{loc} {_fmt_rank(c.ranks[tracked_idx])}")
    return "|".join(parts)


def _log_line(step) -> str:
    bits = [f"t{step.step:>2}"]
    if step.fwd_ms >= 0:
        bits.append(f"fwd {step.fwd_ms:.0f}ms")
    for m in ("lm", "interp"):
        if m in step.methods:
            mt = step.methods[m]
            lat = f" {mt.proj_ms:.1f}ms" if mt.proj_ms >= 0 else ""
            bits.append(f"{m}{lat} |{_five_closest(mt)}|")
    if step.intervened:
        bits.append(f"⚠ INTERVENED→{step.intervention_action.upper()} ({','.join(step.intervention_methods)})")
    return "  ".join(bits)


# ---------------------------------------------------------------------------
# Main run generator
# ---------------------------------------------------------------------------

def run_session(
    prompt, words_raw, scan_mode, max_new_tokens, scale_mode, layers_str, cols_str,
    precise_timing,
    threshold_enabled, threshold_word_idx, threshold_layers_str, threshold_cols_str,
    threshold_rank_below, threshold_rank_above, intervention_action,
):
    b = get_bundle()
    tok = b.tokenizer

    words = [w.strip() for w in words_raw.split(",") if w.strip()]
    if not words:
        yield ("no tracked words provided", _empty_fig(), _empty_fig(), "", {}, "—")
        return
    tracked = [TrackedWord(text=w, token_ids=tokenize_word(tok, w, True), anchor_idx=0) for w in words]

    def _parse_ints(s, default):
        s = s.strip()
        if not s:
            return default
        out = []
        for part in s.split(","):
            part = part.strip()
            step = 1
            if ":" in part:
                part, step_s = part.split(":")
                step = max(1, int(step_s))
            if "-" in part:
                a, z = part.split("-")
                out.extend(range(int(a), int(z) + 1, step))
            else:
                out.append(int(part))
        return sorted(set(out))

    layer_mask = _parse_ints(layers_str, list(range(b.num_layers)))
    methods = {"lm": ("lm",), "interp": ("interp",), "both": ("lm", "interp")}[scan_mode]
    method_list = list(methods)

    cols_parsed = [c.strip() for c in cols_str.split(",") if c.strip()]
    if cols_parsed:
        lm_cols = interp_cols = cols_parsed
    else:
        lm_cols = ["Layer"]      # lm_head is expensive -> residual only
        interp_cols = None       # interp is cheap -> everything

    intervention = None
    if threshold_enabled:
        th_layers = _parse_ints(threshold_layers_str, layer_mask)
        th_cols = [c.strip() for c in threshold_cols_str.split(",") if c.strip()] or ["Layer"]
        rb = int(threshold_rank_below) if int(threshold_rank_below) > 0 else None
        ra = int(threshold_rank_above) if int(threshold_rank_above) > 0 else None
        trigger = make_rank_threshold_trigger(int(threshold_word_idx), th_layers, th_cols, rb, ra)
        intervention = InterventionConfig(trigger_fn=trigger, action=intervention_action)

    hist_by_wm: dict = {}     # (word_idx, method) -> list of (layers, cols) matrices
    trace_snaps: list = []    # per-step step.methods dict (lightweight; trace figs built lazily)
    log_lines: list = []
    generated = ""

    # Rebuilding a make_subplots figure costs ~hundreds of ms; doing it every
    # token dwarfs the model. So render the heavy figures ~12 times across the
    # run (+ a final full render), and stream only log/text on the steps between.
    total = int(max_new_tokens)
    render_every = max(1, total // 12)
    last_hm = _empty_fig("rank heatmaps")
    last_tg = _empty_fig("per-layer trace")
    layer_axis = layer_mask

    def _state(idx):
        return {"snaps": trace_snaps, "words": tracked, "methods": method_list,
                "layer_axis": layer_axis, "idx": idx}

    for i, step in enumerate(stream_trace(
        model=b.model, tokenizer=tok, device=b.device,
        num_layers=b.num_layers, num_heads=b.num_heads,
        prompt=prompt, tracked_words=tracked, max_new_tokens=total,
        layer_mask=layer_mask, lm_cols=lm_cols, interp_cols=interp_cols,
        methods=methods, scale_mode=scale_mode, intervention=intervention,
        precise_timing=bool(precise_timing),
    )):
        generated += step.token_str
        for m, mt in step.methods.items():
            for wi in range(len(tracked)):
                hist_by_wm.setdefault((wi, m), []).append(mt.rank_matrix(wi))
        trace_snaps.append(step.methods)
        log_lines.append(_log_line(step))
        layer_axis = step.methods[method_list[0]].layers_used

        if i % render_every == 0:        # throttled live frame
            last_hm = build_heatmap_grid(hist_by_wm, tracked, method_list, layer_axis)
            last_tg = build_trace_grid(step.methods, tracked, method_list, layer_axis)

        idx = len(trace_snaps) - 1
        yield ("\n".join(log_lines), last_hm, last_tg, generated, _state(idx), f"step {idx} / {idx}")

    # final full render so the finished state is complete
    if trace_snaps:
        last_hm = build_heatmap_grid(hist_by_wm, tracked, method_list, layer_axis)
        last_tg = build_trace_grid(trace_snaps[-1], tracked, method_list, layer_axis)
        idx = len(trace_snaps) - 1
        yield ("\n".join(log_lines), last_hm, last_tg, generated, _state(idx), f"step {idx} / {idx}")


# ---------------------------------------------------------------------------
# Paging — rebuild the per-step trace grid lazily from stored snapshots
# ---------------------------------------------------------------------------

def _page(state, delta):
    if not state or not state.get("snaps"):
        return _empty_fig(), "—", state
    snaps = state["snaps"]
    n = len(snaps)
    idx = max(0, min(n - 1, state.get("idx", n - 1) + delta))
    state["idx"] = idx
    fig = build_trace_grid(snaps[idx], state["words"], state["methods"], state["layer_axis"])
    return fig, f"step {idx} / {n - 1}", state


def page_prev(state):
    return _page(state, -1)


def page_next(state):
    return _page(state, +1)


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------

def build_ui():
    with gr.Blocks(title="Tokescope — interactive observability for transformer internals") as demo:
        gr.Markdown(
            "# Tokescope\n"
            "### Interactive observability for transformer internals\n"
            "Pick a word and watch where the model ranks it as it writes, at every layer. "
            "Tokescope surfaces these intermediate signals and lets them be inspected, "
            "suppressed, or stopped in real time. An alternative, experimental lightweight "
            "vocabulary-projection technique (`interp`) supports real-time exploration with "
            "significantly lower overhead than a repeated full-vocabulary decode at every layer."
        )

        with gr.Row():
            with gr.Column(scale=1):
                prompt = gr.Textbox(label="prompt", lines=3,
                                    value="Give a brief biography of Kurt Cobain and his life, no more than 50 words please.")
                words = gr.Textbox(label="tracked words (focus first, then controls)",
                                   value="suicide,elephant")
                preview = gr.Markdown()
                words.change(render_word_preview, inputs=words, outputs=preview)

                with gr.Accordion("generation & scan", open=True):
                    scan_mode = gr.Radio(
                        ["lm", "interp", "both"], value="lm", label="projection method",
                        info="lm = logit-lens (residual only by default); interp = no lm_head (every head + Layer)",
                    )
                    max_new = gr.Slider(1, 80, value=20, step=1, label="max new tokens")
                    scale_mode = gr.Radio(["raw", "full", "mean"], value="full", label="head scale mode")
                    layers_str = gr.Textbox(
                        label="layers", value="0-25:2",
                        placeholder="blank = all; 'a-z:n' = every nth (default every other, for speed)")
                    cols_str = gr.Textbox(
                        label="cols (override)", value="",
                        placeholder="blank = smart defaults (lm:Layer, interp:all heads+Layer); else forces both")
                    precise_timing = gr.Checkbox(
                        label="precise per-method timing (slower — adds GPU syncs)", value=False)

                with gr.Accordion("intervention", open=False):
                    th_enabled = gr.Checkbox(label="enable threshold", value=False)
                    th_word_idx = gr.Number(label="tracked-word index (0-based)", value=0, precision=0)
                    th_layers = gr.Textbox(label="layers to watch", value="", placeholder="blank = all")
                    th_cols = gr.Textbox(label="cols to watch", value="Layer", placeholder="e.g. Layer")
                    th_below = gr.Number(label="fire if rank <= (0 disables)", value=10, precision=0)
                    th_above = gr.Number(label="fire if rank >= (0 disables)", value=0, precision=0)
                    th_action = gr.Radio(["log", "suppress", "stop"], value="log", label="action on trigger")

                run_btn = gr.Button("run", variant="primary")

            with gr.Column(scale=2):
                heatmap_grid = gr.Plot(label="rank heatmaps — small multiples, shared scale")
                with gr.Row():
                    prev_btn = gr.Button("◀ prev step", size="sm")
                    page_label = gr.Markdown("_(run to populate)_")
                    next_btn = gr.Button("next step ▶", size="sm")
                trace_grid = gr.Plot(label="per-layer rank at selected step")
                page_state = gr.State({})

                generated = gr.Textbox(label="generated", lines=3, interactive=False)
                log = gr.Textbox(label="trace log — fwd + per-method proj latency + 5 closest cells",
                                 lines=12, interactive=False, max_lines=40)

        run_btn.click(
            run_session,
            inputs=[prompt, words, scan_mode, max_new, scale_mode, layers_str, cols_str,
                    precise_timing,
                    th_enabled, th_word_idx, th_layers, th_cols, th_below, th_above, th_action],
            outputs=[log, heatmap_grid, trace_grid, generated, page_state, page_label],
        )
        prev_btn.click(page_prev, inputs=page_state, outputs=[trace_grid, page_label, page_state])
        next_btn.click(page_next, inputs=page_state, outputs=[trace_grid, page_label, page_state])

        demo.load(render_word_preview, inputs=words, outputs=preview)

    return demo


# Module-level `demo` so the HF Spaces runner finds it without guessing.
demo = build_ui()
demo.queue()

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