#!/usr/bin/env python3
"""
Translation Quality Dashboard.

Reads evaluation results produced by scripts/eval/llm_judge.py:
  eval_dir/{model}/*.holistic.jsonl       — HolisticTranslationJudge results
  eval_dir/{model}/*.intent_entity.jsonl  — IntentEntityJudge results

Supports passing a top-level eval folder (e.g. eval_results/apr-21-eval-results/)
which contains per-model subdirectories, or a single model directory directly.

Pages:
  Overview          — Cross-model holistic + intent/entity metrics side by side
  Holistic Deep Dive — Per-model dimension scores, issues, segment inspector
  Intent & Entity   — Intent accuracy, entity scores, per-segment breakdown
  Issues Analysis   — Filterable issues table across models

Usage:
    streamlit run scripts/eval/dashboard.py -- --eval-dir <eval-dir>
    EVAL_DIR=eval_results/apr-21-eval-results streamlit run scripts/eval/dashboard.py
    streamlit run scripts/eval/dashboard.py -- --eval-dir <eval-dir> --exclude-satsangs <id>,<id>
"""

import argparse
import json
import os
from pathlib import Path
from typing import Optional

import numpy as np
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
import streamlit as st
from huggingface_hub import HfApi, snapshot_download

# ── Config ────────────────────────────────────────────────────────────────────

DIMENSIONS = [
    "semantic_accuracy",
    "spiritual_fidelity",
    "expression_quality",
    "asr_robustness",
    "contextual_coherence",
]

DIMENSION_LABELS = {
    "semantic_accuracy": "Semantic Accuracy",
    "spiritual_fidelity": "Spiritual Fidelity",
    "expression_quality": "Expression Quality",
    "asr_robustness": "ASR Robustness",
    "contextual_coherence": "Contextual Coherence",
}

WEIGHTS = {
    "semantic_accuracy": 0.35,
    "spiritual_fidelity": 0.30,
    "expression_quality": 0.10,
    "asr_robustness": 0.15,
    "contextual_coherence": 0.10,
}

ISSUE_TYPES = [
    "word_confusion",
    "terminology_error",
    "negation_flip",
    "omission",
    "insertion",
    "hallucination",
    "content_safety",
]

SEVERITY_COLORS = {"critical": "#e74c3c", "major": "#f39c12", "minor": "#f1c40f"}
SEVERITY_ORDER = {"critical": 0, "major": 1, "minor": 2}
SEVERITY_WEIGHT = {"critical": 1.0, "major": 0.5, "minor": 0.1}

# ── Definitions ───────────────────────────────────────────────────────────────

DIMENSION_DEFINITIONS = {
    "semantic_accuracy": (
        "**Semantic Accuracy (weight 35%)** — Does the hypothesis preserve the speaker's core meaning "
        "compared to the human reference? Omissions of filler are acceptable; adding incorrect content "
        "is penalised heavily. "
        "**5** = meaning fully preserved. **3** = one key clause missing or slightly wrong. "
        "**1** = meaning completely wrong or inverted."
    ),
    "spiritual_fidelity": (
        "**Spiritual Fidelity (weight 30%)** — Are Jain, Hindu/Vedantic, and Shrimad Rajchandra terms "
        "translated correctly within their tradition? Concepts like karma, moksha, and atma carry "
        "tradition-specific meanings — using the wrong tradition's interpretation is a major error. "
        "**5** = all terms accurate. **3** = generic translation that loses tradition-specific nuance. "
        "**1** = core philosophical concepts completely wrong."
    ),
    "expression_quality": (
        "**Expression Quality (weight 10%)** — Is the English natural, grammatically correct, and "
        "accessible to a general Satsang audience? The reference captions use simple English — "
        "overly academic or complex phrasing is penalised. "
        "**5** = reads naturally, matches reference register. **3** = noticeably more complex than reference. "
        "**1** = incomprehensible or wrong genre."
    ),
    "asr_robustness": (
        "**ASR Robustness (weight 15%)** — Did the model correctly recover from noisy or phonetically "
        "garbled Gujarati ASR input? "
        "**5** = seamlessly corrected ASR errors. **3** = translated ASR literally causing mild confusion. "
        "**1** = failed to parse ASR or hallucinated from garbled input."
    ),
    "contextual_coherence": (
        "**Contextual Coherence (weight 10%)** — Is the translation consistent with the preceding "
        "discourse turns? Pronouns resolved, speaker's train of thought maintained. "
        "**5** = perfect continuity. **3** = some inconsistency but understandable in isolation. "
        "**1** = completely disconnected from discourse flow."
    ),
}

WEIGHTED_SCORE_DEFINITION = (
    "**Weighted Score** = (Semantic Accuracy × 0.35) + (Spiritual Fidelity × 0.30) + "
    "(ASR Robustness × 0.15) + (Expression Quality × 0.10) + (Contextual Coherence × 0.10), "
    "then a severity penalty is subtracted per issue (critical −0.5, major −0.25, minor −0.1). "
    "Range: 1.0–5.0. Higher is better."
)

ISSUE_DEFINITIONS = {
    "word_confusion": "A similar-sounding word was swapped (e.g. 'skill' → 'kill').",
    "terminology_error": "A Jain/spiritual term was given the wrong English equivalent.",
    "negation_flip": "A negation was added or dropped, reversing the meaning.",
    "omission": "A key qualifier or clause was dropped such that the core meaning changes.",
    "insertion": "A concept was added that the speaker did not express (worse than omission).",
    "hallucination": "Content was invented with no basis in the ASR or reference.",
    "content_safety": "The translator introduced divisive or offensive language not in the source.",
}

INTENT_ENTITY_DEFINITION = (
    "Intent & Entity evaluation is based on the open-source framework by Sarvam AI: "
    "[github.com/sarvamai/llm_intent_entity](https://github.com/sarvamai/llm_intent_entity/tree/main). "
    "An LLM judge compares the model's hypothesis against the human ground truth on two dimensions:"
)

INTENT_DEFINITION = (
    "**Intent Score (0 or 1)** — Binary. Did the hypothesis preserve the core meaning/action of the sentence? "
    "Score **1** if the subject, action, and intent match (e.g. same actor, same direction, same speech act). "
    "Score **0** if pronouns differ, a statement becomes a question, or the action is reversed. "
    "Equivalent fillers (yes/yeah/hmm) all score 1."
)

ENTITY_DEFINITION = (
    "**Entity Score (0.0–1.0)** — Fraction of key entities from the ground truth correctly preserved. "
    "Entities include names, places, dates, numbers, and domain-specific references. "
    "**1.0** = all entities present and correct. **0.0** = no entities preserved. "
    "Partial credit is given proportionally; substituted entities (John → Tom) are penalised. "
    "Sentences with no entities (greetings, fillers) automatically score 1.0."
)

HALLUCINATION_IMPACT_DEFINITION = (
    "**Hallucination/Safety Impact** = weighted average of severity scores per segment "
    "(critical=1.0, major=0.5, minor=0.1). "
    "A model with only minor hallucinations in 50% of segments scores lower than one with a single critical hallucination. "
    "Range: 0.0 (none) → higher is worse."
)


HF_REPO_ID = "srmdtranslations/model_gu_en_evaluations"


@st.cache_resource
def _list_hf_branches(token: Optional[str]) -> list[str]:
    """List all branches in the HF eval repo."""
    refs = HfApi().list_repo_refs(HF_REPO_ID, repo_type="dataset", token=token)
    return [b.name for b in refs.branches]


@st.cache_resource
def _hf_snapshot(branch: str, token: Optional[str]) -> str:
    """Download the entire branch from HF to local cache; returns local root path."""
    return snapshot_download(
        repo_id=HF_REPO_ID,
        repo_type="dataset",
        revision=branch,
        token=token,
    )


def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(add_help=False)
    parser.add_argument(
        "--eval-dir",
        default=os.environ.get("EVAL_DIR", "eval_results"),
        help="Path to eval results directory",
    )
    parser.add_argument(
        "--exclude-satsangs",
        default=None,
        help="Comma-separated satsang_ids to exclude from all pages (e.g. badly-aligned files)",
    )
    args, _ = parser.parse_known_args()
    return args


_ARGS = _parse_args()
_DEFAULT_EVAL_DIR = _ARGS.eval_dir
_EXCLUDE_SATSANGS = frozenset(
    s.strip() for s in (_ARGS.exclude_satsangs or "").split(",") if s.strip()
)


# ── Data loading ──────────────────────────────────────────────────────────────


def _collect_jsonl_files(eval_dir: str, suffix: str) -> list[tuple[str, Path]]:
    """
    Return (model_name, file_path) pairs for all files ending in `suffix`.
    Handles both flat (eval_dir/model/*.jsonl) and two-level structures.
    """
    root = Path(eval_dir)
    pairs = []
    if not root.exists():
        return pairs
    for subdir in sorted(root.iterdir()):
        if not subdir.is_dir():
            continue
        files = sorted(subdir.glob(f"*{suffix}"))
        if files:
            # model-level dir
            for f in files:
                pairs.append((subdir.name, f))
        else:
            # eval-level dir — go one level deeper
            for model_dir in sorted(subdir.iterdir()):
                if model_dir.is_dir():
                    for f in sorted(model_dir.glob(f"*{suffix}")):
                        pairs.append((model_dir.name, f))
    return pairs


