"""Matplotlib figures for the polygen demo. No Gradio dependency, so the plots
can be rendered and tested independently of the UI.

Figures are built with the object-oriented ``Figure`` API rather than
``pyplot``. ``plt.subplots`` registers every figure in pyplot's global
registry; since the figures are handed to Gradio and never closed, a
long-running Space would leak a figure per request (the "more than 20 figures"
warning). ``Figure`` instances are not registered, so nothing accumulates.

Every figure renders on a transparent background with neutral-grey chrome
(``_blend``) so it reads on either a light or a dark page -- the Space honors
the viewer's color scheme, and a hardcoded white canvas would otherwise show as
a bright rectangle in dark mode, the same light-lock the HTML panels avoid.
"""

import matplotlib

matplotlib.use("Agg")
import numpy as np
from matplotlib.figure import Figure

# EXACT is indigo rather than near-black so the lossless trace stays visible on
# a dark page; the other data colors already read on both schemes.
ACCENT, EXACT_C, COARSE_C, ANOM_C = "#2b2bff", "#4f46e5", "#d08400", "#c0392b"
ORIGINAL_C = "#9aa0ad"  # neutral grey, legible on light and dark
_FG = "#8b8b9a"  # axis/title/legend chrome, tuned to read on both schemes

_SIZES_TITLE = "Stored size (log scale): raw, gzip, polygen EXACT, and the COARSE archive"


def _blend(fig: Figure, *axes) -> None:
    """Transparent canvas + neutral chrome so the figure reads on any page bg.

    Sets the figure and axes backgrounds transparent and recolors titles, axis
    labels, ticks, spines, and legend text to a mid grey that is legible on both
    light and dark schemes. Call after the legend and titles are created.
    """
    fig.patch.set_alpha(0.0)
    for ax in axes:
        ax.patch.set_alpha(0.0)
        ax.title.set_color(_FG)
        ax.xaxis.label.set_color(_FG)
        ax.yaxis.label.set_color(_FG)
        ax.tick_params(colors=_FG)
        for spine in ax.spines.values():
            spine.set_color(_FG)
        legend = ax.get_legend()
        if legend is not None:
            for text in legend.get_texts():
                text.set_color(_FG)


def fig_reconstruction(r: dict):
    """Original vs EXACT (lossless) vs COARSE (smooth fit only), channel 0."""
    fig = Figure(figsize=(8, 3.2))
    ax = fig.subplots()
    ch = r["original"][:, 0]
    view = slice(0, min(len(ch), 2000))  # readable window
    x = np.arange(len(ch))[view]
    ax.plot(x, ch[view], color=ORIGINAL_C, lw=2.4, label="original")
    ax.plot(x, r["exact"][:, 0][view], color=EXACT_C, lw=0.9, label="EXACT decode (lossless)")
    ax.plot(x, r["coarse"][:, 0][view], color=COARSE_C, lw=1.2, ls="--", label="COARSE (shape only)")
    ax.set_title("Reconstruction: EXACT is bit-faithful; COARSE is the compact shape")
    ax.set_xlabel("position")
    ax.legend(loc="upper right", fontsize=8, frameon=False)
    ax.spines[["top", "right"]].set_visible(False)
    fig.tight_layout()
    _blend(fig, ax)
    return fig


def fig_sigma(r: dict):
    """Per-segment sigma_R, with anomalies (mean + 2 sigma) highlighted."""
    fig = Figure(figsize=(8, 2.8))
    ax = fig.subplots()
    s = r["sigma_r"]
    idx = np.arange(len(s))
    thresh = float(np.mean(s) + 2 * np.std(s)) if len(s) > 1 else float("inf")
    colors = [ANOM_C if v > thresh else ACCENT for v in s]
    ax.bar(idx, s, color=colors)
    if np.isfinite(thresh):
        ax.axhline(thresh, color=ANOM_C, ls="--", lw=1, label="alert threshold")
        ax.legend(loc="upper right", fontsize=8, frameon=False)
    ax.set_title("Per-window anomaly score -- free, no extra model")
    ax.set_xlabel("window")
    ax.set_ylabel("anomaly score")
    ax.spines[["top", "right"]].set_visible(False)
    fig.tight_layout()
    _blend(fig, ax)
    return fig


