"""RoPE Explorer Gradio app. Imports only from ``src/``."""

from __future__ import annotations

from html import escape

import numpy as np
import plotly.graph_objects as go
import gradio as gr
import spaces 

from src.absolute_pe import add_positional_encoding
from src.extract import (
    DEFAULT_MODEL,
    MAX_SEQ_LEN,
    MODEL_CHOICES,
    expand_kv_heads,
    extract_from_model,
    get_model_dimensions,
    get_model_sequence_info,
    random_qk,
    select_head,
)
from src.plots import (
    additive_pe_heatmaps,
    attention_bars,
    attention_heatmaps,
    bulk_before_after_delta,
    frequency_strip,
    norm_compare_add_vs_rope,
    norms_and_cosine,
    position_sweep,
    rope_angle_heatmap,
    rope_angle_slices,
    rotation_2d,
    theta_heatmap,
)
from src.rope import (
    attention_scores,
    pair_dim_labels,
    pair_frequencies,
    pair_xy,
    rotate_pair,
    theta_grid,
)

HOWTO_MD = """
# How RoPE works

After each token has a vector from the **embedding table**, attention builds two extra
vectors per token with linear layers (`q_proj`, `k_proj`):

- **Q (query)** — what this token is looking for
- **K (key)** — what this token offers as a match

Attention scores are (scaled) **dot products** `Q · K`. **RoPE rotates those Q and K
vectors in 2D planes before the dot product.** It does **not** add a position vector
onto the raw token embeddings. This app shows embeddings only as context, then focuses
on Q and K before vs after RoPE.

## Pairwise rotation

For even dimension `d`, pair `i` uses frequency

$$\\omega_i = 10000^{-2i/d},\\qquad \\theta(k,i) = k\\,\\omega_i$$

**Interleaved (paper-style)** pairing `(2i, 2i+1)`:

$$
x'_{2i} = x_{2i}\\cos\\theta - x_{2i+1}\\sin\\theta,\\qquad
x'_{2i+1} = x_{2i}\\sin\\theta + x_{2i+1}\\cos\\theta
$$

Hugging Face Llama-like models use the same frequencies but pair `(i, i + d/2)`
(`rotate_half`). Random-matrix mode uses interleaved pairing; real models use the
Llama layout so the numpy implementation can be checksummed against `rotary_emb`.

Relative positions fall out of the algebra: `R(m)^T R(n) = R(n-m)`.

Shaw relative attention (learned bias `b_{m-n}` on scores) is a **different**
mechanism and is not computed here.
"""

BULK_FORMULAS_MD = r"""
### What each Q/K matrix entry means

The selected heatmap entry is `X[h, k, j]`, where `X` is either **Q** or **K**, `h` is
the attention-head index, `k` is the token position, and `j` is the dimension within
that head.

For a real model, the first-layer projections calculate each entry as:

$$
Q^{before}_{k,h,j} = \left(H_k W_Q + b_Q\right)_{h d + j},\qquad
K^{before}_{k,h,j} = \left(H_k W_K + b_K\right)_{h d + j}
$$

Here `H_k` is the token's hidden vector, `d` is the dimension per head, and `W_Q`,
`W_K` are the model's query/key projection weights. In random-matrix mode, the initial
entries are sampled directly: `Q before ~ N(0, 1)` and `K before ~ N(0, 1)` using
independent seeds.

RoPE then transforms either matrix using

$$
\theta(k,i) = k\,base^{-2i/d}
$$

For the real-model Llama layout, the paired dimensions are `(i, i+d/2)`:

$$
X^{after}_{h,k,i} = X^{before}_{h,k,i}\cos\theta -
X^{before}_{h,k,i+d/2}\sin\theta
$$
$$
X^{after}_{h,k,i+d/2} = X^{before}_{h,k,i}\sin\theta +
X^{before}_{h,k,i+d/2}\cos\theta
$$

where `X` means either **Q** or **K**. Random-matrix mode uses adjacent pairs
`(2i, 2i+1)` instead. Therefore, **Before** contains projected or sampled values,
**After** contains rotated values, and **Delta = After − Before**.
"""

PLACEHOLDER = go.Figure().update_layout(
    title="Run **Compute** on the Setup tab first",
    template="plotly_white",
    height=320,
)