def _extract_score(row: dict, dim: str) -> Optional[float]:
    val = row.get(dim)
    if val is None:
        return None
    if isinstance(val, dict):
        return val.get("score")
    if isinstance(val, (int, float)):
        return float(val)
    return None


def _extract_reasoning(row: dict, dim: str) -> str:
    val = row.get(dim)
    if isinstance(val, dict):
        return val.get("reasoning", "")
    return ""


def _extract_issues(row: dict) -> list[dict]:
    issues = row.get("issues")
    if isinstance(issues, list):
        return issues
    return []


@st.cache_data
def load_holistic(eval_dir: str) -> pd.DataFrame:
    """Load all *.holistic.jsonl files into a DataFrame."""
    rows = []
    for model_name, f in _collect_jsonl_files(eval_dir, ".holistic.jsonl"):
        with open(f, encoding="utf-8") as fh:
            for line in fh:
                raw = json.loads(line)
                record = {
                    "model": model_name,
                    "file": f.stem.replace(".holistic", ""),
                    "satsang_id": raw.get("content_id", f.stem.replace(".holistic", "")),
                    "srt_index": raw.get("srt_index"),
                    "seg_index": raw.get("index"),
                    "input": raw.get("input", ""),
                    "reference": raw.get("output", ""),
                    "hypothesis": raw.get("hypothesis", ""),
                    "overall": raw.get("overall"),
                    "confidence": raw.get("confidence", ""),
                    "alignment_confidence": raw.get("alignment_confidence"),
                    "overall_reasoning": raw.get("overall_reasoning", ""),
                    "weighted_score": raw.get("weighted_score"),
                }
                for dim in DIMENSIONS:
                    record[dim] = _extract_score(raw, dim)
                    record[f"{dim}_reasoning"] = _extract_reasoning(raw, dim)
                if record["weighted_score"] is None:
                    scores = [record[d] for d in DIMENSIONS if record[d] is not None]
                    record["weighted_score"] = (
                        sum((record[d] or 0) * WEIGHTS[d] for d in DIMENSIONS)
                        if scores else None
                    )
                issues = _extract_issues(raw)
                record["issues_detail"] = issues
                record["issue_count"] = len(issues)
                record["has_issues"] = len(issues) > 0
                record["critical_count"] = sum(1 for i in issues if i.get("severity") == "critical")
                record["major_count"] = sum(1 for i in issues if i.get("severity") == "major")
                record["minor_count"] = sum(1 for i in issues if i.get("severity") == "minor")
                record["issue_types"] = ", ".join(sorted({i.get("type", "") for i in issues})) if issues else ""
                record["worst_severity"] = (
                    min((i.get("severity", "minor") for i in issues), key=lambda s: SEVERITY_ORDER.get(s, 99))
                    if issues else "none"
                )
                record["hallucination_impact"] = sum(
                    SEVERITY_WEIGHT.get(i["severity"], 0) for i in issues if i.get("type") == "hallucination"
                )
                record["safety_impact"] = sum(
                    SEVERITY_WEIGHT.get(i["severity"], 0) for i in issues if i.get("type") == "content_safety"
                )
                record["misaligned"] = raw.get("alignment_confidence") == 0
                rows.append(record)
    df = pd.DataFrame(rows)
    if _EXCLUDE_SATSANGS and not df.empty:
        df = df[~df["satsang_id"].isin(_EXCLUDE_SATSANGS)].reset_index(drop=True)
    return df


@st.cache_data
def load_intent_entity(eval_dir: str) -> pd.DataFrame:
    """Load all *.intent_entity.jsonl files into a DataFrame."""
    rows = []
    for model_name, f in _collect_jsonl_files(eval_dir, ".intent_entity.jsonl"):
        with open(f, encoding="utf-8") as fh:
            for line in fh:
                raw = json.loads(line)
                record = {
                    "model": model_name,
                    "file": f.stem.replace(".intent_entity", ""),
                    "satsang_id": raw.get("content_id", f.stem.replace(".intent_entity", "")),
                    "srt_index": raw.get("srt_index"),
                    "seg_index": raw.get("index"),
                    "input": raw.get("input", ""),
                    "reference": raw.get("output", ""),
                    "hypothesis": raw.get("hypothesis", ""),
                    "intent_score": raw.get("intent_score"),
                    "intent_explanation": raw.get("intent_explanation", ""),
                    "entity_score": raw.get("entity_score"),
                    "ground_truth_entities": raw.get("ground_truth_entities", ""),
                    "preserved_entities": raw.get("preserved_entities", ""),
                    "missing_entities": raw.get("missing_entities", ""),
                    "entity_explanation": raw.get("entity_explanation", ""),
                }
                rows.append(record)
    df = pd.DataFrame(rows)
    if _EXCLUDE_SATSANGS and not df.empty:
        df = df[~df["satsang_id"].isin(_EXCLUDE_SATSANGS)].reset_index(drop=True)
    return df


# ── Chart helpers ─────────────────────────────────────────────────────────────


def radar_chart(means: dict[str, float], title: str = "") -> go.Figure:
    dims = list(means.keys())
    vals = list(means.values())
    dims.append(dims[0])
    vals.append(vals[0])
    fig = go.Figure()
    fig.add_trace(go.Scatterpolar(r=vals, theta=dims, fill="toself", name=title or "Score"))
    fig.update_layout(
        polar=dict(radialaxis=dict(visible=True, range=[0, 5])),
        showlegend=False,
        title=title,
        height=380,
    )
    return fig


def model_comparison_radar(df: pd.DataFrame) -> go.Figure:
    fig = go.Figure()
    for model in sorted(df["model"].unique()):
        mdf = df[df["model"] == model]
        means = {DIMENSION_LABELS[d]: mdf[d].mean() for d in DIMENSIONS if mdf[d].notna().any()}
        if not means:
            continue
        dims = list(means.keys()) + [list(means.keys())[0]]
        vals = list(means.values()) + [list(means.values())[0]]
        fig.add_trace(go.Scatterpolar(r=vals, theta=dims, fill="toself", name=model))
    fig.update_layout(
        polar=dict(radialaxis=dict(visible=True, range=[0, 5])),
        title="Dimension Scores by Model",
        height=460,
    )
    return fig


def weighted_score_boxplot(df: pd.DataFrame) -> go.Figure:
    fig = px.box(
        df, x="model", y="weighted_score", color="model",
        title="Weighted Score Distribution",
        labels={"weighted_score": "Weighted Score", "model": "Model"},
        points="outliers",
    )
    fig.update_layout(showlegend=False, height=380)
    return fig


def score_distribution_chart(df: pd.DataFrame, dimension: str) -> go.Figure:
    counts = df.groupby(["model", dimension]).size().reset_index(name="count")
    counts[dimension] = counts[dimension].astype(int)
    fig = px.bar(
        counts, x=dimension, y="count", color="model", barmode="group",
        title=f"{DIMENSION_LABELS.get(dimension, dimension)} Score Distribution",
        labels={dimension: "Score (1–5)", "count": "Segments"},
    )
    fig.update_layout(xaxis=dict(dtick=1, range=[0.5, 5.5]), height=350)
    return fig


def satsang_heatmap(df: pd.DataFrame, model: str) -> go.Figure:
    mdf = df[df["model"] == model]
    pivot = mdf.groupby("satsang_id")[DIMENSIONS].mean()
    pivot.columns = [DIMENSION_LABELS[d] for d in DIMENSIONS]
    if len(pivot) > 30:
        pivot = pivot.head(30)
    fig = px.imshow(
        pivot.values,
        x=pivot.columns.tolist(),
        y=[s[:14] + "…" if len(s) > 14 else s for s in pivot.index.tolist()],
        color_continuous_scale="RdYlGn",
        zmin=1, zmax=5,
        title=f"Per-Satsang Scores — {model}",
        aspect="auto",
    )
    fig.update_layout(height=max(380, len(pivot) * 22))
    return fig


# ── Definition helpers ────────────────────────────────────────────────────────


def _holistic_legend():
    with st.expander("📖 How to read holistic scores", expanded=False):
        st.markdown(WEIGHTED_SCORE_DEFINITION)
        st.markdown("---")
        for defn in DIMENSION_DEFINITIONS.values():
            st.markdown(defn)
        st.markdown("---")
        st.markdown("**Issue severity** — critical: completely inverts meaning or safety risk. "
                    "major: misleads listener. minor: noticeable but recoverable.")
        cols = st.columns(len(ISSUE_DEFINITIONS))
        for col, (itype, desc) in zip(cols, ISSUE_DEFINITIONS.items()):
            col.markdown(f"**`{itype}`**  \n{desc}")


