ui/tab_plot.py

# ui/tab_plot.py 
"""
Tab5: Plot — Publication-quality figure generation
- Plotly native  → interactive browser preview (12×1, full-width, fast)
- matplotlib     → PNG / PDF / SVG export (300 dpi, paper-ready)

NO nested gr.Tabs() — avoids Gradio rendering bugs with Accordion+Tabs nesting.
Two side-by-side buttons instead: ⚡ Interactive  |  🖨️ Export
"""

import os
import zipfile

import gradio as gr
import pandas as pd

from db.schema import init_db
from db.reader import get_layer_metrics, get_analyzed_models
from core.plotter import plot_single_model, plot_compare_models, save_figure
from core.plotter_plotly import plotly_single, plotly_compare

_OUT_DIR = "/tmp/wang_plots"
os.makedirs(_OUT_DIR, exist_ok=True)


# ─────────────────────────────────────────────────────────────────────────────
# DB helpers
# ─────────────────────────────────────────────────────────────────────────────

def _get_model_choices() -> list[str]:
    try:
        conn = init_db()
        df   = get_analyzed_models(conn)
        return df["model_id"].tolist() if not df.empty else []
    except Exception:
        return []


def _load_df(model_id, modality, start_layer, end_layer) -> pd.DataFrame:
    conn = init_db()
    return get_layer_metrics(
        conn,
        model_id    = model_id,
        modality    = modality if modality != "all" else None,
        layer_type  = None,
        start_layer = int(start_layer),
        end_layer   = int(end_layer),
    )


def _infer_dims(df: pd.DataFrame) -> tuple[int, int]:
    head_dim = 128
    d_model  = 5120
    if not df.empty:
        if "head_dim" in df.columns and df["head_dim"].notna().any():
            head_dim = int(df["head_dim"].dropna().median())
        if "d_model" in df.columns and df["d_model"].notna().any():
            d_model  = int(df["d_model"].dropna().median())
    return head_dim, d_model


def _short(model_id: str) -> str:
    return model_id.split("/")[-1] if "/" in model_id else model_id


def _safe_path(tag: str) -> str:
    return os.path.join(_OUT_DIR, tag.replace("/", "_").replace(" ", "_"))


def _make_zip(paths: list) -> str | None:
    valid = [p for p in paths if p and os.path.exists(p)]
    if not valid:
        return None
    zp = os.path.join(_OUT_DIR, "wang_laws_figures.zip")
    with zipfile.ZipFile(zp, "w", zipfile.ZIP_DEFLATED) as zf:
        for p in valid:
            zf.write(p, os.path.basename(p))
    return zp


# ─────────────────────────────────────────────────────────────────────────────
# Single-model handlers
# ─────────────────────────────────────────────────────────────────────────────

def gen_single_plotly(model_id, modality, start_l, end_l, show_band,
                      progress=gr.Progress()):
    if not model_id:
        return None, "Please select a model."
    progress(0.2, desc="Loading data from DB...")
    df = _load_df(model_id, modality, start_l, end_l)
    if df.empty:
        return None, f"No data for {model_id}. Run Tab 2 analysis first."
    progress(0.7, desc="Building Plotly figure...")
    fig = plotly_single(df, _short(model_id), show_band=show_band)
    status = (
        f"✅ {model_id}  |  {df['layer'].nunique()} layers  "
        f"{len(df)} head-records  |  modality={modality}"
    )
    progress(1.0)
    return fig, status


def gen_single_export(model_id, modality, start_l, end_l, show_band,
                      progress=gr.Progress()):
    if not model_id:
        return "Please select a model.", None, None, None, None, None
    progress(0.15, desc="Loading data from DB...")
    df = _load_df(model_id, modality, start_l, end_l)
    if df.empty:
        return f"No data for {model_id}.", None, None, None, None, None
    head_dim, d_model = _infer_dims(df)
    progress(0.40, desc="Rendering matplotlib figure (18×20 in, 300 dpi)...")
    import matplotlib.pyplot as plt
    fig = plot_single_model(
        df, _short(model_id),
        show_band=show_band,
        head_dim=head_dim, d_model=d_model,
    )
    progress(0.78, desc="Saving PNG / PDF / SVG...")
    base  = _safe_path(f"single_{_short(model_id)}_L{int(start_l)}-{int(end_l)}")
    paths = save_figure(fig, base)
    plt.close(fig)
    zip_p  = _make_zip(paths)
    status = (
        f"✅ Exported: {', '.join(os.path.basename(p) for p in paths)}\n"
        f"   head_dim={head_dim}  d_model={d_model}"
    )
    progress(1.0)
    png = paths[0] if len(paths) > 0 else None
    pdf = paths[1] if len(paths) > 1 else None
    svg = paths[2] if len(paths) > 2 else None
    # 6 values: status, preview(=png), png_dl, pdf_dl, svg_dl, zip
    return status, png, png, pdf, svg, zip_p