@spaces.GPU
def gpu_test():
    return "GPU available"

def _safe_slider_max(n: int) -> int:
    """Gradio sliders need max > min; keep a one-step range even at edge cases."""
    return max(int(n), 1)


def _head_mapping_markdown(total_dim: int, n_heads: int, head_dim: int) -> str:
    rows = [
        "**Projected Q/K dimensions handled by each query head**",
        "",
        "These are output dimensions after `q_proj`/`k_proj`; RoPE uses the local dimensions within each head.",
        "",
        "| Head | Projected dimensions |",
        "|---:|---:|",
    ]
    for head_index in range(n_heads):
        start = head_index * head_dim
        end = min(start + head_dim - 1, total_dim - 1)
        rows.append(f"| {head_index} | `{start}–{end}` |")
    return "\n".join(rows)


def update_dimension(source: str, model_name: str):
    if source.startswith("Random"):
        return (
            gr.update(minimum=4, maximum=128, value=32, step=2, interactive=True),
            gr.update(value=32),
            gr.update(value=1),
            gr.update(value=32),
            _head_mapping_markdown(32, 1, 32),
        )
    try:
        total_dim, n_heads, head_dim = get_model_dimensions(model_name)
        return (
            gr.update(
                minimum=4,
                maximum=max(128, head_dim),
                value=head_dim,
                step=2,
                interactive=False,
            ),
            gr.update(value=total_dim),
            gr.update(value=n_heads),
            gr.update(value=head_dim),
            _head_mapping_markdown(total_dim, n_heads, head_dim),
        )
    except Exception:
        return gr.update(), gr.update(), gr.update(), gr.update(), gr.update()


def update_sequence_length(source: str, model_name: str, sentence: str):
    if source.startswith("Random"):
        return gr.update(minimum=1, maximum=MAX_SEQ_LEN, value=16, interactive=True)
    try:
        token_count, context_limit = get_model_sequence_info(model_name, sentence)
        return gr.update(
            minimum=1,
            maximum=max(1, context_limit),
            value=token_count,
            interactive=False,
        )
    except Exception:
        return gr.update()


def update_random_dimension_display(source: str, dim: int):
    if not source.startswith("Random"):
        return gr.update(), gr.update(), gr.update(), gr.update()
    random_dim = max(4, int(dim))
    return (
        gr.update(value=random_dim),
        gr.update(value=1),
        gr.update(value=random_dim),
        _head_mapping_markdown(random_dim, 1, random_dim),
    )


def compute(
    source: str,
    sentence: str,
    model_name: str,
    seq_len: int,
    dim: int,
    seed: int,
    base: float,
    progress=gr.Progress(track_tqdm=False),
):
    try:
        if source.startswith("Random"):
            progress(0.4, desc="Sampling random Q/K")
            data = random_qk(int(seq_len), int(dim), seed=int(seed), base=float(base))
        else:
            progress(0.2, desc=f"Loading {model_name} (first time downloads weights)")
            data = extract_from_model(model_name, sentence)
        seq = int(select_head(data["q_before"], 0).shape[0])
        head_dim = int(select_head(data["q_before"], 0).shape[1])
        n_pairs = head_dim // 2
        n_heads = max(int(data["n_q_heads"]) - 1, 0)
        checksum = data["checksum"]
        if checksum is None:
            status = (
                f"Random Q/K · seq={seq} · dim={head_dim} · base={data['base']:g} · "
                f"style={data['style']}"
            )
        else:
            status = (
                f"Model `{data['model_name']}` · {seq} tokens · head_dim={head_dim} · "
                f"Q heads={data['n_q_heads']} · KV heads={data['n_kv_heads']} · "
                f"rope_theta={data['base']:g} · "
                f"max |numpy RoPE − model rotary| on Q = **{checksum:.3e}**"
            )
        token_labels = ", ".join(data["tokens"][:seq])
        status = status + f"\n\nTokens: `{token_labels}`"
        return (
            data,
            status,
            gr.update(maximum=_safe_slider_max(n_heads), value=0),
            gr.update(maximum=_safe_slider_max(seq - 1), value=0),
            gr.update(maximum=_safe_slider_max(n_pairs - 1), value=0),
            gr.update(maximum=_safe_slider_max(seq - 1), value=0),
        )
    except Exception as exc:
        return (
            None,
            f"**Error:** {exc}",
            gr.update(),
            gr.update(),
            gr.update(),
            gr.update(),
        )