def _intent_entity_legend():
    with st.expander("📖 How to read intent & entity scores", expanded=False):
        st.markdown(
            INTENT_ENTITY_DEFINITION + "  \n  \n"
            + INTENT_DEFINITION + "  \n  \n"
            + ENTITY_DEFINITION
        )
        st.markdown(
            "**Reading the charts:**  \n"
            "- **Intent Accuracy bar chart** — proportion of segments where intent was preserved. "
            "Closer to 1.0 is better.  \n"
            "- **Entity Score histogram** — distribution of per-segment entity preservation. "
            "A left-skewed distribution (peak near 1.0) indicates strong entity retention.  \n"
            "- **Intent vs. Entity scatter** — each dot is one segment. "
            "Top-right (intent=1, entity≈1) is ideal. "
            "Dots in the bottom-left indicate both intent and entity failures — the most critical segments to inspect."
        )


def _issues_legend():
    with st.expander("📖 Issue types & severity explained", expanded=False):
        st.markdown("**Issue types:**")
        for itype, desc in ISSUE_DEFINITIONS.items():
            st.markdown(f"- **`{itype}`** — {desc}")
        st.markdown("---")
        st.markdown(HALLUCINATION_IMPACT_DEFINITION)
        st.markdown(
            "**Reading the charts:**  \n"
            "- **Issues by Severity per Model** — compare total critical/major/minor counts across models. "
            "A model with fewer critical issues is safer even if its weighted score is similar.  \n"
            "- **Issues by Type per Model** — reveals *what kind* of errors each model makes. "
            "High `hallucination` counts need closer human review than high `omission` counts.  \n"
            "- **Hallucination & Safety Impact** — unlike raw counts, this weights severity so one critical "
            "hallucination outweighs many minor ones. Lower is better."
        )


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


def render_overview(hdf: pd.DataFrame, iedf: pd.DataFrame):
    st.header("Overview — All Models")
    st.caption(
        "Side-by-side comparison of all models. "
        "**Weighted Score** combines 5 holistic dimensions (higher = better, max 5.0). "
        "**Intent Accuracy** and **Entity Score** are from the Sarvam AI intent/entity framework (higher = better, max 1.0)."
    )
    _holistic_legend()
    _intent_entity_legend()

    # ── Summary table ─────────────────────────────────────────────────────────
    models = sorted(set(
        list(hdf["model"].unique() if not hdf.empty else [])
        + list(iedf["model"].unique() if not iedf.empty else [])
    ))
    rows = []
    for model in models:
        row = {"Model": model}
        if not hdf.empty and model in hdf["model"].values:
            mh_all = hdf[hdf["model"] == model]
            mh = mh_all[mh_all["weighted_score"].notna()]
            row["Segments (H)"] = len(mh)
            row["Misaligned"] = int(mh_all["misaligned"].sum())
            row["Unjudged"] = int((mh_all["weighted_score"].isna() & ~mh_all["misaligned"]).sum())
            row["Weighted Score"] = round(mh["weighted_score"].mean(), 3)
            row["Overall (1–5)"] = round(mh["overall"].mean(), 3) if mh["overall"].notna().any() else None
            row["Issues"] = int(mh["issue_count"].sum())
            row["Critical"] = int(mh["critical_count"].sum())
            row["Major"] = int(mh["major_count"].sum())
            row["Minor"] = int(mh["minor_count"].sum())
            for d in DIMENSIONS:
                row[DIMENSION_LABELS[d]] = round(mh[d].mean(), 3) if mh[d].notna().any() else None
        if not iedf.empty and model in iedf["model"].values:
            mi = iedf[iedf["model"] == model]
            row["Segments (IE)"] = len(mi)
            row["Intent Accuracy"] = f"{mi['intent_score'].mean()*100:.1f}%" if mi["intent_score"].notna().any() else None
            row["Entity Score (mean)"] = f"{mi['entity_score'].mean()*100:.1f}%" if mi["entity_score"].notna().any() else None
        rows.append(row)

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

    col1, col2 = st.columns(2)
    if not hdf.empty:
        with col1:
            st.plotly_chart(model_comparison_radar(hdf), use_container_width=True)
            st.caption("Each axis is a dimension (1–5). Larger filled area = better overall quality. "
                       "Uneven shapes reveal where a model under-performs.")
        with col2:
            st.plotly_chart(weighted_score_boxplot(hdf), use_container_width=True)
            st.caption("Box shows 25th–75th percentile. Whiskers extend to min/max (outliers shown as dots). "
                       "Taller boxes = more variable quality; outliers at the bottom are problem segments.")

    if not iedf.empty:
        st.subheader("Intent & Entity — Model Comparison")
        col1, col2 = st.columns(2)
        with col1:
            intent_df = iedf.groupby("model")["intent_score"].mean().reset_index()
            intent_df["intent_score"] = intent_df["intent_score"] * 100
            fig = px.bar(
                intent_df, x="model", y="intent_score",
                title="Intent Accuracy by Model",
                labels={"intent_score": "Accuracy (%)", "model": "Model"},
                color="model",
            )
            fig.update_layout(yaxis=dict(range=[0, 100], ticksuffix="%"), showlegend=False, height=350)
            st.plotly_chart(fig, use_container_width=True)
        with col2:
            entity_df = iedf.groupby("model")["entity_score"].mean().reset_index()
            entity_df["entity_score"] = entity_df["entity_score"] * 100
            fig = px.bar(
                entity_df, x="model", y="entity_score",
                title="Entity Score (mean) by Model",
                labels={"entity_score": "Score (%)", "model": "Model"},
                color="model",
            )
            fig.update_layout(yaxis=dict(range=[0, 100], ticksuffix="%"), showlegend=False, height=350)
            st.plotly_chart(fig, use_container_width=True)


