src/app.py

"""
LLM Subject Extraction — Streamlit Demo

Analyses and manually evaluates LLM-based subject extraction results stored
in results/llm_evaluation/<run_id>/.

Pages
-----
  Overview        — high-level info & quick stats
  📊 Statistics   — aggregate metrics & per-document distributions
  🔍 Document     — aligned GT vs prediction view for one document
  ⚔️ Comparison   — cross-run metric comparison
"""

import streamlit as st
import sys
from pathlib import Path

# ── Path setup ───────────────────────────────────────────────────────────────
_DEMO_DIR = Path(__file__).parent
sys.path.insert(0, str(_DEMO_DIR))

from components.data_loader import LLMEvalDataLoader
from components.statistics import render_statistics
from components.visualizer import render_document_view
from components.comparison import render_comparison

# ── Results root — adjust if running from a different working directory ───────
_REPO_ROOT = _DEMO_DIR.parent.parent
_RESULTS_ROOT = _REPO_ROOT / "results" / "llm_evaluation"

# ── Page config ───────────────────────────────────────────────────────────────
st.set_page_config(
    page_title="LLM Subject Extraction · Evaluation Demo",
    page_icon="🧠",
    layout="wide",
    initial_sidebar_state="expanded",
)

# ── Custom CSS ────────────────────────────────────────────────────────────────
st.markdown(
    """
<style>
/* ---------- global ---------- */
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&display=swap');

html, body, [class*="css"] {
    font-family: 'Inter', sans-serif;
}

/* ---------- header ---------- */
.main-header {
    background: linear-gradient(135deg, #1a1a2e 0%, #16213e 50%, #0f3460 100%);
    color: white;
    padding: 2rem 2.5rem;
    border-radius: 16px;
    margin-bottom: 1.5rem;
    box-shadow: 0 8px 32px rgba(0,0,0,0.18);
}
.main-header h1 {
    margin: 0 0 0.4rem 0;
    font-size: 2rem;
    font-weight: 700;
    letter-spacing: -0.5px;
}
.main-header p {
    margin: 0;
    font-size: 0.95rem;
    opacity: 0.78;
}

/* ---------- section header ---------- */
.section-header {
    font-size: 1.4rem;
    font-weight: 600;
    color: #0f3460;
    border-left: 5px solid #e94560;
    padding-left: 0.8rem;
    margin: 1.5rem 0 1rem 0;
}

/* ---------- sidebar ---------- */
[data-testid="stSidebar"] {
    background: linear-gradient(180deg, #1a1a2e 0%, #16213e 100%);
}
[data-testid="stSidebar"] * {
    color: #eee !important;
}
[data-testid="stSidebar"] .stSelectbox label,
[data-testid="stSidebar"] .stMultiSelect label,
[data-testid="stSidebar"] .stRadio label {
    color: #ccc !important;
    font-size: 0.88rem;
}
[data-testid="stSidebar"] hr {
    border-color: #334 !important;
}

/* ---------- metric cards ---------- */
[data-testid="metric-container"] {
    background: #f8faff;
    border-radius: 10px;
    padding: 0.6rem 0.9rem;
    border: 1px solid #e0e8f5;
    box-shadow: 0 2px 8px rgba(15,52,96,0.06);
}

/* ---------- overview cards ---------- */
.overview-card {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    color: white;
    padding: 1.3rem 1.5rem;
    border-radius: 14px;
    text-align: center;
    box-shadow: 0 4px 16px rgba(102,126,234,0.3);
    transition: transform 0.2s;
}
.overview-card:hover { transform: translateY(-3px); }
.overview-card h3 { margin: 0 0 0.4rem 0; font-size: 1.1rem; font-weight: 600; }
.overview-card p { margin: 0; font-size: 0.85rem; opacity: 0.85; }

/* ---------- run badge ---------- */
.run-badge {
    display:inline-block;
    background:#0f3460;
    color:white;
    padding:3px 10px;
    border-radius:8px;
    font-size:0.78rem;
    margin:2px;
}
</style>
""",
    unsafe_allow_html=True,
)


