src/components/comparison.py

"""
Run Comparison Component — LLM Subject Extraction Demo

Side-by-side comparison of aggregate metrics across multiple evaluation runs.
"""

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

METRICS = {
    "Agenda Items": {
        "boundary_precision": "Boundary Precision",
        "boundary_recall": "Boundary Recall",
        "boundary_f1": "Boundary F1",
        "boundary_similarity": "Boundary Similarity",
        "bed_fmeasure": "BED F-measure",
        "segeval_pk": "Segeval Pk",
        "segeval_windowdiff": "Segeval WindowDiff",
    },
    "Subjects": {
        "boundary_precision": "Boundary Precision",
        "boundary_recall": "Boundary Recall",
        "boundary_f1": "Boundary F1",
        "boundary_similarity": "Boundary Similarity",
        "bed_fmeasure": "BED F-measure",
        "segeval_pk": "Segeval Pk",
        "segeval_windowdiff": "Segeval WindowDiff",
        "theme_accuracy": "Theme Accuracy",
        "topic_f1": "Topic F1",
    },
}

# Lower-is-better metrics (for colouring)
LOWER_IS_BETTER = {"segeval_pk", "segeval_windowdiff"}


def _build_comparison_df(
    runs_data: List[Dict[str, Any]], section: str
) -> pd.DataFrame:
    rows = []
    section_key = "agenda_items" if section == "Agenda Items" else "subjects"
    metric_keys = METRICS[section]

    for run in runs_data:
        agg = run.get("aggregate", {}).get(section_key, {})
        cfg = run.get("config", {}).get("pipeline_config", {})
        row = {
            "Run ID": run["run_id"],
            "Model": cfg.get("model_name", "—"),
            "Backend": cfg.get("backend", "—"),
        }
        for key, label in metric_keys.items():
            row[label] = agg.get(key)
        rows.append(row)

    return pd.DataFrame(rows)


def _highlight_best(df: pd.DataFrame, section: str):
    """Highlight best value in each metric column."""
    metric_labels = list(METRICS[section].values())
    numeric_cols = [c for c in metric_labels if c in df.columns]

    def _style_col(col_name: str):
        col_key = [k for k, v in METRICS[section].items() if v == col_name]
        lower_is_better = col_key and col_key[0] in LOWER_IS_BETTER
        return lower_is_better

    def apply_color(col):
        lb = _style_col(col.name)
        try:
            best = col.dropna().min() if lb else col.dropna().max()
        except Exception:
            return [""] * len(col)
        return [
            "background-color:#d5f5d5;font-weight:bold;" if v == best else ""
            for v in col
        ]

    styled = df.style
    for c in numeric_cols:
        styled = styled.apply(apply_color, subset=[c])
    return styled.format({c: "{:.4f}" for c in numeric_cols if c in df.columns})


def render_comparison(
    runs_data: List[Dict[str, Any]], run_labels: List[str]
) -> None:
    """Render the cross-run comparison page."""

    if len(runs_data) < 2:
        st.info("Select at least two runs in the sidebar to compare them.")
        return

    section = st.radio(
        "Compare",
        ["Agenda Items", "Subjects"],
        horizontal=True,
        key="compare_section",
    )

    df = _build_comparison_df(runs_data, section)

    st.markdown(f"### 📋 {section} — Metric Table")
    styled = _highlight_best(df, section)
    st.dataframe(styled, use_container_width=True)

    # ── Bar chart comparison ──────────────────────────────────────────────────
    st.markdown(f"### 📊 {section} — Bar Chart")

    metric_labels = list(METRICS[section].values())
    numeric_cols = [c for c in metric_labels if c in df.columns]

    selected_metrics = st.multiselect(
        "Select metrics to compare",
        numeric_cols,
        default=numeric_cols[:4],
        key="compare_metrics",
    )

    if not selected_metrics:
        return

    plot_df = df[["Model"] + selected_metrics].melt(
        id_vars=["Model"], var_name="Metric", value_name="Value"
    )
    fig = px.bar(
        plot_df,
        x="Metric",
        y="Value",
        color="Model",
        barmode="group",
        title=f"{section} Metrics Comparison",
    )
    fig.update_layout(height=420, margin=dict(t=40, b=30))
    st.plotly_chart(fig, use_container_width=True)

    # ── Radar chart comparison ────────────────────────────────────────────────
    st.markdown(f"### 🕸️ {section} — Radar Chart")

    radar_metrics = [m for m in selected_metrics if m not in ("Segeval Pk", "Segeval WindowDiff")]
    if len(radar_metrics) < 3:
        st.info("Select at least 3 non-segeval metrics for the radar chart.")
        return

    fig2 = go.Figure()
    color_palette = px.colors.qualitative.Set2

    for idx, row in df.iterrows():
        values = [row.get(m, 0) or 0 for m in radar_metrics]
        values.append(values[0])
        fig2.add_trace(
            go.Scatterpolar(
                r=values,
                theta=radar_metrics + [radar_metrics[0]],
                fill="toself",
                name=row["Model"],
                line_color=color_palette[idx % len(color_palette)],
                fillcolor=color_palette[idx % len(color_palette)].replace("rgb", "rgba").replace(")", ",0.18)"),
            )
        )

    fig2.update_layout(
        polar=dict(radialaxis=dict(visible=True, range=[0, 1])),
        showlegend=True,
        height=420,
        margin=dict(t=40, b=20),
    )
    st.plotly_chart(fig2, use_container_width=True)