def render_holistic(hdf: pd.DataFrame):
    st.header("Holistic Judge — Deep Dive")
    st.caption(
        "Per-model analysis of 5 evaluation dimensions scored 1–5 by an LLM judge. "
        "Use the **Segment Inspector** at the bottom to read the judge's reasoning for any individual segment."
    )
    _holistic_legend()
    if hdf.empty:
        st.info("No *.holistic.jsonl files found.")
        return

    models = sorted(hdf["model"].unique())
    selected_model = st.selectbox("Model", models)
    mdf_all = hdf[hdf["model"] == selected_model]
    mdf = mdf_all[mdf_all["weighted_score"].notna()]
    n_misaligned = int(mdf_all["misaligned"].sum())
    n_unjudged = int((mdf_all["weighted_score"].isna() & ~mdf_all["misaligned"]).sum())

    # Metrics row
    col1, col2, col3, col4, col5, col6, col7 = st.columns(7)
    col1.metric("Weighted Score", f"{mdf['weighted_score'].mean():.3f}" if len(mdf) else "—")
    col2.metric("Overall (1–5)", f"{mdf['overall'].mean():.2f}" if mdf["overall"].notna().any() else "—")
    col3.metric(
        "Segments",
        len(mdf_all),
        help=f"{len(mdf)} judged, {n_unjudged} unjudged, {n_misaligned} misaligned"
        if (n_unjudged or n_misaligned) else None,
    )
    col4.metric("Misaligned", n_misaligned, help="Rows with alignment_confidence == 0 (skipped at inference)")
    col5.metric("Total Issues", int(mdf["issue_count"].sum()))
    col6.metric("Critical", int(mdf["critical_count"].sum()))
    halluc_affected = mdf[mdf["hallucination_impact"] > 0]
    halluc_pct = 100 * len(halluc_affected) / len(mdf) if len(mdf) else 0
    halluc_impact_affected = halluc_affected["hallucination_impact"].mean() if len(halluc_affected) else 0
    col7.metric(
        "Halluc. Segs",
        f"{halluc_pct:.1f}%",
        help=f"Mean impact (affected only): {halluc_impact_affected:.2f}. "
             f"Denominator excludes {n_unjudged} unjudged segment(s)." if n_unjudged
             else f"Mean impact (affected only): {halluc_impact_affected:.2f}",
    )

    # Radar
    means = {DIMENSION_LABELS[d]: mdf[d].mean() for d in DIMENSIONS if mdf[d].notna().any()}
    st.plotly_chart(radar_chart(means, title=selected_model), use_container_width=True)
    st.caption("Radar chart: each axis represents one dimension (mean score, 1–5). "
               "A balanced pentagon indicates consistent quality across all dimensions. "
               "A dip on one axis shows where this model struggles most.")

    if mdf["satsang_id"].nunique() > 1:
        st.plotly_chart(satsang_heatmap(hdf, selected_model), use_container_width=True)
        st.caption("Heatmap: each row is a satsang, each column is a dimension. "
                   "Green = high score, red = low score. Dark red cells are the worst-performing satsang/dimension combinations.")

    # Score distribution tabs
    st.subheader("Score Distributions")
    st.caption("Each tab shows how scores for that dimension are distributed across segments. "
               "A good model should have most segments at 4–5. "
               "A long tail at 1–2 indicates systematic failures on that dimension.")
    tabs = st.tabs([DIMENSION_LABELS[d] for d in DIMENSIONS])
    for tab, dim in zip(tabs, DIMENSIONS):
        with tab:
            st.plotly_chart(score_distribution_chart(mdf, dim), use_container_width=True)

    # Satsang drill-down
    st.subheader("Satsang Drill-Down")
    satsangs = sorted(mdf["satsang_id"].unique())
    selected_satsang = st.selectbox("Satsang", satsangs)
    sdf = mdf[mdf["satsang_id"] == selected_satsang].sort_values("srt_index")

    # Score progression
    prog_df = sdf[["srt_index", "weighted_score"] + DIMENSIONS].melt(
        id_vars=["srt_index"], var_name="metric", value_name="score",
    )
    prog_df["metric"] = prog_df["metric"].map(lambda m: DIMENSION_LABELS.get(m, m))
    fig = px.line(
        prog_df, x="srt_index", y="score", color="metric",
        title="Score Progression",
        labels={"srt_index": "Segment Index", "score": "Score"},
    )
    fig.update_layout(height=360)
    st.plotly_chart(fig, use_container_width=True)

    # Segment table
    n_ctx = st.slider("Past context lines (N)", min_value=0, max_value=5, value=2, key="hol_ctx_n")
    sdf = sdf.copy()
    inputs_list = sdf["input"].tolist()
    hyps_list = sdf["hypothesis"].tolist()
    sdf["ctx_asr"] = [
        " ↵ ".join(inputs_list[max(0, i - n_ctx):i]) if n_ctx > 0 else ""
        for i in range(len(sdf))
    ]
    sdf["ctx_hyp"] = [
        " ↵ ".join(hyps_list[max(0, i - n_ctx):i]) if n_ctx > 0 else ""
        for i in range(len(sdf))
    ]
    detail_cols = ["srt_index", "weighted_score"] + DIMENSIONS + [
        "issue_count", "critical_count", "major_count", "worst_severity", "issue_types",
        "ctx_asr", "input", "ctx_hyp", "hypothesis",
    ]
    display = sdf[detail_cols].copy()
    display.columns = [
        "Seg#", "Weighted",
        *[DIMENSION_LABELS[d] for d in DIMENSIONS],
        "Issues", "Critical", "Major", "Worst", "Types",
        "Past ASR (ctx)", "ASR", "Past Hypothesis (ctx)", "Hypothesis",
    ]
    st.dataframe(display, use_container_width=True, hide_index=True, height=380)

    # Segment inspector
    st.subheader("Segment Inspector")
    seg_options = sdf["srt_index"].dropna().tolist()
    if not seg_options:
        st.info("No segments with srt_index.")
        return
    seg_idx = st.selectbox("Segment", seg_options, format_func=lambda x: f"Segment {x}")
    seg = sdf[sdf["srt_index"] == seg_idx].iloc[0]

    col_l, col_r = st.columns(2)
    with col_l:
        st.markdown("**Gujarati ASR (Input)**")
        st.text(seg["input"])
        st.markdown("**Human Reference**")
        st.text(seg["reference"])
    with col_r:
        st.markdown("**Model Hypothesis**")
        st.text(seg["hypothesis"])

    for dim in DIMENSIONS:
        reasoning = seg.get(f"{dim}_reasoning", "")
        score = seg.get(dim)
        if reasoning:
            with st.expander(f"{DIMENSION_LABELS[dim]}: {score}/5"):
                st.write(reasoning)

    if seg.get("overall_reasoning"):
        with st.expander(f"Overall: {seg.get('overall')}/5"):
            st.write(seg["overall_reasoning"])

    issues = seg.get("issues_detail") or []
    if issues:
        with st.expander(f"Issues ({len(issues)})"):
            for issue in sorted(issues, key=lambda i: SEVERITY_ORDER.get(i.get("severity", "minor"), 99)):
                sev = issue.get("severity", "")
                itype = issue.get("type", "")
                desc = issue.get("description", "")
                color = SEVERITY_COLORS.get(sev, "#999")
                st.markdown(
                    f"<span style='color:{color}'>**{sev.upper()}**</span> · `{itype}`<br>{desc}",
                    unsafe_allow_html=True,
                )
                st.divider()


def render_intent_entity(iedf: pd.DataFrame):
    st.header("Intent & Entity Judge")
    st.caption(
        "Based on the open-source Sarvam AI evaluation framework: "
        "[github.com/sarvamai/llm_intent_entity](https://github.com/sarvamai/llm_intent_entity/tree/main). "
        "Measures whether the translated output preserves the **intent** (what is being communicated) "
        "and the **entities** (who/what/when/where) of the ground truth."
    )
    _intent_entity_legend()
    if iedf.empty:
        st.info("No *.intent_entity.jsonl files found.")
        return

    models = sorted(iedf["model"].unique())
    selected_model = st.selectbox("Model", models)
    mdf = iedf[iedf["model"] == selected_model]

    scored = mdf[mdf["intent_score"].notna() & mdf["entity_score"].notna()]

    col1, col2, col3, col4 = st.columns(4)
    col1.metric("Segments", len(mdf))
    col2.metric("Scored", len(scored))
    col3.metric(
        "Intent Accuracy",
        f"{scored['intent_score'].mean()*100:.1f}%" if len(scored) else "—",
    )
    col4.metric(
        "Entity Score (mean)",
        f"{scored['entity_score'].mean()*100:.1f}%" if len(scored) else "—",
    )

    # Score distributions
    col1, col2 = st.columns(2)
    with col1:
        intent_counts = scored["intent_score"].value_counts().reset_index()
        intent_counts.columns = ["intent_score", "count"]
        intent_counts["intent_score"] = intent_counts["intent_score"] * 100
        fig = px.bar(
            intent_counts, x="intent_score", y="count",
            title="Intent Score Distribution (0% = failed, 100% = preserved)",
            labels={"intent_score": "Intent Score (%)", "count": "Segments"},
            color="intent_score",
            color_continuous_scale=["#e74c3c", "#2ecc71"],
        )
        fig.update_layout(showlegend=False, height=300, xaxis=dict(ticksuffix="%"))
        st.plotly_chart(fig, use_container_width=True)
        st.caption("Binary score. The bar at **100%** should dominate. "
                   "A large bar at **0%** means the model frequently changed who is doing what, "
                   "reversed roles, or changed statements to questions.")

    with col2:
        entity_pct = scored.copy()
        entity_pct["entity_score"] = entity_pct["entity_score"] * 100
        fig = px.histogram(
            entity_pct, x="entity_score", nbins=20,
            title="Entity Score Distribution",
            labels={"entity_score": "Entity Score (%)", "count": "Segments"},
        )
        fig.update_layout(height=300, xaxis=dict(range=[0, 100], ticksuffix="%"))
        st.plotly_chart(fig, use_container_width=True)
        st.caption("A peak near **100%** means entities (names, dates, places) "
                   "are well preserved. A peak near **0%** indicates systematic entity loss. "
                   "Segments with no entities in the ground truth automatically score 100%.")

    # Scatter: intent vs entity
    st.subheader("Intent vs. Entity Score per Segment")
    scatter_df = scored.copy()
    scatter_df["entity_pct"] = scatter_df["entity_score"] * 100
    scatter_df["intent_pct"] = scatter_df["intent_score"] * 100
    fig = px.scatter(
        scatter_df, x="entity_pct", y="intent_pct",
        hover_data=["srt_index", "hypothesis"],
        opacity=0.6,
        title="Intent Score vs. Entity Score",
        labels={"entity_pct": "Entity Score (%)", "intent_pct": "Intent Score (%)"},
    )
    fig.update_layout(
        yaxis=dict(tickvals=[0, 100], ticksuffix="%"),
        xaxis=dict(range=[0, 100], ticksuffix="%"),
        height=360,
    )
    st.plotly_chart(fig, use_container_width=True)
    st.caption(
        "Each dot is one segment. **Top-right** (intent=100%, entity≈100%) = ideal. "
        "**Bottom-left** (intent=0%, entity≈0%) = critical failures — the speaker's meaning AND key facts are lost. "
        "**Top-left** (intent=100%, entity low) = meaning preserved but specific details dropped. "
        "**Bottom-right** (intent=0%, entity high) = entities correct but the action/subject is wrong. "
        "Hover over any dot to see the hypothesis."
    )

    # Satsang drill-down
    st.subheader("Satsang Drill-Down")
    satsangs = sorted(mdf["satsang_id"].unique())
    selected_satsang = st.selectbox("Satsang", satsangs)
    sdf = mdf[mdf["satsang_id"] == selected_satsang].sort_values("srt_index")

    # Score progression
    prog_sdf = sdf[["srt_index", "intent_score", "entity_score"]].copy()
    prog_sdf["intent_score"] = prog_sdf["intent_score"] * 100
    prog_sdf["entity_score"] = prog_sdf["entity_score"] * 100
    prog = prog_sdf.melt(id_vars=["srt_index"], var_name="metric", value_name="score")
    fig = px.line(
        prog, x="srt_index", y="score", color="metric",
        title="Intent & Entity Score Progression",
        labels={"srt_index": "Segment Index", "score": "Score (%)"},
    )
    fig.update_layout(yaxis=dict(range=[-5, 105], ticksuffix="%"), height=320)
    st.plotly_chart(fig, use_container_width=True)

    # Segment table
    table_cols = [
        "srt_index", "intent_score", "entity_score",
        "input", "reference", "hypothesis",
        "ground_truth_entities", "preserved_entities", "missing_entities",
        "intent_explanation", "entity_explanation",
    ]
    existing_cols = [c for c in table_cols if c in sdf.columns]
    col_labels = {"input": "ASR", "reference": "Ground Truth", "hypothesis": "Hypothesis"}
    st.dataframe(
        sdf[existing_cols].rename(columns=col_labels),
        use_container_width=True, hide_index=True, height=380,
    )

    # Segment inspector
    st.subheader("Segment Inspector")
    seg_options = sdf["srt_index"].dropna().tolist()
    if seg_options:
        seg_idx = st.selectbox("Segment", seg_options, format_func=lambda x: f"Segment {x}")
        seg = sdf[sdf["srt_index"] == seg_idx].iloc[0]
        col_l, col_r = st.columns(2)
        with col_l:
            st.markdown("**Gujarati ASR (Input)**")
            st.text(seg["input"])
            st.markdown("**Human Reference**")
            st.text(seg["reference"])
        with col_r:
            st.markdown("**Model Hypothesis**")
            st.text(seg["hypothesis"])
        c1, c2 = st.columns(2)
        with c1:
            score_color = "green" if seg["intent_score"] == 1 else "red"
            st.markdown(
                f"**Intent Score:** <span style='color:{score_color}'>{seg['intent_score']}</span> "
                f"— {seg.get('intent_explanation', '')}",
                unsafe_allow_html=True,
            )
        with c2:
            entity_val = seg.get("entity_score")
            entity_disp = f"{entity_val*100:.1f}%" if entity_val is not None else "—"
            st.markdown(f"**Entity Score:** {entity_disp} — {seg.get('entity_explanation', '')}")
        st.markdown(f"**GT Entities:** {seg.get('ground_truth_entities', '—')}")
        st.markdown(f"**Preserved:** {seg.get('preserved_entities', '—')}")
        st.markdown(f"**Missing:** {seg.get('missing_entities', '—')}")