def _qk_slice(data: dict, which: str, head: int):
    before = data["q_before"] if which == "Q" else data["k_before"]
    after = data["q_after"] if which == "Q" else data["k_after"]
    return select_head(before, head), select_head(after, head)


def update_bulk(data, which, head, mod_2pi):
    if not data:
        fig = PLACEHOLDER
        return fig, fig, fig, fig
    before, after = _qk_slice(data, which, int(head))
    dim = before.shape[-1]
    seq = before.shape[0]
    return (
        bulk_before_after_delta(before, after, tokens=data["tokens"]),
        norms_and_cosine(before, after, tokens=data["tokens"]),
        theta_heatmap(seq, dim, data["base"], mod_2pi=bool(mod_2pi)),
        frequency_strip(dim, data["base"]),
    )


def update_angle_explorer(data, head, token, pair, display_mode):
    if not data:
        return (
            PLACEHOLDER,
            PLACEHOLDER,
            "Compute on the Setup tab first.",
            gr.update(maximum=1, value=0),
            gr.update(maximum=1, value=0),
        )
    before = select_head(data["q_before"], int(head))
    seq, dim = before.shape
    n_pairs = dim // 2
    token = int(np.clip(token, 0, seq - 1))
    pair = int(np.clip(pair, 0, n_pairs - 1))
    base = float(data["base"])
    theta = float(theta_grid(seq, dim, base)[token, pair])
    if display_mode == "turns":
        shown_theta = theta / (2 * np.pi)
        display_label = "θ / 2π (turns)"
    elif display_mode == "wrapped":
        shown_theta = float(np.mod(theta, 2 * np.pi))
        display_label = "θ mod 2π (radians)"
    else:
        shown_theta = theta
        display_label = "θ (radians)"
    omega = float(pair_frequencies(dim, base)[pair])
    d0, d1 = pair_dim_labels(pair, dim, style=data["style"])
    detail = fr"""
### Selected rotation angle

- **Head:** `{int(head)}`
- **Token position:** `k = {token}`
- **Pair:** `i = {pair}` → ({d0}, {d1})
- **Frequency:** `ωᵢ = {omega:.8f}` radians per token position
- **Raw angle:** `θ({token}, {pair}) = {theta:.8f}` radians
- **Displayed angle ({display_label}):** `{shown_theta:.8f}`

$$\theta(k,i) = k\,base^{{-2i/d}} = {token}\,({base:g})^{{-2\times{pair}/{dim}}}$$

Every increase of one token position adds `ωᵢ` radians for this pair. Lower-frequency
pairs change more slowly as `k` increases.
"""
    return (
        rope_angle_heatmap(seq, dim, base, token, pair, display_mode),
        rope_angle_slices(seq, dim, base, token, pair, display_mode),
        detail,
        gr.update(maximum=_safe_slider_max(seq - 1), value=token),
        gr.update(maximum=_safe_slider_max(n_pairs - 1), value=pair),
    )