# ── Session state init ────────────────────────────────────────────────────────
def _init_state() -> None:
    if "loader" not in st.session_state:
        st.session_state.loader = LLMEvalDataLoader(str(_RESULTS_ROOT))
    if "selected_run" not in st.session_state:
        st.session_state.selected_run = None
    if "compare_runs" not in st.session_state:
        st.session_state.compare_runs = []


# ── Sidebar ───────────────────────────────────────────────────────────────────
def _render_sidebar(loader: LLMEvalDataLoader) -> str:
    with st.sidebar:
        st.markdown("## 🧠 LLM Extraction Demo")
        st.divider()

        # Navigation
        page = st.radio(
            "Navigate",
            ["🏠 Overview", "📊 Statistics", "🔍 Document Explorer", "⚔️ Run Comparison"],
            key="nav_page",
        )
        st.divider()

        # Run selection
        st.markdown("### 📂 Evaluation Run")
        runs = loader.get_available_runs()

        if not runs:
            st.error("No evaluation runs found.\nCheck that results/llm_evaluation/ exists.")
            return page

        run_labels = [
            f"{r['model_name']} — {r['timestamp'][:10]}" for r in runs
        ]
        run_ids = [r["run_id"] for r in runs]

        selected_idx = st.selectbox(
            "Primary Run",
            range(len(run_labels)),
            format_func=lambda i: run_labels[i],
            key="run_selector",
        )
        st.session_state.selected_run = run_ids[selected_idx]

        # Quick info from selected run
        cfg_run = runs[selected_idx]
        st.caption(f"Backend: **{cfg_run['backend']}**")
        st.caption(f"Model: **{cfg_run['model_name']}**")
        st.caption(f"Split: **{cfg_run['split']}**")
        st.caption(f"Timestamp: {cfg_run['timestamp'][:19]}")

        st.divider()

        # Document selection (for Document Explorer page)
        if page == "🔍 Document Explorer":
            st.markdown("### 📄 Document")
            run_data = loader.load_run(st.session_state.selected_run)
            if run_data:
                doc_ids = sorted(run_data.get("documents", {}).keys())
                municipalities = sorted(set(
                    loader.parse_municipality(d) for d in doc_ids
                ))
                muni_filter = st.multiselect(
                    "Filter by municipality",
                    municipalities,
                    default=[],
                    key="muni_filter",
                )
                if muni_filter:
                    doc_ids = [
                        d for d in doc_ids
                        if loader.parse_municipality(d) in muni_filter
                    ]

                selected_doc = st.selectbox(
                    "Document",
                    doc_ids,
                    key="doc_selector",
                )
                st.session_state.selected_doc = selected_doc

        # Comparison run selection
        if page == "⚔️ Run Comparison":
            st.markdown("### ⚔️ Compare Runs")
            compare_idxs = st.multiselect(
                "Select runs to compare",
                range(len(run_labels)),
                default=list(range(min(3, len(run_labels)))),
                format_func=lambda i: run_labels[i],
                key="compare_selector",
            )
            st.session_state.compare_runs = [run_ids[i] for i in compare_idxs]

    return page


# ── Pages ─────────────────────────────────────────────────────────────────────