def render_issues(hdf: pd.DataFrame):
    st.header("Issues Analysis")
    st.caption(
        "Issues are flagged by the holistic judge within each scored segment. "
        "Each issue has a **type** (what went wrong), a **severity** (how bad it is), "
        "and a quoted description."
    )
    _issues_legend()
    if hdf.empty:
        st.info("No holistic results to analyse.")
        return

    models = sorted(hdf["model"].unique())

    def _pct_segs(mdf: pd.DataFrame, issue_type: Optional[str], severity: str) -> str:
        n = len(mdf)
        if n == 0:
            return "0.0%"
        count = sum(
            1 for issues in mdf["issues_detail"]
            if any(
                (issue_type is None or i.get("type") == issue_type)
                and i.get("severity") == severity
                for i in (issues or [])
            )
        )
        return f"{100 * count / n:.1f}%"

    def _pct_val(mdf: pd.DataFrame, issue_type: Optional[str], severity: str) -> float:
        return float(_pct_segs(mdf, issue_type, severity).rstrip("%"))

    # ── Section 1: High-level summary table ───────────────────────────────────
    st.subheader("High-Level Summary")
    st.caption(
        "Issue-rate percentages use **judged segments** as the denominator "
        "(unjudged and misaligned rows are excluded). "
        "**Misaligned** = alignment_confidence == 0 (skipped at inference). "
        "**Unjudged** = judge produced no score for some other reason."
    )
    agg_rows = []
    for model in models:
        mdf_all = hdf[hdf["model"] == model]
        mdf = mdf_all[mdf_all["weighted_score"].notna()]
        all_issues = [i for issues in mdf["issues_detail"] for i in issues]
        n = len(mdf)
        agg_rows.append({
            "Model": model,
            "Segments": len(mdf_all),
            "Misaligned": int(mdf_all["misaligned"].sum()),
            "Unjudged": int((mdf_all["weighted_score"].isna() & ~mdf_all["misaligned"]).sum()),
            "With Issues": int(mdf["has_issues"].sum()),
            "% With Issues": f"{100 * mdf['has_issues'].mean():.1f}%" if n else "—",
            "Total Issues": len(all_issues),
            "% Critical segs": _pct_segs(mdf, None, "critical"),
            "% Major segs": _pct_segs(mdf, None, "major"),
            "% Minor segs": _pct_segs(mdf, None, "minor"),
            "% Halluc Segs": f"{100 * (mdf['hallucination_impact'] > 0).mean():.1f}%" if n else "—",
            "Total Halluc Impact": round(mdf["hallucination_impact"].sum(), 3),
            "Halluc Impact (affected)": round(
                mdf.loc[mdf["hallucination_impact"] > 0, "hallucination_impact"].mean(), 3
            ) if (mdf["hallucination_impact"] > 0).any() else 0.0,
            "% Safety Segs": f"{100 * (mdf['safety_impact'] > 0).mean():.1f}%" if n else "—",
            "Total Safety Impact": round(mdf["safety_impact"].sum(), 3),
        })
    st.dataframe(pd.DataFrame(agg_rows), use_container_width=True, hide_index=True)

    # ── Section 2: Issue Distribution tables ──────────────────────────────────
    st.subheader("Issue Distribution")

    def _pct_segs_any_sev(mdf: pd.DataFrame, issue_type: str) -> str:
        n = len(mdf)
        if n == 0:
            return "0.0%"
        count = sum(
            1 for issues in mdf["issues_detail"]
            if any(i.get("type") == issue_type for i in (issues or []))
        )
        return f"{100 * count / n:.1f}%"

    # Table 1: issue type × model (count + % segs)
    st.markdown("**Issue counts and % of segments affected, by type and model**")
    st.caption(
        "Count = total flagged instances of that type. "
        "% Segs = % of all evaluated segments that contain at least one instance "
        "(denominator = total segments for that model)."
    )
    dist_rows = []
    for itype in ISSUE_TYPES:
        row = {"Issue Type": itype}
        for model in models:
            mdf_m = hdf[(hdf["model"] == model) & hdf["weighted_score"].notna()]
            all_issues_m = [i for issues in mdf_m["issues_detail"] for i in issues]
            count = sum(1 for i in all_issues_m if i.get("type") == itype)
            pct = _pct_segs_any_sev(mdf_m, itype)
            row[f"{model} — Count"] = count
            row[f"{model} — % Segs"] = pct
        dist_rows.append(row)
    st.dataframe(pd.DataFrame(dist_rows), use_container_width=True, hide_index=True)

    # Table 2: severity breakdown per model
    st.markdown("**Severity breakdown by model**")
    st.caption(
        "Count = total issues of that severity across all segments. "
        "% Segs = % of segments with at least one issue at that severity "
        "(a segment with 2 critical issues counts once)."
    )
    sev_dist_rows = []
    for model in models:
        mdf_all = hdf[hdf["model"] == model]
        mdf_m = mdf_all[mdf_all["weighted_score"].notna()]
        all_issues_m = [i for issues in mdf_m["issues_detail"] for i in issues]
        sev_dist_rows.append({
            "Model": model,
            "Total Segments": len(mdf_all),
            "Misaligned": int(mdf_all["misaligned"].sum()),
            "Unjudged": int((mdf_all["weighted_score"].isna() & ~mdf_all["misaligned"]).sum()),
            "Total Issues": len(all_issues_m),
            "Critical — Count": sum(1 for i in all_issues_m if i.get("severity") == "critical"),
            "Critical — % Segs": _pct_segs(mdf_m, None, "critical"),
            "Major — Count": sum(1 for i in all_issues_m if i.get("severity") == "major"),
            "Major — % Segs": _pct_segs(mdf_m, None, "major"),
            "Minor — Count": sum(1 for i in all_issues_m if i.get("severity") == "minor"),
            "Minor — % Segs": _pct_segs(mdf_m, None, "minor"),
        })
    st.dataframe(pd.DataFrame(sev_dist_rows), use_container_width=True, hide_index=True)

    # ── Section 3: Hallucination & Safety Impact ───────────────────────────────
    st.subheader("Hallucination & Safety Impact")
    with st.expander("📖 How these numbers are calculated", expanded=False):
        st.markdown(
            "Each hallucination/safety issue is assigned a **severity weight**: "
            "critical = 1.0, major = 0.5, minor = 0.1.  \n"
            "The **impact score for a segment** = sum of weights of all hallucination issues in that segment "
            "(e.g. one critical + one minor = 1.1).  \n\n"
            "**Columns explained:**\n"
            "- **Segs w/ Halluc** — how many segments contain at least one hallucination "
            "(denominator = all evaluated segments for that model).\n"
            "- **% Affected** — `segs w/ halluc / total segs × 100`.\n"
            "- **% Critical / Major / Minor** — % of *all* segments that have at least one "
            "hallucination of that severity. Denominator = total segments.\n"
            "- **Impact Sum** — total severity weight summed across every hallucination issue "
            "in every segment. Larger = more or worse hallucinations overall.\n"
            "- **Impact / Affected Seg** — `Impact Sum / segs w/ halluc`. "
            "Tells you how severe the hallucinations are *within* affected segments, "
            "ignoring the many clean segments. High value = hallucinations that do appear tend to be serious.\n"
            "- **Impact / All Segs** — `Impact Sum / total segs`. "
            "Diluted by all clean segments — useful for comparing models on the same dataset size."
        )

    for impact_col, label in [("hallucination_impact", "Hallucination"), ("safety_impact", "Safety")]:
        st.markdown(f"**{label}**")
        impact_rows = []
        for model in models:
            mdf_m = hdf[(hdf["model"] == model) & hdf["weighted_score"].notna()]
            n = len(mdf_m)
            affected = mdf_m[mdf_m[impact_col] > 0]
            n_aff = len(affected)
            impact_sum = round(mdf_m[impact_col].sum(), 2)
            itype = "hallucination" if label == "Hallucination" else "content_safety"
            n_crit = sum(1 for issues in mdf_m["issues_detail"] if any(
                i.get("type") == itype and i.get("severity") == "critical" for i in (issues or [])))
            n_maj = sum(1 for issues in mdf_m["issues_detail"] if any(
                i.get("type") == itype and i.get("severity") == "major" for i in (issues or [])))
            n_min = sum(1 for issues in mdf_m["issues_detail"] if any(
                i.get("type") == itype and i.get("severity") == "minor" for i in (issues or [])))
            impact_rows.append({
                "Model": model,
                "Total Segs": n,
                "Segs w/ Issue": f"{n_aff} / {n}",
                "% Affected": f"{100 * n_aff / n:.1f}%" if n else "0.0%",
                "% Critical segs": f"{100 * n_crit / n:.1f}%  ({n_crit}/{n})" if n else "0%",
                "% Major segs": f"{100 * n_maj / n:.1f}%  ({n_maj}/{n})" if n else "0%",
                "% Minor segs": f"{100 * n_min / n:.1f}%  ({n_min}/{n})" if n else "0%",
                "Impact Sum": impact_sum,
                "Impact / Affected Seg": f"{affected[impact_col].mean():.2f}  (÷ {n_aff} affected)" if n_aff else "—",
                "Impact / All Segs": f"{impact_sum / n:.3f}  (÷ {n} total)" if n else "—",
            })
        st.dataframe(pd.DataFrame(impact_rows), use_container_width=True, hide_index=True)

    st.divider()

    # ── Section 4: Per-issue-type breakdown ────────────────────────────────────
    st.subheader("Per-Issue-Type Breakdown")
    st.caption(
        "Each section below covers one error type: definition, cross-model stats, "
        "severity split, and a table of all affected segments for the selected model."
    )

    detail_model = st.selectbox("Model for segment tables", models, key="detail_model_sel")
    detail_mdf = hdf[hdf["model"] == detail_model]

    for itype in ISSUE_TYPES:
        defn = ISSUE_DEFINITIONS.get(itype, "")
        # count across all models for the header badge
        total_across_models = sum(
            sum(1 for i in (row.get("issues_detail") or []) if i.get("type") == itype)
            for _, row in hdf.iterrows()
        )
        label = f"**`{itype}`** — {defn}  ·  {total_across_models} total instances"
        with st.expander(label, expanded=False):
            # Cross-model stats table
            stats_rows = []
            for model in models:
                mdf_m = hdf[hdf["model"] == model]
                all_issues_m = [i for issues in mdf_m["issues_detail"] for i in issues]
                count = sum(1 for i in all_issues_m if i.get("type") == itype)
                stats_rows.append({
                    "Model": model,
                    "Count": count,
                    "% Segs": _pct_segs(mdf_m, itype, "critical")[:-1] + "% crit / "
                              + _pct_segs(mdf_m, itype, "major")[:-1] + "% maj / "
                              + _pct_segs(mdf_m, itype, "minor") + " min",
                    "% Critical": _pct_segs(mdf_m, itype, "critical"),
                    "% Major": _pct_segs(mdf_m, itype, "major"),
                    "% Minor": _pct_segs(mdf_m, itype, "minor"),
                })
            st.dataframe(
                pd.DataFrame(stats_rows)[["Model", "Count", "% Critical", "% Major", "% Minor"]],
                use_container_width=True, hide_index=True,
            )

            # Severity pie for this type (all models combined)
            sev_counts = {}
            for model in models:
                mdf_m = hdf[hdf["model"] == model]
                for sev in ["critical", "major", "minor"]:
                    sev_counts[sev] = sev_counts.get(sev, 0) + sum(
                        1 for issues in mdf_m["issues_detail"]
                        for i in (issues or [])
                        if i.get("type") == itype and i.get("severity") == sev
                    )
            sev_counts = {k: v for k, v in sev_counts.items() if v > 0}
            if sev_counts:
                col_pie, _ = st.columns([1, 2])
                with col_pie:
                    fig = px.pie(
                        names=list(sev_counts.keys()),
                        values=list(sev_counts.values()),
                        title=f"{itype} — severity split (all models)",
                        hole=0.4,
                        color=list(sev_counts.keys()),
                        color_discrete_map=SEVERITY_COLORS,
                    )
                    fig.update_traces(textposition="inside", textinfo="percent+label")
                    fig.update_layout(height=260, showlegend=False, margin=dict(t=40, b=10, l=10, r=10))
                    st.plotly_chart(fig, use_container_width=True)

            # Segment table for the selected model
            seg_rows = []
            for _, row in detail_mdf.sort_values(["satsang_id", "srt_index"]).iterrows():
                type_issues = [
                    i for i in (row.get("issues_detail") or [])
                    if i.get("type") == itype
                ]
                for issue in sorted(type_issues, key=lambda i: SEVERITY_ORDER.get(i.get("severity", "minor"), 99)):
                    seg_rows.append({
                        "Satsang": row["satsang_id"],
                        "Seg #": row["srt_index"],
                        "Score": round(row["weighted_score"], 2) if row["weighted_score"] is not None else None,
                        "Severity": issue.get("severity", ""),
                        "Description": issue.get("description", "")[:160],
                        "Hypothesis": row["hypothesis"][:100],
                    })
            if seg_rows:
                st.caption(
                    f"**{len(seg_rows)} instance(s)** of `{itype}` in **{detail_model}** "
                    f"— sorted critical → major → minor."
                )
                st.dataframe(pd.DataFrame(seg_rows), use_container_width=True, hide_index=True, height=300)
            else:
                st.info(f"No `{itype}` issues found for {detail_model}.")

    st.divider()

    # ── Section 5: Filterable all-issues table ─────────────────────────────────
    st.subheader("All Segments with Issues")
    sel_model = st.selectbox("Model", models, key="issues_model_sel")
    mdf = hdf[hdf["model"] == sel_model]
    flagged = mdf[mdf["has_issues"]]

    type_filter = st.multiselect("Filter by type", ISSUE_TYPES, default=[], key="type_filter")
    sev_filter = st.multiselect(
        "Filter by severity", ["critical", "major", "minor"],
        default=["critical", "major"], key="sev_filter",
    )
    n_ctx = st.slider("Past context lines (N)", min_value=0, max_value=5, value=2, key="issues_ctx_n")

    # Build (satsang_id, srt_index) → (past_asr, past_hyp) from ALL segments (not just flagged),
    # so context is correct even when preceding segments have no issues.
    ctx_lookup: dict[tuple, tuple[str, str]] = {}
    for sat_id, gdf in mdf.groupby("satsang_id"):
        gdf_sorted = gdf.sort_values("srt_index")
        inputs = gdf_sorted["input"].tolist()
        hyps = gdf_sorted["hypothesis"].tolist()
        for i, srt_idx in enumerate(gdf_sorted["srt_index"]):
            ctx_lookup[(sat_id, srt_idx)] = (
                " ↵ ".join(inputs[max(0, i - n_ctx):i]) if n_ctx > 0 else "",
                " ↵ ".join(hyps[max(0, i - n_ctx):i]) if n_ctx > 0 else "",
            )

    table_rows = []
    for _, row in flagged.sort_values(["satsang_id", "srt_index"]).iterrows():
        issues = row.get("issues_detail") or []
        filtered = [
            i for i in issues
            if (not type_filter or i.get("type") in type_filter)
            and (not sev_filter or i.get("severity") in sev_filter)
        ]
        if not filtered:
            continue
        worst = min((i.get("severity", "minor") for i in filtered), key=lambda s: SEVERITY_ORDER.get(s, 99))
        ctx_asr, ctx_hyp = ctx_lookup.get((row["satsang_id"], row["srt_index"]), ("", ""))
        table_rows.append({
            "Satsang": row["satsang_id"],
            "Seg #": row["srt_index"],
            "Weighted": round(row["weighted_score"], 2) if row["weighted_score"] is not None else None,
            "Past ASR (ctx)": ctx_asr,
            "ASR": row.get("input", ""),
            "Past Hypothesis (ctx)": ctx_hyp,
            "Ground Truth": row.get("reference", ""),
            "Hypothesis": row["hypothesis"],
            "# Issues": len(filtered),
            "Worst": worst,
            "Types": ", ".join(sorted({i.get("type", "") for i in filtered})),
            "Descriptions": " | ".join(i.get("description", "") for i in filtered),
        })

    if table_rows:
        st.markdown(f"**{len(table_rows)} segments** match filters")
        st.dataframe(pd.DataFrame(table_rows), use_container_width=True, hide_index=True, height=420)
    else:
        st.info("No segments match the current filters.")

    st.divider()

    # ── Section 6: Cross-model segment comparison ──────────────────────────────
    st.subheader("Cross-Model Segment Comparison")
    seg_key = ["satsang_id", "srt_index"]
    if not hdf.empty:
        all_segs = hdf[seg_key + ["input"]].drop_duplicates().sort_values(seg_key).reset_index(drop=True)
        result = all_segs.rename(columns={"satsang_id": "Satsang", "srt_index": "Seg #"})
        result["Segment Text"] = result["input"].str[:100]
        result = result.drop(columns=["input"])
        for model in models:
            mdf_m = hdf[hdf["model"] == model][seg_key + ["issue_count", "worst_severity"]]
            mdf_m = mdf_m.rename(columns={
                "issue_count": f"{model} — Issues",
                "worst_severity": f"{model} — Worst",
            })
            result = result.merge(mdf_m, left_on=["Satsang", "Seg #"], right_on=seg_key, how="left")
            result = result.drop(columns=[c for c in seg_key if c in result.columns and c not in ["Satsang", "Seg #"]])
        show_flagged = st.checkbox("Show only segments with at least one issue", value=True)
        issue_cols = [c for c in result.columns if c.endswith("— Issues")]
        if show_flagged and issue_cols:
            mask = result[issue_cols].fillna(0).gt(0).any(axis=1)
            result = result[mask]
        st.dataframe(result, use_container_width=True, hide_index=True, height=480)