# ─────────────────────────────────────────────────────────────────────────────
# Compare handlers
# ─────────────────────────────────────────────────────────────────────────────

def gen_compare_plotly(model_a, model_b, modality, start_l, end_l,
                       show_band, show_delta, progress=gr.Progress()):
    if not model_a or not model_b:
        return None, "Please select both models."
    if model_a == model_b:
        return None, "Please select two different models."
    progress(0.15, desc="Loading Model A...")
    df_a = _load_df(model_a, modality, start_l, end_l)
    progress(0.35, desc="Loading Model B...")
    df_b = _load_df(model_b, modality, start_l, end_l)
    if df_a.empty:
        return None, f"No data for Model A ({model_a})."
    if df_b.empty:
        return None, f"No data for Model B ({model_b})."
    progress(0.65, desc="Building Plotly comparison figure...")
    fig = plotly_compare(df_a, df_b, _short(model_a), _short(model_b),
                         show_band=show_band, show_delta=show_delta)
    status = (
        f"✅ {_short(model_a)} vs {_short(model_b)}  |  "
        f"A: {len(df_a)} records  B: {len(df_b)} records  |  modality={modality}"
    )
    progress(1.0)
    return fig, status


def gen_compare_export(model_a, model_b, modality, start_l, end_l,
                       show_band, show_delta, progress=gr.Progress()):
    if not model_a or not model_b or model_a == model_b:
        return "Select two different models.", None, None, None, None, None
    progress(0.10, desc="Loading data...")
    df_a = _load_df(model_a, modality, start_l, end_l)
    df_b = _load_df(model_b, modality, start_l, end_l)
    if df_a.empty or df_b.empty:
        return "Missing data for one or both models.", None, None, None, None, None
    head_dim_a, d_model_a = _infer_dims(df_a)
    head_dim_b, d_model_b = _infer_dims(df_b)
    head_dim = (head_dim_a + head_dim_b) // 2
    d_model  = (d_model_a  + d_model_b)  // 2
    progress(0.40, desc="Rendering matplotlib figure...")
    import matplotlib.pyplot as plt
    fig = plot_compare_models(
        df_a, df_b, _short(model_a), _short(model_b),
        show_band=show_band, show_delta=show_delta,
        head_dim=head_dim, d_model=d_model,
    )
    progress(0.78, desc="Saving PNG / PDF / SVG...")
    base  = _safe_path(
        f"compare_{_short(model_a)}_vs_{_short(model_b)}_L{int(start_l)}-{int(end_l)}"
    )
    paths = save_figure(fig, base)
    plt.close(fig)
    zip_p  = _make_zip(paths)
    status = (
        f"✅ Exported: {', '.join(os.path.basename(p) for p in paths)}\n"
        f"   head_dim≈{head_dim}  d_model≈{d_model}"
    )
    progress(1.0)
    png = paths[0] if len(paths) > 0 else None
    pdf = paths[1] if len(paths) > 1 else None
    svg = paths[2] if len(paths) > 2 else None
    # 6 values: status, preview(=png), png_dl, pdf_dl, svg_dl, zip
    return status, png, png, pdf, svg, zip_p


# ─────────────────────────────────────────────────────────────────────────────
# Tab5 UI  —  NO nested gr.Tabs() inside Accordion
# ─────────────────────────────────────────────────────────────────────────────