def fig_sizes(r: dict, title: str = _SIZES_TITLE):
    """Stored size: raw / gzip / polygen EXACT / polygen COARSE (log scale).

    ``title`` is overridable so a modality tab (where EXACT can sit well below
    gzip) does not inherit the numeric tab's gzip-parity framing.
    """
    fig = Figure(figsize=(8, 2.8))
    ax = fig.subplots()
    labels = ["raw", "gzip", "polygen\nEXACT", "polygen\nCOARSE"]
    vals = [r["raw_size"] / 1024, r["gzip_size"] / 1024, r["token_size"] / 1024, r["coeff_size"] / 1024]
    ax.bar(labels, vals, color=[ORIGINAL_C, "#6f7682", EXACT_C, COARSE_C])
    ax.set_yscale("log")
    ax.set_ylabel("KB (log scale)")
    ax.set_title(title)
    for i, v in enumerate(vals):
        ax.text(i, v, f"{v:.1f}", ha="center", va="bottom", fontsize=8, color=_FG)
    ax.spines[["top", "right"]].set_visible(False)
    fig.tight_layout()
    _blend(fig, ax)
    return fig


def fig_sizes_archive(r: dict):
    """Stored size without a gzip baseline: the page, the lossless token, and the
    queryable analytics archive (log scale).

    Used where gzip is not the right comparison (a PDF is already a compressed
    container), so the figure shows polygen's own shrink instead.
    """
    fig = Figure(figsize=(8, 2.8))
    ax = fig.subplots()
    labels = ["page\n(raw)", "polygen\ntoken", "analytics\narchive"]
    vals = [r["raw_size"] / 1024, r["token_size"] / 1024, r["coeff_size"] / 1024]
    ax.bar(labels, vals, color=[ORIGINAL_C, EXACT_C, COARSE_C])
    ax.set_yscale("log")
    ax.set_ylabel("KB (log scale)")
    ax.set_title("Stored size: the page, the lossless token, and the queryable archive")
    for i, v in enumerate(vals):
        ax.text(i, v, f"{v:.1f}", ha="center", va="bottom", fontsize=8, color=_FG)
    ax.spines[["top", "right"]].set_visible(False)
    fig.tight_layout()
    _blend(fig, ax)
    return fig


def fig_image_pair(r: dict):
    """Original image vs the version recovered through the polygen tokens."""
    orig = np.asarray(r["original_image"])
    recon = np.asarray(r["recon_image"])
    fig = Figure(figsize=(8, 3.6))
    axes = fig.subplots(1, 2)
    for ax, im, title in (
        (axes[0], orig, "original"),
        (axes[1], recon, "recovered from tokens"),
    ):
        if im.ndim == 2:
            ax.imshow(im, cmap="gray", vmin=0, vmax=255)
        else:
            ax.imshow(im)
        ax.set_title(title, fontsize=9)
        ax.axis("off")
    psnr = r.get("psnr", float("inf"))
    quality = "lossless" if not np.isfinite(psnr) else f"{psnr:.1f} dB"
    kind = "Page" if r.get("modality") == "pdf" else "Image"
    fig.suptitle(f"{kind} through polygen tokens ({quality})", fontsize=10, color=_FG)
    fig.tight_layout()
    _blend(fig, *axes)
    return fig