def update_individual(data, which, head, token, pair, sweep):
    if not data:
        return "Compute on the Setup tab first.", PLACEHOLDER, PLACEHOLDER
    before, after = _qk_slice(data, which, int(head))
    token = int(np.clip(token, 0, before.shape[0] - 1))
    n_pairs = before.shape[1] // 2
    pair = int(np.clip(pair, 0, n_pairs - 1))
    style = data["style"]
    xb, yb = pair_xy(before, token, pair, style=style)
    xa, ya = pair_xy(after, token, pair, style=style)
    theta = float(theta_grid(before.shape[0], before.shape[1], data["base"])[token, pair])
    cos_t, sin_t = float(np.cos(theta)), float(np.sin(theta))
    xe_chk, xo_chk = rotate_pair(np.array([xb]), np.array([yb]), np.array([theta]))
    d0, d1 = pair_dim_labels(pair, before.shape[1], style=style)
    table = f"""
### Token `{token}` · pair `{pair}` (`{d0}`, `{d1}`)

| | {d0} | {d1} |
|---|---:|---:|
| before | {xb:.6f} | {yb:.6f} |
| after | {xa:.6f} | {ya:.6f} |
| check (`rotate_pair`) | {float(xe_chk):.6f} | {float(xo_chk):.6f} |

**θ(k,i) = {theta:.6f} rad** · cos = {cos_t:.6f} · sin = {sin_t:.6f}

`x'_even = x_even cos θ − x_odd sin θ`  
`x'_odd  = x_even sin θ + x_odd cos θ`
"""
    neighbors = [n for n in (token - 1, token + 1, token + 2) if 0 <= n < before.shape[0]]
    rot = rotation_2d(before, after, token, pair, style, theta, neighbor_tokens=neighbors)
    if sweep:
        omega = float(pair_frequencies(before.shape[1], data["base"])[pair])
        sweep_fig = position_sweep(xb, yb, omega, before.shape[0], token)
    else:
        sweep_fig = PLACEHOLDER
        sweep_fig.update_layout(title="Enable “replay same pair at every k” to see position-only spin")
    return table, rot, sweep_fig


def _attention_context_markdown(data: dict | None) -> str:
    if not data:
        return "Compute on the Setup tab to see the sentence and tokenization."
    text = data.get("text") or "Random matrix mode does not use a sentence."
    token_lines = " | ".join(f"{i}: {token}" for i, token in enumerate(data["tokens"]))
    return (
        "### Input sentence and tokens\n\n"
        f"**Sentence:** <code>{escape(str(text))}</code>\n\n"
        f"**Tokenized form (index: token):** <code>{escape(token_lines)}</code>"
    )


def update_attention(data, head, query_token):
    if not data:
        return PLACEHOLDER, PLACEHOLDER, "", _attention_context_markdown(None)
    q_b = select_head(data["q_before"], int(head))
    q_a = select_head(data["q_after"], int(head))
    k_b_all = expand_kv_heads(data["k_before"], data["n_q_heads"])
    k_a_all = expand_kv_heads(data["k_after"], data["n_q_heads"])
    k_b = select_head(k_b_all, int(head))
    k_a = select_head(k_a_all, int(head))
    sb = attention_scores(q_b, k_b)
    sa = attention_scores(q_a, k_a)
    qt = int(np.clip(query_token, 0, sb.shape[0] - 1))
    note = (
        "Additive PE changes values by **addition**. RoPE encodes **relative** offset "
        "because `R(m)^T R(n) = R(n−m)`: the score depends on the position difference, "
        "not on absolute indices alone."
    )
    return (
        attention_heatmaps(sb, sa, tokens=data["tokens"]),
        attention_bars(sb[qt], sa[qt], qt, tokens=data["tokens"]),
        note,
        _attention_context_markdown(data),
    )


def update_compare(data):
    if not data:
        return PLACEHOLDER, PLACEHOLDER, ""
    emb = np.asarray(data["embeddings"], dtype=np.float64)
    # Compare tab always uses additive PE on the embedding matrix (may be wider than a head).
    pe, combined = add_positional_encoding(emb, base=data["base"])
    q_b = select_head(data["q_before"], 0)
    q_a = select_head(data["q_after"], 0)
    heat = additive_pe_heatmaps(emb, pe, combined)
    norms = norm_compare_add_vs_rope(emb, combined, q_b, q_a)
    copy = """
**Absolute sinusoidal PE** *adds* a position-shaped vector, so both **norm and direction** change.

**RoPE** *rotates* query/key pairs: **norm stays**, and the relative angle depends on `m − n`.

Shaw-style relative attention (`q_m^T k_n + b_{m-n}`) is a third, learned-bias mechanism — not shown as a plot.
"""
    return heat, norms, copy


def toggle_source(source: str):
    is_random = source.startswith("Random")
    return (
        gr.update(visible=True, interactive=is_random),
        gr.update(visible=True, interactive=is_random),
        gr.update(visible=is_random),
        gr.update(visible=not is_random),
        gr.update(visible=not is_random),
    )


