src/components/statistics.py

"""
Statistics Panel — LLM Subject Extraction Demo

Shows aggregate and per-municipality metrics for a selected evaluation run.
"""

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


def _fmt(val, decimals: int = 3) -> str:
    if isinstance(val, (int, float)):
        return f"{val:.{decimals}f}"
    return str(val)


def render_statistics(run_data: Dict[str, Any]) -> None:
    """Render the statistics dashboard for a loaded run."""

    aggregate = run_data.get("aggregate", {})
    config = run_data.get("config", {})
    documents = run_data.get("documents", {})

    # ── Run Info ──────────────────────────────────────────────────────────────
    st.markdown("### ⚙️ Run Configuration")

    pcfg = config.get("pipeline_config", {})
    cfg_cols = st.columns(4)
    cfg_cols[0].metric("Backend", pcfg.get("backend", "—"))
    cfg_cols[1].metric("Model", pcfg.get("model_name", "—"))
    cfg_cols[2].metric("Temperature", pcfg.get("temperature", "—"))
    cfg_cols[3].metric("Split", config.get("split", "—"))

    extra_cols = st.columns(4)
    extra_cols[0].metric("Few-shot", str(pcfg.get("use_few_shot", "—")))
    extra_cols[1].metric("Fix Gaps", str(pcfg.get("fix_gaps", "—")))
    extra_cols[2].metric("Max Tokens", pcfg.get("max_tokens", "—"))
    extra_cols[3].metric("Tolerance (sent)", config.get("tolerance_sentences", "—"))

    st.divider()

    # ── High-level counts ─────────────────────────────────────────────────────
    st.markdown("### 📦 Dataset Overview")

    agg_cols = st.columns(4)
    agg_cols[0].metric("Total Documents", aggregate.get("total_documents", "—"))
    agg_cols[1].metric("Successful", aggregate.get("successful", "—"))
    agg_cols[2].metric("Failed", aggregate.get("failed", 0))
    agg_cols[3].metric(
        "Municipalities", len(aggregate.get("municipalities", {}))
    )

    overall = aggregate.get("overall", {})
    ov_cols = st.columns(4)
    ov_cols[0].metric("Agenda Items GT", overall.get("total_agenda_items_gt", "—"))
    ov_cols[1].metric("Agenda Items Predicted", overall.get("total_agenda_items_predicted", "—"))
    ov_cols[2].metric("Subjects GT", overall.get("total_subjects_gt", "—"))
    ov_cols[3].metric("Subjects Predicted", overall.get("total_subjects_predicted", "—"))

    extra2_cols = st.columns(2)
    extra2_cols[0].metric(
        "Avg Subjects/Item (Predicted)",
        _fmt(overall.get("avg_subjects_per_item_predicted", 0), 2),
    )
    extra2_cols[1].metric(
        "Avg Subjects/Item (GT)",
        _fmt(overall.get("avg_subjects_per_item_gt", 0), 2),
    )

    st.divider()

    # ── Metric tables ─────────────────────────────────────────────────────────
    def _metric_table(title: str, section_key: str):
        st.markdown(f"### 📐 {title} Metrics")
        m = aggregate.get(section_key, {})
        if not m:
            st.info(f"No {title.lower()} metrics available.")
            return

        metric_rows = [
            ("Boundary Precision", m.get("boundary_precision")),
            ("Boundary Recall", m.get("boundary_recall")),
            ("Boundary F1", m.get("boundary_f1")),
            ("Boundary Similarity", m.get("boundary_similarity")),
            ("BED Precision", m.get("bed_precision")),
            ("BED Recall", m.get("bed_recall")),
            ("BED F-measure", m.get("bed_fmeasure")),
            ("Segeval Pk", m.get("segeval_pk")),
            ("Segeval WindowDiff", m.get("segeval_windowdiff")),
        ]
        # subject-only extras
        if section_key == "subjects":
            metric_rows += [
                ("Theme Accuracy", m.get("theme_accuracy")),
                ("Topic Precision", m.get("topic_precision")),
                ("Topic Recall", m.get("topic_recall")),
                ("Topic F1", m.get("topic_f1")),
            ]

        count_rows = [
            ("True Positives", m.get("true_positives")),
            ("False Positives", m.get("false_positives")),
            ("False Negatives", m.get("false_negatives")),
            ("Num Predicted", m.get("num_predicted")),
            ("Num Ground Truth", m.get("num_ground_truth")),
        ]

        left, right = st.columns(2)

        with left:
            st.markdown("**Boundary / BED / Segeval**")
            df = pd.DataFrame(
                [(name, _fmt(val)) for name, val in metric_rows if val is not None],
                columns=["Metric", "Value"],
            )
            st.dataframe(df, use_container_width=True, hide_index=True)

        with right:
            st.markdown("**Counts**")
            df2 = pd.DataFrame(
                [(name, val) for name, val in count_rows if val is not None],
                columns=["Metric", "Value"],
            )
            st.dataframe(df2, use_container_width=True, hide_index=True)

            # Radar chart of key metrics
            radar_metrics = {
                "Boundary P": m.get("boundary_precision", 0),
                "Boundary R": m.get("boundary_recall", 0),
                "Boundary F1": m.get("boundary_f1", 0),
                "BED F-measure": m.get("bed_fmeasure", 0),
                "BS": m.get("boundary_similarity", 0),
            }
            labels = list(radar_metrics.keys())
            values = list(radar_metrics.values())
            values.append(values[0])  # close the polygon

            fig = go.Figure(
                data=[
                    go.Scatterpolar(
                        r=values,
                        theta=labels + [labels[0]],
                        fill="toself",
                        line_color="#4C72B0",
                        fillcolor="rgba(76,114,176,0.25)",
                        name=title,
                    )
                ],
                layout=go.Layout(
                    polar=dict(radialaxis=dict(visible=True, range=[0, 1])),
                    showlegend=False,
                    margin=dict(l=30, r=30, t=40, b=20),
                    height=280,
                ),
            )
            st.plotly_chart(fig, use_container_width=True)

    _metric_table("Agenda Items", "agenda_items")
    st.divider()
    _metric_table("Subjects", "subjects")
    st.divider()

    # ── Municipality breakdown ────────────────────────────────────────────────
    st.markdown("### 🏛️ Municipality Breakdown")
    municipalities = aggregate.get("municipalities", {})
    if municipalities:
        muni_rows = []
        for muni, info in municipalities.items():
            muni_rows.append(
                {
                    "Municipality": muni,
                    "Documents": info.get("documents", 0),
                    "Agenda Items GT": info.get("agenda_items_gt", 0),
                    "Subjects GT": info.get("subjects_gt", 0),
                }
            )
        mdf = pd.DataFrame(muni_rows)
        st.dataframe(mdf, use_container_width=True, hide_index=True)

        fig = px.bar(
            mdf,
            x="Municipality",
            y=["Agenda Items GT", "Subjects GT"],
            barmode="group",
            title="Ground-Truth Counts by Municipality",
            color_discrete_sequence=["#4C72B0", "#DD8452"],
        )
        fig.update_layout(height=350, margin=dict(t=40, b=20))
        st.plotly_chart(fig, use_container_width=True)

    # ── Per-document metric distribution ─────────────────────────────────────
    if documents:
        st.divider()
        st.markdown("### 📄 Per-Document Metric Distribution")

        doc_rows = []
        for doc_id, doc_data in documents.items():
            ev = doc_data.get("evaluation", {})
            ai = ev.get("agenda_items", {})
            subj = ev.get("subjects", {})
            doc_rows.append(
                {
                    "Document": doc_id,
                    "Municipality": doc_data.get("municipality", ""),
                    "AI Boundary F1": ai.get("boundary_f1"),
                    "AI BED F-measure": ai.get("bed_fmeasure"),
                    "Subj Boundary F1": subj.get("boundary_f1"),
                    "Subj BED F-measure": subj.get("bed_fmeasure"),
                    "Subj Boundary Sim": subj.get("boundary_similarity"),
                    "AI #Predicted": ai.get("num_predicted"),
                    "AI #GT": ai.get("num_ground_truth"),
                    "Subj #Predicted": subj.get("num_predicted"),
                    "Subj #GT": subj.get("num_ground_truth"),
                }
            )

        doc_df = pd.DataFrame(doc_rows)
        st.dataframe(doc_df, use_container_width=True, hide_index=True)

        # Distribution plots
        metric_col = st.selectbox(
            "Plot distribution for metric",
            ["AI Boundary F1", "AI BED F-measure", "Subj Boundary F1", "Subj BED F-measure"],
            key="stats_dist_metric",
        )
        fig2 = px.histogram(
            doc_df,
            x=metric_col,
            color="Municipality",
            nbins=15,
            title=f"Distribution of {metric_col}",
            opacity=0.8,
        )
        fig2.update_layout(height=350, margin=dict(t=40, b=20))
        st.plotly_chart(fig2, use_container_width=True)

        # Scatter: AI vs Subject performance
        fig3 = px.scatter(
            doc_df,
            x="AI Boundary F1",
            y="Subj Boundary F1",
            color="Municipality",
            hover_data=["Document"],
            title="Agenda Items vs Subject Boundary F1 per Document",
            size_max=12,
        )
        fig3.update_layout(height=400, margin=dict(t=40, b=20))
        st.plotly_chart(fig3, use_container_width=True)