def build_tab_plot():
    with gr.Tab("📈 Plot"):
        gr.Markdown("""
        ## Wang's Five Laws — Figures

        | Button | Engine | Speed | Output |
        |--------|--------|-------|--------|
        | ⚡ **Interactive** | Native Plotly 12×1 full-width | ~2 s | In-page, hover/zoom |
        | 🖨️ **Export** | Matplotlib 18×20 in @ 300 dpi | ~30 s | PNG · PDF · SVG download |

        > Run **Tab 2 (Analyze)** first to populate the database.
        """)

        # ── Shared controls ───────────────────────────────────────────────────
        with gr.Row():
            modality_sel = gr.Dropdown(
                ["language", "vision", "audio", "all"],
                value="language", label="Modality", scale=1,
            )
            start_l = gr.Number(value=0,  precision=0, label="Start Layer", scale=1)
            end_l   = gr.Number(value=47, precision=0, label="End Layer",   scale=1)
            show_band_chk = gr.Checkbox(
                value=True, label="Show IQR band", scale=1
            )

        gr.Markdown("---")

        # ══ Single model ══════════════════════════════════════════════════════
        with gr.Accordion("📊 Single Model", open=True):

            choices = _get_model_choices()
            single_model = gr.Dropdown(
                choices=choices,
                value=choices[0] if choices else None,
                allow_custom_value=True,
                label="Model",
                info="Refresh page after new analysis to update this list.",
            )

            # Two side-by-side buttons — no nested Tabs
            with gr.Row():
                single_plotly_btn = gr.Button(
                    "⚡ Interactive (Plotly)", variant="primary", scale=1
                )
                single_export_btn = gr.Button(
                    "🖨️ Export PNG / PDF / SVG", variant="secondary", scale=1
                )

            single_status = gr.Textbox(
                lines=2, interactive=False, label="Status"
            )

            # Interactive output — always visible, populated on demand
            single_plotly_fig = gr.Plot(label="Interactive figure")

            # Export outputs — always visible, populated on demand
            gr.Markdown("#### 🖨️ Export outputs")
            single_preview = gr.Image(
                type="filepath", label="PNG preview (click to enlarge)", height=350
            )
            with gr.Row():
                dl_s_png = gr.File(label="⬇ PNG (300 dpi)")
                dl_s_pdf = gr.File(label="⬇ PDF (vector)")
                dl_s_svg = gr.File(label="⬇ SVG (vector)")
                dl_s_zip = gr.File(label="⬇ ZIP (all formats)")

        gr.Markdown("---")

        # ══ Two-model comparison ══════════════════════════════════════════════
        with gr.Accordion("📊 Two-Model Comparison", open=False):

            with gr.Row():
                model_a = gr.Dropdown(
                    choices=choices,
                    value=choices[0] if len(choices) > 0 else None,
                    allow_custom_value=True,
                    label="Model A  (solid line)", scale=2,
                )
                model_b = gr.Dropdown(
                    choices=choices,
                    value=choices[1] if len(choices) > 1 else None,
                    allow_custom_value=True,
                    label="Model B  (dashed line)", scale=2,
                )
                show_delta_chk = gr.Checkbox(
                    value=True, label="Show Δ fill (B − A)", scale=1
                )

            with gr.Row():
                cmp_plotly_btn = gr.Button(
                    "⚡ Interactive (Plotly)", variant="primary", scale=1
                )
                cmp_export_btn = gr.Button(
                    "🖨️ Export PNG / PDF / SVG", variant="secondary", scale=1
                )

            cmp_status = gr.Textbox(
                lines=2, interactive=False, label="Status"
            )

            cmp_plotly_fig = gr.Plot(label="Interactive comparison figure")

            gr.Markdown("#### 🖨️ Export outputs")
            cmp_preview = gr.Image(
                type="filepath", label="PNG preview", height=350
            )
            with gr.Row():
                dl_c_png = gr.File(label="⬇ PNG (300 dpi)")
                dl_c_pdf = gr.File(label="⬇ PDF (vector)")
                dl_c_svg = gr.File(label="⬇ SVG (vector)")
                dl_c_zip = gr.File(label="⬇ ZIP (all formats)")

        gr.Markdown("""
        ---
        **Reading the figures**
        - **IQR band** — 25%–75% quantile across attention heads per layer.
          Narrow band → heads are consistent → model is well-organized.
        - **Dotted vertical lines** — global (K=V shared) layers (Gemma-4 only).
        - **Dashed horizontal lines** — theoretical ideals (r=1, SSR=0, α=1)
          or random baselines (cosU: 1/√d_head · cosV: 1/√d_model).
        - **Super-orthogonality** (Law 4) — cosU(Q–V) and cosU(K–V) sit *below*
          the random baseline; pretraining actively pushes V away from Q/K.
        """)

        # ── Wiring ────────────────────────────────────────────────────────────

        single_plotly_btn.click(
            fn=gen_single_plotly,
            inputs=[single_model, modality_sel, start_l, end_l, show_band_chk],
            outputs=[single_plotly_fig, single_status],
        )
        single_export_btn.click(
            fn=gen_single_export,
            inputs=[single_model, modality_sel, start_l, end_l, show_band_chk],
            outputs=[single_status, single_preview,
                     dl_s_png, dl_s_pdf, dl_s_svg, dl_s_zip],
        )
        cmp_plotly_btn.click(
            fn=gen_compare_plotly,
            inputs=[model_a, model_b, modality_sel,
                    start_l, end_l, show_band_chk, show_delta_chk],
            outputs=[cmp_plotly_fig, cmp_status],
        )
        cmp_export_btn.click(
            fn=gen_compare_export,
            inputs=[model_a, model_b, modality_sel,
                    start_l, end_l, show_band_chk, show_delta_chk],
            outputs=[cmp_status, cmp_preview,
                     dl_c_png, dl_c_pdf, dl_c_svg, dl_c_zip],
        )