with gr.Blocks(title="RoPE Explorer") as demo:
    state = gr.State(None)
    gr.Markdown("# RoPE Explorer")
    gr.Markdown(
        "Interactive view of **Rotary Position Embedding**: random Q/K matrices or "
        "query/key vectors from a small ungated Hugging Face model."
    )
    with gr.Tabs():
        with gr.Tab("How RoPE works"):
            gr.Markdown(HOWTO_MD)

        with gr.Tab("Setup"):
            source = gr.Radio(
                ["Random matrix", "Real model"],
                value="Random matrix",
                label="Source",
            )
            with gr.Row():
                sentence = gr.Textbox(
                    value="RoPE rotates query and key vectors.",
                    label="Sentence (real model)",
                    visible=False,
                )
                model_name = gr.Dropdown(
                    MODEL_CHOICES,
                    value=DEFAULT_MODEL,
                    label="Model (ungated, Llama-like)",
                    visible=False,
                )
            with gr.Row():
                seq_len = gr.Slider(1, MAX_SEQ_LEN, value=16, step=1, label="Sequence length")
                dim = gr.Slider(4, 128, value=32, step=2, label="Dimension (even; per attention head)")
                seed = gr.Number(value=42, label="Seed", precision=0)
            with gr.Row():
                total_dim = gr.Number(value=32, label="Total dimension", precision=0, interactive=False)
                attention_heads = gr.Number(value=1, label="Attention heads", precision=0, interactive=False)
                head_dim = gr.Number(value=32, label="Dimension per attention head", precision=0, interactive=False)
            gr.Markdown(
                "**Why these numbers differ:** `total dimension = attention heads × dimension per head`. "
                "RoPE rotates each query/key head separately, so its Dimension slider uses "
                "the per-head value, not the model's total dimension."
            )
            head_mapping = gr.Markdown(_head_mapping_markdown(32, 1, 32))
            base = gr.Number(
                value=10000,
                label="RoPE base (overridden by config.rope_theta for real models)",
            )
            compute_btn = gr.Button("Compute", variant="primary")
            status = gr.Markdown("Choose a source and click Compute.")
            head = gr.Slider(minimum=0, maximum=2, step=1, value=0, label="Head index (real models)")

        with gr.Tab("Bulk changes"):
            gr.Markdown(BULK_FORMULAS_MD)
            which = gr.Radio(["Q", "K"], value="Q", label="Tensor")
            mod_2pi = gr.Checkbox(False, label="θ heatmap: wrap mod 2π")
            bulk_main = gr.Plot(label="Before / after / delta")
            bulk_norm = gr.Plot(label="Norms and cosine")
            bulk_theta = gr.Plot(label="θ(k, i)")
            bulk_freq = gr.Plot(label="ω_i")

        with gr.Tab("RoPE angles"):
            gr.Markdown(
                "Explore how the rotation angle changes with token position `k` and "
                "dimension pair `i`. The selected head is shared with the Setup tab. "
                "A marker identifies the selected `(k, i)` cell in the heatmap."
            )
            with gr.Row():
                angle_token = gr.Slider(0, 15, step=1, value=0, label="Token position k")
                angle_pair = gr.Slider(0, 15, step=1, value=0, label="Pair index i")
                angle_display = gr.Radio(
                    [
                        ("Absolute angle (radians)", "absolute"),
                        ("Angle / 2π (turns)", "turns"),
                        ("Wrapped angle mod 2π", "wrapped"),
                    ],
                    value="absolute",
                    label="Angle display",
                )
            angle_heatmap = gr.Plot(label="RoPE angle heatmap")
            angle_slices = gr.Plot(label="Selected pair/token slices")
            angle_detail = gr.Markdown("Compute on the Setup tab first.")

        with gr.Tab("Individual changes"):
            with gr.Row():
                token_k = gr.Slider(0, 15, step=1, value=0, label="Token index k")
                pair_i = gr.Slider(0, 15, step=1, value=0, label="Pair index i")
            sweep = gr.Checkbox(True, label="Replay the same content pair at every position k")
            pair_table = gr.Markdown()
            pair_plot = gr.Plot()
            sweep_plot = gr.Plot()

        with gr.Tab("Attention effect"):
            attention_context = gr.Markdown(
                "Compute on the Setup tab to see the sentence and tokenization."
            )
            query_token = gr.Slider(0, 15, step=1, value=0, label="Query token")
            attn_heat = gr.Plot()
            attn_bar = gr.Plot()
            attn_note = gr.Markdown()
            gr.Markdown(
                """
### How to read the Attention logits bars

The selected query token is compared with every key token. Each bar is the raw
dot product `Q_query · K_key` for one key position:

- A **higher bar** means the key is a stronger match for this query relative to the other keys in the same chart.
- A **lower or negative bar** means a weaker or opposing match.
- The gap between **without RoPE** and **with RoPE** shows how position-aware rotation changes that comparison.
- These bars are **logits, not probabilities**. Applying softmax across all bars for one query would convert them into attention weights.

Use the hover text to see the exact key token and position behind each bar. The
query token is the position shown in the chart title.
"""
            )
            gr.Markdown(
                """
### How to read an attention score

Each heatmap cell is the raw dot product `Q_query · K_key` for the query token on
the y-axis and key token on the x-axis:

- **Zero** means the two vectors are orthogonal, so this query/key pair has no directional match.
- **Positive** means the vectors point partly in the same direction, indicating a compatible match.
- **Negative** means the vectors point partly in opposite directions, indicating an incompatible match.
- **Magnitude** shows how strong the alignment or opposition is. Larger absolute values mean a stronger raw signal.

These are raw, unnormalized scores, not probabilities. Compare scores within the
same query row; the model would apply softmax across that row to turn them into
relative attention weights. Vector lengths also affect the magnitude, so a larger
score does not represent a universal threshold of importance.
"""
            )

        with gr.Tab("Compare to additive PE"):
            pe_heat = gr.Plot()
            pe_norm = gr.Plot()
            pe_note = gr.Markdown()

    compute_btn.click(
        compute,
        inputs=[source, sentence, model_name, seq_len, dim, seed, base],
        outputs=[state, status, head, token_k, pair_i, query_token],
    )
    source.change(
        toggle_source,
        inputs=[source],
        outputs=[seq_len, dim, seed, sentence, model_name],
    )
    source.change(
        update_dimension,
        inputs=[source, model_name],
        outputs=[dim, total_dim, attention_heads, head_dim, head_mapping],
    )
    source.change(
        update_sequence_length,
        inputs=[source, model_name, sentence],
        outputs=[seq_len],
    )
    model_name.change(
        update_dimension,
        inputs=[source, model_name],
        outputs=[dim, total_dim, attention_heads, head_dim, head_mapping],
    )
    model_name.change(
        update_sequence_length,
        inputs=[source, model_name, sentence],
        outputs=[seq_len],
    )
    sentence.change(
        update_sequence_length,
        inputs=[source, model_name, sentence],
        outputs=[seq_len],
    )
    dim.change(
        update_random_dimension_display,
        inputs=[source, dim],
        outputs=[total_dim, attention_heads, head_dim, head_mapping],
    )

    bulk_inputs = [state, which, head, mod_2pi]
    bulk_outputs = [bulk_main, bulk_norm, bulk_theta, bulk_freq]
    for ctrl in bulk_inputs:
        ctrl.change(update_bulk, inputs=bulk_inputs, outputs=bulk_outputs)

    angle_inputs = [state, head, angle_token, angle_pair, angle_display]
    angle_outputs = [angle_heatmap, angle_slices, angle_detail, angle_token, angle_pair]
    for ctrl in angle_inputs:
        ctrl.change(update_angle_explorer, inputs=angle_inputs, outputs=angle_outputs)

    ind_inputs = [state, which, head, token_k, pair_i, sweep]
    ind_outputs = [pair_table, pair_plot, sweep_plot]
    for ctrl in ind_inputs:
        ctrl.change(update_individual, inputs=ind_inputs, outputs=ind_outputs)

    attn_inputs = [state, head, query_token]
    attn_outputs = [attn_heat, attn_bar, attn_note, attention_context]
    for ctrl in attn_inputs:
        ctrl.change(update_attention, inputs=attn_inputs, outputs=attn_outputs)

    state.change(update_compare, inputs=[state], outputs=[pe_heat, pe_norm, pe_note])

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)