def fig_spectrogram(r: dict):
    """Log-mel spectrogram features: one column per analysis frame."""
    log_mel = np.asarray(r["log_mel"]).T  # (bands, frames)
    fig = Figure(figsize=(8, 3.2))
    ax = fig.subplots()
    im = ax.imshow(log_mel, aspect="auto", origin="lower", cmap="magma")
    ax.set_title("Spectrogram features the tokenizer sees (one column per frame)")
    ax.set_xlabel("frame")
    ax.set_ylabel("frequency band")
    cbar = fig.colorbar(im, ax=ax, fraction=0.025, pad=0.01)
    cbar.ax.tick_params(colors=_FG)
    cbar.outline.set_edgecolor(_FG)
    fig.tight_layout()
    _blend(fig, ax)
    return fig


def fig_text_features(r: dict):
    """Sparse numeric feature vector extracted from the text."""
    vec = np.asarray(r["features"])
    idx = np.arange(vec.shape[0])
    nonzero = vec != 0
    fig = Figure(figsize=(8, 2.8))
    ax = fig.subplots()
    ax.bar(idx[nonzero], vec[nonzero], color=ACCENT, width=max(1, vec.shape[0] // 300))
    n_active = int(nonzero.sum())
    ax.set_title(f"Text as a numeric feature vector ({n_active} of {vec.shape[0]} active)")
    ax.set_xlabel("feature index")
    ax.set_ylabel("weight")
    ax.set_xlim(0, vec.shape[0])
    ax.spines[["top", "right"]].set_visible(False)
    fig.tight_layout()
    _blend(fig, ax)
    return fig


def fig_numeric_overview(r: dict):
    """Reconstruction (top) and the per-window anomaly score (bottom), one figure."""
    fig = Figure(figsize=(8, 4.6))
    ax_top, ax_bot = fig.subplots(2, 1)
    ch = r["original"][:, 0]
    view = slice(0, min(len(ch), 2000))
    x = np.arange(len(ch))[view]
    ax_top.plot(x, ch[view], color=ORIGINAL_C, lw=2.4, label="original")
    ax_top.plot(x, r["exact"][:, 0][view], color=EXACT_C, lw=0.9, label="EXACT (lossless)")
    ax_top.plot(x, r["coarse"][:, 0][view], color=COARSE_C, lw=1.2, ls="--", label="COARSE (shape only)")
    ax_top.set_title("Reconstruction: EXACT is bit-faithful; COARSE is the compact shape")
    ax_top.legend(loc="upper right", fontsize=8, frameon=False)
    ax_top.spines[["top", "right"]].set_visible(False)
    s = r["sigma_r"]
    idx = np.arange(len(s))
    thresh = float(np.mean(s) + 2 * np.std(s)) if len(s) > 1 else float("inf")
    colors = [ANOM_C if v > thresh else ACCENT for v in s]
    ax_bot.bar(idx, s, color=colors)
    if np.isfinite(thresh):
        ax_bot.axhline(thresh, color=ANOM_C, ls="--", lw=1, label="alert threshold")
        ax_bot.legend(loc="upper right", fontsize=8, frameon=False)
    ax_bot.set_title("Per-window anomaly score -- free, no extra model")
    ax_bot.set_xlabel("window")
    ax_bot.set_ylabel("anomaly")
    ax_bot.spines[["top", "right"]].set_visible(False)
    fig.tight_layout()
    _blend(fig, ax_top, ax_bot)
    return fig


def fig_bytes_signal(r: dict):
    """The file's raw bytes as a 1-D signal (the 'tokenize anything' view)."""
    vals = np.asarray(r["byte_values"])
    view = slice(0, min(len(vals), 4000))
    fig = Figure(figsize=(8, 3.0))
    ax = fig.subplots()
    ax.plot(np.arange(len(vals))[view], vals[view], color=ACCENT, lw=0.6)
    ax.set_title("Raw bytes as a 1-D signal -- any file tokenizes")
    ax.set_xlabel("byte index")
    ax.set_ylabel("byte value (0-255)")
    ax.set_ylim(-5, 260)
    ax.spines[["top", "right"]].set_visible(False)
    fig.tight_layout()
    _blend(fig, ax)
    return fig