# ── Severity Bubbles ──────────────────────────────────────────────────────────

_BUBBLE_COLORS = {
    "critical": "#E53935",
    "major":    "#FFC107",
    "green":    "#4CAF50",
}

# Darker shades used as ball borders for depth
_BUBBLE_BORDER_COLORS = {
    "critical": "#B71C1C",
    "major":    "#A37800",
    "green":    "#1B5E20",
}


def _bubble_points(n: int, seed: int = 42):
    """Uniform random (x, y) inside unit disk. sqrt(r) ensures area-uniform distribution."""
    rng = np.random.default_rng(seed)
    r = np.sqrt(rng.uniform(0.0, 1.0, size=n))
    theta = rng.uniform(0.0, 2 * np.pi, size=n)
    return r * np.cos(theta), r * np.sin(theta)


def _dot_size(n: int) -> float:
    if n < 500:
        return 16
    if n < 2000:
        return 12
    if n < 5000:
        return 8
    return 6


def _build_bubble_fig(
    model: str,
    model_df: pd.DataFrame,
    selected_point_idx: Optional[int] = None,
) -> go.Figure:
    sev = model_df["worst_severity"].fillna("none")
    colors = [
        _BUBBLE_COLORS["critical"] if s == "critical"
        else _BUBBLE_COLORS["major"] if s == "major"
        else _BUBBLE_COLORS["green"]
        for s in sev
    ]
    n_red    = int((sev == "critical").sum())
    n_yellow = int((sev == "major").sum())
    n_green  = len(model_df) - n_red - n_yellow
    n_total  = len(model_df)

    x, y = _bubble_points(n_total)
    theta_c = np.linspace(0, 2 * np.pi, 300)

    sev_arr = sev.to_numpy()
    dot_sz = _dot_size(n_total)

    # Build per-severity index masks so we can layer: green → major → critical
    masks = {
        "green":    (sev_arr != "critical") & (sev_arr != "major"),
        "major":    sev_arr == "major",
        "critical": sev_arr == "critical",
    }
    dot_colors = {
        "green":    _BUBBLE_COLORS["green"],
        "major":    _BUBBLE_COLORS["major"],
        "critical": _BUBBLE_COLORS["critical"],
    }

    # Map original DataFrame position → trace-local index for each severity layer.
    # This is needed so click events (point_index within a trace) can be mapped back.
    orig_indices: dict[str, list[int]] = {k: [] for k in masks}

    fig = go.Figure()
    for layer in ("green", "major", "critical"):
        mask = masks[layer]
        idxs = np.where(mask)[0]
        orig_indices[layer] = idxs.tolist()
        if len(idxs) == 0:
            continue
        layer_customdata = [
            (
                model_df["srt_index"].iloc[i],
                sev_arr[i],
                model_df["hypothesis"].iloc[i],
                model_df["reference"].iloc[i],
                int(i),  # original positional index in mdf for selection mapping
            )
            for i in idxs
        ]
        fig.add_trace(go.Scatter(
            x=x[idxs], y=y[idxs],
            mode="markers",
            marker=dict(
                color=dot_colors[layer],
                size=dot_sz,
                opacity=0.88,
                line=dict(width=0),
            ),
            customdata=layer_customdata,
            hovertemplate=(
                f"<b style='font-size:16px'>Seg %{{customdata[0]}}</b>"
                f" &nbsp;<span style='color:{dot_colors[layer]}'>●</span>"
                f" <span style='color:{dot_colors[layer]};font-weight:bold'>"
                f"%{{customdata[1]}}</span><br>"
                "<br><span style='color:rgba(255,255,255,0.25)'>──────────────────────</span><br>"
                "<span style='color:#aaaaaa;font-size:12px'>HYPOTHESIS</span><br>"
                "<span style='font-size:15px'>%{customdata[2]}</span><br><br>"
                "<span style='color:#aaaaaa;font-size:12px'>GROUND TRUTH</span><br>"
                "<span style='font-size:15px'>%{customdata[3]}</span>"
                "<extra></extra>"
            ),
            hoverlabel=dict(
                bgcolor="rgba(10,10,10,0.98)",
                bordercolor=dot_colors[layer],
                font=dict(size=15, color="#ffffff", family="sans-serif"),
                align="left",
                namelength=0,
            ),
            name=layer,
            showlegend=False,
        ))

    # Overlay the selected dot at 2× size with a white outline (topmost trace)
    if selected_point_idx is not None and 0 <= selected_point_idx < n_total:
        fig.add_trace(go.Scatter(
            x=[x[selected_point_idx]],
            y=[y[selected_point_idx]],
            mode="markers",
            marker=dict(
                color=colors[selected_point_idx],
                size=dot_sz * 2,
                opacity=1.0,
                line=dict(width=0),
            ),
            hoverinfo="skip",
            showlegend=False,
        ))

    fig.add_trace(go.Scatter(
        x=np.cos(theta_c),
        y=np.sin(theta_c),
        mode="lines",
        line=dict(color="rgba(255,255,255,0.5)", width=2),
        hoverinfo="skip",
        showlegend=False,
    ))
    title_html = (
        f"<b style='color:#ffffff'>{model}</b><br>"
        f"<span style='color:{_BUBBLE_COLORS['critical']}'>● {n_red} critical</span>  "
        f"<span style='color:{_BUBBLE_COLORS['major']}'>● {n_yellow} major</span>  "
        f"<span style='color:{_BUBBLE_COLORS['green']}'>● {n_green} minor/clean</span>  "
        f"<span style='color:#aaaaaa'>({n_total} total)</span>"
    )
    fig.update_layout(
        title=dict(text=title_html, x=0.5, xanchor="center", font=dict(size=13)),
        xaxis=dict(visible=False, range=[-1.12, 1.12], scaleanchor="y", scaleratio=1),
        yaxis=dict(visible=False, range=[-1.12, 1.12]),
        plot_bgcolor="#0d0d0d",
        paper_bgcolor="#0d0d0d",
        height=680,
        margin=dict(t=90, b=20, l=20, r=20),
        showlegend=False,
    )
    return fig


