visualizer.py

import plotly.graph_objects as go
import plotly.express as px
from plotly.subplots import make_subplots
import pandas as pd
from typing import List, Dict, Any


# Color scale for heatmap: dark (low prob) → teal (high prob)
HEATMAP_COLORSCALE = [
    [0.0,  "#1a1a2e"],
    [0.15, "#16213e"],
    [0.35, "#0f3460"],
    [0.60, "#1a6b6f"],
    [0.80, "#01696f"],
    [1.0,  "#4fc3a1"],
]

ENTROPY_LINE_COLOR = "#4fc3a1"
CONFIDENCE_COLORS = {
    "CONFIDENT": "#6daa45",
    "NEUTRAL":   "#e8af34",
    "UNCERTAIN": "#a13544",
}


def build_heatmap(results: List[Dict[str, Any]], top_k: int = 10) -> go.Figure:
    """
    Build a 2D token probability heatmap.

    X-axis: Decoding step
    Y-axis: Top-K token candidates (ranked by probability at each step)
    Cell color: Probability value

    Args:
        results: List of step result dicts from decoder
        top_k: Number of top tokens to show (Y-axis height)

    Returns:
        Plotly Figure
    """
    steps = [r["step"] for r in results]
    n_steps = len(steps)

    # Build matrices: z[token_rank][step] = prob
    z_matrix = []
    text_matrix = []
    y_labels = [f"Top {i+1}" for i in range(top_k)]

    for rank in range(top_k):
        row_z = []
        row_text = []
        for r in results:
            candidates = r["top_candidates"]
            if rank < len(candidates):
                prob = candidates[rank]["prob"]
                token = candidates[rank]["token"].replace("\n", "\\n").replace(" ", "·")
                row_z.append(prob)
                row_text.append(f"{token}<br>{prob:.4f}")
            else:
                row_z.append(0.0)
                row_text.append("—")
        z_matrix.append(row_z)
        text_matrix.append(row_text)

    step_labels = [f"Step {s}" for s in steps]

    fig = go.Figure(data=go.Heatmap(
        z=z_matrix,
        x=step_labels,
        y=y_labels,
        text=text_matrix,
        hovertemplate="<b>%{y}</b> @ %{x}<br>Token: %{text}<extra></extra>",
        colorscale=HEATMAP_COLORSCALE,
        showscale=True,
        colorbar=dict(
            title=dict(text="Probability", side="right"),
            tickformat=".3f",
            thickness=14,
        ),
        zmin=0,
        zmax=max(
            max(row) for row in z_matrix if row
        ),
    ))

    fig.update_layout(
        title=dict(
            text="Token Probability Heatmap — Top-K Candidates per Decoding Step",
            font=dict(size=14),
        ),
        xaxis=dict(
            title="Decoding Step",
            tickangle=-45,
            tickfont=dict(size=10),
            showgrid=False,
        ),
        yaxis=dict(
            title="Token Rank",
            tickfont=dict(size=10),
            autorange="reversed",
        ),
        height=420,
        margin=dict(l=80, r=80, t=60, b=80),
        paper_bgcolor="#0e1117",
        plot_bgcolor="#0e1117",
        font=dict(color="#cdccca"),
    )

    return fig