def _page_overview(loader: LLMEvalDataLoader) -> None:
    st.markdown(
        """
<div class="main-header">
    <h1>🧠 LLM Subject Extraction · Evaluation Demo</h1>
    <p>Analyse and manually evaluate LLM-based agenda-item & subject segmentation results.</p>
</div>
""",
        unsafe_allow_html=True,
    )

    run_data = loader.load_run(st.session_state.selected_run)
    if not run_data:
        st.error("Could not load the selected run.")
        return

    agg = run_data.get("aggregate", {})
    cfg = run_data.get("config", {})
    pcfg = cfg.get("pipeline_config", {})

    # Cards row
    card_cols = st.columns(4)
    cards = [
        ("📄 Documents", str(agg.get("total_documents", "—")), "Processed in this run"),
        ("✅ Successful", str(agg.get("successful", "—")), "Without errors"),
        ("❌ Failed", str(agg.get("failed", 0)), "Extraction errors"),
        ("🏛️ Municipalities", str(len(agg.get("municipalities", {}))), "In the test set"),
    ]
    for col, (emoji_title, value, desc) in zip(card_cols, cards):
        with col:
            st.markdown(
                f'<div class="overview-card"><h3>{emoji_title}</h3>'
                f'<p style="font-size:1.8rem;font-weight:700;margin:0.3rem 0">{value}</p>'
                f'<p>{desc}</p></div>',
                unsafe_allow_html=True,
            )

    st.divider()

    # Key metrics snapshot
    st.markdown('<p class="section-header">📐 Key Metrics Snapshot</p>', unsafe_allow_html=True)

    ai = agg.get("agenda_items", {})
    subj = agg.get("subjects", {})

    m1, m2, m3, m4, m5, m6 = st.columns(6)
    m1.metric("AI Boundary F1", f"{ai.get('boundary_f1', 0):.3f}")
    m2.metric("AI BED F-measure", f"{ai.get('bed_fmeasure', 0):.3f}")
    m3.metric("AI Segeval Pk", f"{ai.get('segeval_pk', 0):.3f}")
    m4.metric("Subj Boundary F1", f"{subj.get('boundary_f1', 0):.3f}")
    m5.metric("Subj BED F-measure", f"{subj.get('bed_fmeasure', 0):.3f}")
    m6.metric("Subj Theme Acc", f"{subj.get('theme_accuracy', 0):.3f}")

    st.divider()

    # Getting started
    st.markdown('<p class="section-header">🚀 Getting Started</p>', unsafe_allow_html=True)
    st.markdown(
        """
1. **Select an Evaluation Run** in the sidebar (primary run for exploration).
2. **📊 Statistics** — Full metric breakdown, municipality analysis, per-document distributions.
3. **🔍 Document Explorer** — Step through individual documents, view aligned predictions vs ground truth.
4. **⚔️ Run Comparison** — Compare aggregate metrics across multiple runs with radar & bar charts.
"""
    )


def _page_statistics(loader: LLMEvalDataLoader) -> None:
    st.markdown('<p class="section-header">📊 Statistics</p>', unsafe_allow_html=True)
    run_data = loader.load_run(st.session_state.selected_run)
    if not run_data:
        st.error("Could not load run data.")
        return
    render_statistics(run_data)


def _page_document_explorer(loader: LLMEvalDataLoader) -> None:
    st.markdown('<p class="section-header">🔍 Document Explorer</p>', unsafe_allow_html=True)

    run_id = st.session_state.selected_run
    doc_id = st.session_state.get("selected_doc")

    if not doc_id:
        st.info("Select a document from the sidebar.")
        return

    doc_data = loader.get_document(run_id, doc_id)
    if not doc_data:
        st.error(f"Could not load document: {doc_id}")
        return

    render_document_view(doc_data, loader.get_topic_color)


def _page_comparison(loader: LLMEvalDataLoader) -> None:
    st.markdown('<p class="section-header">⚔️ Run Comparison</p>', unsafe_allow_html=True)

    compare_run_ids = st.session_state.get("compare_runs", [])
    if not compare_run_ids:
        st.info("Select runs to compare from the sidebar.")
        return

    runs_data = []
    run_labels = []
    for rid in compare_run_ids:
        rd = loader.load_run(rid)
        if rd:
            runs_data.append(rd)
            cfg = rd.get("config", {}).get("pipeline_config", {})
            run_labels.append(cfg.get("model_name", rid))

    render_comparison(runs_data, run_labels)


# ── Main ──────────────────────────────────────────────────────────────────────

def main() -> None:
    _init_state()
    loader: LLMEvalDataLoader = st.session_state.loader

    page = _render_sidebar(loader)

    if page == "🏠 Overview":
        _page_overview(loader)
    elif page == "📊 Statistics":
        _page_statistics(loader)
    elif page == "🔍 Document Explorer":
        _page_document_explorer(loader)
    elif page == "⚔️ Run Comparison":
        _page_comparison(loader)


if __name__ == "__main__":
    main()