def _render_bubble_segment_details(model_df: pd.DataFrame, point_idx: int) -> None:
    """Show hypothesis and ground truth for the clicked segment."""
    if point_idx < 0 or point_idx >= len(model_df):
        return
    row = model_df.iloc[point_idx]
    sev = row.get("worst_severity") or "none"
    sev_color = (
        _BUBBLE_COLORS["critical"] if sev == "critical"
        else _BUBBLE_COLORS["major"] if sev == "major"
        else _BUBBLE_COLORS["green"]
    )
    st.markdown(
        f"<span style='color:{sev_color}'>●</span> **Seg {row['srt_index']}** — `{sev}`",
        unsafe_allow_html=True,
    )
    st.markdown(f"**Hypothesis:** {row['hypothesis']}")
    st.markdown(f"**Ground truth:** {row['reference']}")


def _render_bubble_stats(mdf: pd.DataFrame) -> None:
    """Render a vertical stats panel showing critical/major/clean % breakdowns."""
    sev = mdf["worst_severity"].fillna("none")
    n_total   = len(mdf)
    n_crit    = int((sev == "critical").sum())
    n_major   = int((sev == "major").sum())
    n_clean   = n_total - n_crit - n_major
    def pct(n: int) -> str:
        return f"{100 * n / n_total:.1f}%" if n_total else "—"

    def _card(color: str, border: str, label: str, count: int, percent: str) -> str:
        return (
            f"<div style='"
            f"background:rgba(255,255,255,0.04);"
            f"border-left:4px solid {border};"
            f"border-radius:6px;"
            f"padding:12px 14px;"
            f"margin-bottom:12px;"
            f"'>"
            f"<div style='color:{color};font-weight:700;font-size:13px;letter-spacing:0.05em'>"
            f"{label}</div>"
            f"<div style='color:#ffffff;font-size:28px;font-weight:700;line-height:1.2'>"
            f"{percent}</div>"
            f"<div style='color:#aaaaaa;font-size:12px'>{count:,} segments</div>"
            f"</div>"
        )

    st.markdown(
        "<div style='padding-top:80px'>"
        + _card(_BUBBLE_COLORS["critical"], _BUBBLE_BORDER_COLORS["critical"],
                "CRITICAL", n_crit, pct(n_crit))
        + _card(_BUBBLE_COLORS["major"], _BUBBLE_BORDER_COLORS["major"],
                "MAJOR", n_major, pct(n_major))
        + _card(_BUBBLE_COLORS["green"], _BUBBLE_BORDER_COLORS["green"],
                "MINOR / CLEAN", n_clean, pct(n_clean))
        + f"<div style='color:#666;font-size:11px;margin-top:4px'>{n_total:,} total</div>"
        + "</div>",
        unsafe_allow_html=True,
    )