def build_entropy_chart(results: List[Dict[str, Any]]) -> go.Figure:
    """
    Build an entropy line chart with confidence-level color markers.

    X-axis: Decoding step
    Y-axis: Shannon entropy (nats)
    Markers: Colored by confidence level (green/yellow/red)

    Args:
        results: List of step result dicts from decoder

    Returns:
        Plotly Figure
    """
    steps = [r["step"] for r in results]
    entropies = [r["entropy"] for r in results]
    confidence = [r["confidence_level"] for r in results]
    tokens = [r["chosen_token"].replace("\n", "\\n") for r in results]

    marker_colors = [CONFIDENCE_COLORS[c] for c in confidence]

    fig = go.Figure()

    # Line trace
    fig.add_trace(go.Scatter(
        x=steps,
        y=entropies,
        mode="lines",
        line=dict(color=ENTROPY_LINE_COLOR, width=2),
        name="Entropy",
        showlegend=False,
    ))

    # Marker trace (colored by confidence)
    fig.add_trace(go.Scatter(
        x=steps,
        y=entropies,
        mode="markers",
        marker=dict(
            color=marker_colors,
            size=9,
            line=dict(width=1, color="#1a1a2e"),
        ),
        customdata=list(zip(tokens, confidence)),
        hovertemplate=(
            "<b>Step %{x}</b><br>"
            "Entropy: %{y:.4f} nats<br>"
            "Token: <b>%{customdata[0]}</b><br>"
            "Confidence: %{customdata[1]}<extra></extra>"
        ),
        name="Step",
        showlegend=False,
    ))

    # Threshold reference lines
    fig.add_hline(y=1.0, line_dash="dot", line_color="#6daa45",
                  annotation_text="Confident threshold", annotation_font_size=10,
                  annotation_position="bottom right")
    fig.add_hline(y=2.5, line_dash="dot", line_color="#a13544",
                  annotation_text="Uncertain threshold", annotation_font_size=10,
                  annotation_position="bottom right")

    # Legend annotations
    for label, color in CONFIDENCE_COLORS.items():
        fig.add_trace(go.Scatter(
            x=[None], y=[None],
            mode="markers",
            marker=dict(color=color, size=9),
            name=label,
            showlegend=True,
        ))

    fig.update_layout(
        title=dict(
            text="Entropy per Decoding Step — Model Confidence Over Time",
            font=dict(size=14),
        ),
        xaxis=dict(
            title="Decoding Step",
            showgrid=True,
            gridcolor="#262523",
            tickmode="linear",
            dtick=max(1, len(steps) // 10),
        ),
        yaxis=dict(
            title="Entropy (nats)",
            showgrid=True,
            gridcolor="#262523",
        ),
        legend=dict(
            orientation="h",
            yanchor="bottom",
            y=1.02,
            xanchor="right",
            x=1,
            font=dict(size=11),
        ),
        height=320,
        margin=dict(l=60, r=60, t=60, b=50),
        paper_bgcolor="#0e1117",
        plot_bgcolor="#0e1117",
        font=dict(color="#cdccca"),
    )

    return fig


def build_step_bar(step_result: Dict[str, Any]) -> go.Figure:
    """
    Build a horizontal bar chart for top-K candidates at a single step.
    Used in the step detail panel.

    Args:
        step_result: Single step dict from decoder results

    Returns:
        Plotly Figure
    """
    candidates = step_result["top_candidates"]
    tokens = [c["token"].replace("\n", "\\n").replace(" ", "·") for c in candidates]
    probs = [c["prob"] for c in candidates]

    # Color the chosen token (rank 0) differently
    colors = [ENTROPY_LINE_COLOR if i == 0 else "#393836" for i in range(len(tokens))]

    fig = go.Figure(go.Bar(
        x=probs,
        y=tokens,
        orientation="h",
        marker_color=colors,
        hovertemplate="<b>%{y}</b><br>Prob: %{x:.6f}<extra></extra>",
    ))

    fig.update_layout(
        title=dict(
            text=f"Step {step_result['step']} — Top-{len(candidates)} Candidates",
            font=dict(size=13),
        ),
        xaxis=dict(
            title="Probability",
            showgrid=True,
            gridcolor="#262523",
        ),
        yaxis=dict(
            autorange="reversed",
            tickfont=dict(size=11),
        ),
        height=350,
        margin=dict(l=100, r=40, t=50, b=50),
        paper_bgcolor="#0e1117",
        plot_bgcolor="#0e1117",
        font=dict(color="#cdccca"),
    )

    return fig