def _render_bubble_model(m: str, mdf: pd.DataFrame) -> None:
    """Render one bubble chart with click-to-highlight and segment detail panel."""
    sel_key = f"bubble_sel_{m}"
    selected_idx = st.session_state.get(sel_key)

    col_chart, col_stats = st.columns([5, 1])

    with col_stats:
        _render_bubble_stats(mdf)

    with col_chart:
        fig = _build_bubble_fig(m, mdf, selected_point_idx=selected_idx)
        event = st.plotly_chart(
            fig,
            use_container_width=True,
            on_select="rerun",
            key=f"bubble_chart_{m}",
        )

    # Update selection from click event.
    # curve_number 0/1/2 = green/major/critical layers; curve_number for the
    # selected-dot overlay and circle outline are always the last two — skip them.
    # The original mdf position is stored in customdata[4].
    pts = [
        p for p in (event.selection.points if event.selection else [])
        if len(p.get("customdata", [])) >= 5
    ]
    if pts:
        new_idx = int(pts[0]["customdata"][4])
        if new_idx != selected_idx:
            st.session_state[sel_key] = new_idx
            selected_idx = new_idx

    if selected_idx is not None:
        with st.expander("Segment details", expanded=True):
            _render_bubble_segment_details(mdf, selected_idx)


def render_severity_bubbles(hdf: pd.DataFrame):
    st.header("Severity Bubbles")
    st.caption(
        "Each dot is one evaluated segment. Hover to inspect. Click to pin the segment details below."
    )
    # CSS: scale up the SVG marker path on hover so dots visually grow
    st.markdown(
        """
        <style>
        .stPlotlyChart g.trace.scatter .point path {
            transition: transform 0.12s ease, opacity 0.12s ease;
            transform-box: fill-box;
            transform-origin: center;
        }
        .stPlotlyChart g.trace.scatter .point path:hover {
            transform: scale(1.6);
            opacity: 1 !important;
        }
        </style>
        """,
        unsafe_allow_html=True,
    )
    if hdf.empty:
        st.info("No holistic results found.")
        return

    st.markdown(
        f"<span style='color:{_BUBBLE_COLORS['critical']}'>●</span> **Critical** &nbsp;&nbsp;"
        f"<span style='color:{_BUBBLE_COLORS['major']}'>●</span> **Major** &nbsp;&nbsp;"
        f"<span style='color:{_BUBBLE_COLORS['green']}'>●</span> **Minor / No issue**",
        unsafe_allow_html=True,
    )

    models = sorted(hdf["model"].unique())
    m = st.selectbox("Model", models, key="bubble_model_select")
    _render_bubble_model(m, hdf[hdf["model"] == m].reset_index(drop=True))


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


def main():
    st.set_page_config(
        page_title="Translation Quality Dashboard v2",
        layout="wide",
        initial_sidebar_state="expanded",
    )
    st.title("Translation Quality Dashboard v2")

    # HF Spaces injects secrets as env vars; fall back to os.environ when no secrets.toml
    try:
        token: Optional[str] = st.secrets.get("HF_TOKEN") or None
        default_branch: str = st.secrets.get("HF_DEFAULT_BRANCH", "main")
        default_subfolder: str = st.secrets.get(
            "HF_DEFAULT_SUBFOLDER",
            "",
        )
    except Exception:
        token = os.environ.get("HF_TOKEN") or None
        default_branch = os.environ.get("HF_DEFAULT_BRANCH", "main")
        default_subfolder = os.environ.get(
            "HF_DEFAULT_SUBFOLDER",
            "",
        )

    use_hf = st.sidebar.checkbox("Load from HuggingFace", value=True)
    if use_hf:
        try:
            branches = _list_hf_branches(token)
        except Exception:
            branches = [default_branch]
        branch = st.sidebar.selectbox(
            "HF Branch",
            branches,
            index=branches.index(default_branch) if default_branch in branches else 0,
        )
        if st.sidebar.button("Refresh data from HF"):
            st.cache_resource.clear()
            st.cache_data.clear()
            st.rerun()
        with st.sidebar:
            with st.spinner("Loading data from HuggingFace…"):
                local_root = _hf_snapshot(branch, token)
        subfolders = ["— (all)"] + sorted([
            d for d in os.listdir(local_root)
            if os.path.isdir(os.path.join(local_root, d)) and not d.startswith(".")
        ])
        default_idx = subfolders.index(default_subfolder) if default_subfolder in subfolders else 0
        subfolder = st.sidebar.selectbox("Subfolder (optional)", subfolders, index=default_idx)
        eval_dir = local_root if subfolder == "— (all)" else os.path.join(local_root, subfolder)
    else:
        eval_dir = st.sidebar.text_input("Eval results directory", value=_DEFAULT_EVAL_DIR)

    hdf = load_holistic(eval_dir)
    iedf = load_intent_entity(eval_dir)

    if hdf.empty and iedf.empty:
        st.error(
            f"No results found in `{eval_dir}/`. "
            "Run `scripts/eval/llm_judge.py` to generate *.holistic.jsonl or *.intent_entity.jsonl files."
        )
        return

    n_models = len(set(
        list(hdf["model"].unique() if not hdf.empty else [])
        + list(iedf["model"].unique() if not iedf.empty else [])
    ))
    n_segs_h = len(hdf)
    n_segs_ie = len(iedf)
    st.sidebar.markdown(
        f"**{n_models}** models  \n"
        f"**{n_segs_h:,}** holistic segments  \n"
        f"**{n_segs_ie:,}** intent/entity segments"
    )
    st.sidebar.divider()
    st.sidebar.markdown(
        "**Metric frameworks**  \n"
        "Holistic judge: custom 5-dimension LLM-as-judge  \n"
        "Intent & Entity: [Sarvam AI llm_intent_entity](https://github.com/sarvamai/llm_intent_entity/tree/main)  \n"
        "  \n"
        "**Score ranges**  \n"
        "Holistic dimensions: 1–5 (higher = better)  \n"
        "Weighted score: 1.0–5.0 (higher = better)  \n"
        "Intent accuracy: 0–100% (higher = better)  \n"
        "Entity score: 0–100% (higher = better)  \n"
        "Halluc./Safety impact: 0+ (lower = better)"
    )

    # Model filter
    all_models = sorted(set(
        list(hdf["model"].unique() if not hdf.empty else [])
        + list(iedf["model"].unique() if not iedf.empty else [])
    ))
    selected_models = st.sidebar.multiselect("Filter models", all_models, default=all_models)
    if selected_models:
        if not hdf.empty:
            hdf = hdf[hdf["model"].isin(selected_models)]
        if not iedf.empty:
            iedf = iedf[iedf["model"].isin(selected_models)]

    page = st.sidebar.radio(
        "Page",
        ["Overview", "Holistic Deep Dive", "Intent & Entity", "Issues Analysis", "Severity Bubbles"],
    )

    if page == "Overview":
        render_overview(hdf, iedf)
    elif page == "Holistic Deep Dive":
        render_holistic(hdf)
    elif page == "Intent & Entity":
        render_intent_entity(iedf)
    elif page == "Issues Analysis":
        render_issues(hdf)
    elif page == "Severity Bubbles":
        render_severity_bubbles(hdf)


if __name__ == "__main__":
    main()