src/components/visualizer.py

"""
Document Visualizer — LLM Subject Extraction Demo

Renders aligned agenda-item / subject comparisons for a single document.
Shows predicted vs ground-truth side by side with IoU scores, topic badges,
and full text preview.
"""

import html as _html
import streamlit as st
import plotly.express as px
import pandas as pd
from typing import Dict, Any, List, Optional


# ---------------------------------------------------------------------------
# Colour helpers
# ---------------------------------------------------------------------------

_MATCH_COLORS = {
    "high":   "#2ecc71",   # IoU ≥ 0.8
    "medium": "#f39c12",   # IoU ≥ 0.5
    "low":    "#e74c3c",   # IoU < 0.5
}

def _iou_color(iou: Optional[float]) -> str:
    if iou is None:
        return "#95a5a6"
    if iou >= 0.8:
        return _MATCH_COLORS["high"]
    if iou >= 0.5:
        return _MATCH_COLORS["medium"]
    return _MATCH_COLORS["low"]


def _topic_badge(topic: str, color: str) -> str:
    return (
        f'<span style="background:{color};color:white;padding:3px 10px;'
        f'border-radius:12px;font-size:0.78rem;margin:2px 4px 2px 0;'
        f'display:inline-block;">{_html.escape(topic)}</span>'
    )


def _card(
    title: str,
    text: str,
    topics: List[str],
    theme: str,
    border_color: str,
    topic_color_fn,
    extra_html: str = "",
    missing: bool = False,
) -> str:
    """Build an HTML card for one side of the comparison."""
    if missing:
        return (
            f'<div style="border:2px dashed {border_color};border-radius:10px;'
            f'padding:14px;margin:6px 0;background:#fafafa;color:#aaa;font-style:italic;">'
            f'<strong>{_html.escape(title)}</strong><br/>—</div>'
        )

    badges = "".join(_topic_badge(t, topic_color_fn(t)) for t in topics)
    theme_html = (
        f'<div style="background:#edf6ff;border-left:4px solid #2196f3;'
        f'padding:7px 10px;margin:8px 0;font-style:italic;font-size:0.85rem;color:#111;">'
        f'📝 {_html.escape(theme)}</div>'
        if theme else ""
    )
    text_html = (
        f'<div style="font-size:0.82rem;color:#333;line-height:1.55;'
        f'white-space:pre-wrap;max-height:180px;overflow-y:auto;">'
        f'{_html.escape(text)}</div>'
    )
    return (
        f'<div style="border:2px solid {border_color};border-radius:10px;'
        f'padding:14px;margin:6px 0;background:white;">'
        f'<div style="font-weight:700;color:{border_color};font-size:0.95rem;margin-bottom:6px;">'
        f'{_html.escape(title)}</div>'
        f'{badges}<br/>{theme_html}{text_html}{extra_html}'
        f'</div>'
    )


# ---------------------------------------------------------------------------
# Per-document evaluation summary bar
# ---------------------------------------------------------------------------

def _render_doc_metrics(evaluation: Dict[str, Any]) -> None:
    ai = evaluation.get("agenda_items", {})
    subj = evaluation.get("subjects", {})

    st.markdown("#### 📊 Document Metrics")
    cols = st.columns(6)
    cols[0].metric("AI Boundary F1", f"{ai.get('boundary_f1', 0):.3f}")
    cols[1].metric("AI BED F-measure", f"{ai.get('bed_fmeasure', 0):.3f}")
    cols[2].metric("AI Boundary Sim", f"{ai.get('boundary_similarity', 0):.3f}")
    cols[3].metric("Subj Boundary F1", f"{subj.get('boundary_f1', 0):.3f}")
    cols[4].metric("Subj BED F-measure", f"{subj.get('bed_fmeasure', 0):.3f}")
    cols[5].metric("Subj Theme Acc", f"{subj.get('theme_accuracy', 0):.3f}")

    overall = evaluation.get("overall", {})
    o_cols = st.columns(4)
    o_cols[0].metric("AI Predicted", overall.get("total_agenda_items_predicted", "—"))
    o_cols[1].metric("AI GT", overall.get("total_agenda_items_gt", "—"))
    o_cols[2].metric("Subj Predicted", overall.get("total_subjects_predicted", "—"))
    o_cols[3].metric("Subj GT", overall.get("total_subjects_gt", "—"))


# ---------------------------------------------------------------------------
# Agenda-item alignment view
# ---------------------------------------------------------------------------

def _render_agenda_alignment(
    alignment: List[Dict[str, Any]],
    topic_color_fn,
    view_mode: str,
    filter_iou: float,
) -> None:
    st.markdown("### 🗂️ Agenda Item Alignment")

    if not alignment:
        st.info("No agenda-item alignment data found for this document.")
        return

    # Filter
    shown = [
        a for a in alignment
        if (a.get("iou") is None or a.get("iou", 1.0) >= filter_iou)
    ]

    if not shown:
        st.warning(f"No agenda items with IoU ≥ {filter_iou:.2f}.")
        return

    col_gt, col_pred = st.columns(2)
    col_gt.markdown("**📋 Ground Truth**")
    col_pred.markdown("**🤖 Predicted**")

    for entry in shown:
        pred = entry.get("pred", {})
        gt = entry.get("matched_gt", {})
        iou = entry.get("iou")
        border = _iou_color(iou)
        iou_label = f"IoU: {iou:.3f}" if iou is not None else "Unmatched"

        # Build subject previews for GT and Pred
        def _subj_list(subjects: List[Dict]) -> str:
            items = []
            for s in subjects:
                theme = s.get("theme", "")
                if theme:
                    items.append(f"• {theme}")
            return "\n".join(items) if items else ""

        gt_subjects_preview = _subj_list(gt.get("subjects", []))
        pred_subjects_count = len(pred.get("subjects", []))

        gt_topics: List[str] = []
        for s in gt.get("subjects", []):
            gt_topics.extend(s.get("topics", []))
        gt_topics = list(dict.fromkeys(gt_topics))  # dedup, keep order

        iou_badge = (
            f'<div style="text-align:right;font-size:0.78rem;color:{border};'
            f'font-weight:bold;margin-top:6px;">{_html.escape(iou_label)}</div>'
        )

        # GT card
        with col_gt:
            st.html(
                _card(
                    title=gt.get("item_title", f"Item #{entry.get('matched_gt_idx', '?')}"),
                    text=gt_subjects_preview or gt.get("text_preview", ""),
                    topics=gt_topics,
                    theme="",
                    border_color=border,
                    topic_color_fn=topic_color_fn,
                    extra_html=iou_badge,
                )
            )

        # Pred card
        with col_pred:
            pred_title = pred.get("item_title", f"Item #{entry.get('pred_idx', '?')}")
            pred_text = pred.get("text_preview", "")
            st.html(
                _card(
                    title=pred_title,
                    text=pred_text,
                    topics=[],
                    theme="",
                    border_color=border,
                    topic_color_fn=topic_color_fn,
                    extra_html=(
                        f'<div style="font-size:0.78rem;color:#777;margin-top:4px;">'
                        f'{pred_subjects_count} subject(s) extracted</div>'
                        + iou_badge
                    ),
                )
            )


# ---------------------------------------------------------------------------
# Subject alignment view
# ---------------------------------------------------------------------------

def _render_subject_alignment(
    alignment: List[Dict[str, Any]],
    topic_color_fn,
    filter_iou: float,
) -> None:
    st.markdown("### 🔍 Subject Alignment")

    if not alignment:
        st.info("No subject alignment data found for this document.")
        return

    shown = [
        a for a in alignment
        if (a.get("iou") is None or a.get("iou", 1.0) >= filter_iou)
    ]

    if not shown:
        st.warning(f"No subjects with IoU ≥ {filter_iou:.2f}.")
        return

    col_gt, col_pred = st.columns(2)
    col_gt.markdown("**📋 Ground Truth Subjects**")
    col_pred.markdown("**🤖 Predicted Subjects**")

    for i, entry in enumerate(shown):
        pred = entry.get("pred", {})
        gt = entry.get("matched_gt", {})
        iou = entry.get("iou")
        theme_match: Optional[bool] = entry.get("theme_match")
        border = _iou_color(iou)
        iou_label = f"IoU: {iou:.3f}" if iou is not None else "Unmatched"

        topic_overlap = entry.get("topic_overlap", [])
        topic_pred_only = entry.get("topic_pred_only", [])
        topic_gt_only = entry.get("topic_gt_only", [])

        iou_badge = (
            f'<div style="text-align:right;font-size:0.78rem;color:{border};'
            f'font-weight:bold;margin-top:6px;">{_html.escape(iou_label)}</div>'
        )
        theme_match_html = ""
        if theme_match is not None:
            icon = "✅" if theme_match else "❌"
            theme_match_html = (
                f'<div style="font-size:0.78rem;margin-top:4px;">'
                f'{icon} Theme match</div>'
            )

        # GT side
        with col_gt:
            st.html(
                _card(
                    title=gt.get("subject_id", f"Subject {i+1}"),
                    text=gt.get("text_preview", gt.get("text", "")[:300]),
                    topics=gt.get("topics", []),
                    theme=gt.get("theme", ""),
                    border_color=border,
                    topic_color_fn=topic_color_fn,
                    extra_html=iou_badge,
                )
            )

        # Pred side
        with col_pred:
            pred_theme = pred.get("theme", "")
            pred_text = pred.get("text_preview", pred.get("text", "")[:300])
            pred_topics = pred.get("topics", [])

            # Topic diff annotations
            topic_diff = ""
            if topic_overlap or topic_pred_only or topic_gt_only:
                parts = []
                if topic_overlap:
                    parts.append(
                        "✅ " + ", ".join(
                            f'<span style="color:#27ae60">{_html.escape(t)}</span>'
                            for t in topic_overlap
                        )
                    )
                if topic_pred_only:
                    parts.append(
                        "➕ " + ", ".join(
                            f'<span style="color:#e67e22">{_html.escape(t)}</span>'
                            for t in topic_pred_only
                        )
                    )
                if topic_gt_only:
                    parts.append(
                        "➖ " + ", ".join(
                            f'<span style="color:#e74c3c">{_html.escape(t)}</span>'
                            for t in topic_gt_only
                        )
                    )
                topic_diff = (
                    '<div style="font-size:0.78rem;margin-top:6px;line-height:1.7;">'
                    + " &nbsp;|&nbsp; ".join(parts)
                    + "</div>"
                )

            if not pred_theme and not pred_text:
                st.html(
                    _card(
                        title="Missing Prediction",
                        text="",
                        topics=[],
                        theme="",
                        border_color=border,
                        topic_color_fn=topic_color_fn,
                        missing=True,
                    )
                )
            else:
                st.html(
                    _card(
                        title=f"Predicted Subject {i+1}",
                        text=pred_text,
                        topics=pred_topics,
                        theme=pred_theme,
                        border_color=border,
                        topic_color_fn=topic_color_fn,
                        extra_html=topic_diff + theme_match_html + iou_badge,
                    )
                )


# ---------------------------------------------------------------------------
# IoU distribution chart for a document
# ---------------------------------------------------------------------------

def _render_iou_chart(alignment: List[Dict[str, Any]], title: str) -> None:
    ious = [e.get("iou") for e in alignment if e.get("iou") is not None]
    if not ious:
        return
    df = pd.DataFrame({"IoU": ious})
    fig = px.histogram(
        df, x="IoU", nbins=15,
        title=title, range_x=[0, 1],
        color_discrete_sequence=["#4C72B0"],
    )
    fig.update_layout(height=260, margin=dict(t=36, b=20, l=20, r=10))
    st.plotly_chart(fig, use_container_width=True)


# ---------------------------------------------------------------------------
# Predictions-only view  (clean reading layout)
# ---------------------------------------------------------------------------

def _render_predictions_view(
    agenda_alignment: List[Dict[str, Any]],
    subjects_alignment: List[Dict[str, Any]],
    topic_color_fn,
) -> None:
    """Clean reading view: agenda item titles + all LLM-predicted subjects."""
    st.markdown("### 📋 HSeg Predictions")

    if not agenda_alignment:
        st.info("No prediction data available for this document.")
        return

    # Build lookup from subjects_alignment to retrieve actual predicted themes/topics
    pred_subject_lookup = {}
    for sa in subjects_alignment:
        spred = sa.get("pred")
        if spred and "start" in spred and "end" in spred:
            pred_subject_lookup[(spred["start"], spred["end"])] = {
                "theme": spred.get("theme", ""),
                "topics": spred.get("topics", [])
            }

    for entry in agenda_alignment:
        pred = entry.get("pred", {})
        item_title = pred.get("item_title", f"Item #{entry.get('pred_idx', '?')}")
        subjects: List[Dict[str, Any]] = pred.get("subjects", [])

        # ── Agenda item header ────────────────────────────────────────────
        st.markdown(
            f'<div style="background:linear-gradient(90deg,#0f3460,#16213e);'
            f'color:white;padding:10px 16px;border-radius:10px;'
            f'font-weight:700;font-size:1rem;margin:14px 0 6px 0;">'
            f'📁 {_html.escape(item_title)}</div>',
            unsafe_allow_html=True,
        )

        if not subjects:
            st.caption("_No subjects predicted for this item._")
            continue

        # ── Subject cards (full width, stacked) ──────────────────────────
        for idx, subj in enumerate(subjects):
            s_start = subj.get("start")
            s_end = subj.get("end")

            theme = subj.get("theme", "")
            topics: List[str] = subj.get("topics", [])

            # Check if subjects_alignment has the populated theme/topics
            if (s_start, s_end) in pred_subject_lookup:
                lookup_data = pred_subject_lookup[(s_start, s_end)]
                if lookup_data["theme"]:
                    theme = lookup_data["theme"]
                if lookup_data["topics"]:
                    topics = lookup_data["topics"]

            text: str = subj.get("text", subj.get("text_preview", ""))

            badges = "".join(_topic_badge(t, topic_color_fn(t)) for t in topics)
            badges_html = (
                f'<div style="margin:6px 0 4px 0;">{badges}</div>' if badges else ""
            )
            theme_html = (
                f'<div style="background:#edf6ff;border-left:4px solid #2196f3;'
                f'padding:7px 10px;margin:6px 0;font-style:italic;'
                f'font-size:0.88rem;color:#111;">'
                f'📝 {_html.escape(theme)}</div>'
                if theme else ""
            )
            text_html = (
                f'<div style="font-size:0.85rem;color:#222;line-height:1.6;'
                f'white-space:pre-wrap;margin-top:6px;">'
                f'{_html.escape(text)}</div>'
                if text else ""
            )

            st.html(
                f'<div style="border:1px solid #dde4ef;border-radius:8px;'
                f'padding:12px 16px;margin:4px 0 4px 24px;background:#fafcff;">'
                f'<div style="font-size:0.78rem;color:#888;margin-bottom:4px;">'
                f'Subject {idx + 1}</div>'
                f'{badges_html}{theme_html}{text_html}'
                f'</div>'
            )


# ---------------------------------------------------------------------------
# Main entry point
# ---------------------------------------------------------------------------

def render_document_view(doc_data: Dict[str, Any], topic_color_fn) -> None:
    """Render the full document visualiser."""

    doc_id = doc_data.get("minute_id", "Unknown")
    st.markdown(f"### 📄 {doc_id}")

    if doc_data.get("status") != "success":
        st.error(f"Document status: {doc_data.get('status', 'unknown')}")
        return

    evaluation = doc_data.get("evaluation", {})
    aligned = doc_data.get("aligned_comparison", {})
    agenda_alignment = aligned.get("agenda_items_alignment", [])
    subject_alignment = aligned.get("subjects_alignment", [])

    # Metrics summary
    _render_doc_metrics(evaluation)
    st.divider()

    # Controls
    ctrl_cols = st.columns(3)
    with ctrl_cols[0]:
        view_section = st.radio(
            "View",
            ["📋 Predictions", "Agenda Items", "Subjects", "Both"],
            key=f"view_section_{doc_id}",
            horizontal=True,
        )
    with ctrl_cols[1]:
        filter_iou = st.slider(
            "Min IoU filter",
            0.0, 1.0, 0.0, 0.05,
            key=f"iou_filter_{doc_id}",
        )
    with ctrl_cols[2]:
        show_iou_chart = st.checkbox("Show IoU distribution", value=True, key=f"show_iou_{doc_id}")

    if show_iou_chart:
        c1, c2 = st.columns(2)
        with c1:
            _render_iou_chart(agenda_alignment, "Agenda Item IoU Distribution")
        with c2:
            _render_iou_chart(subject_alignment, "Subject IoU Distribution")

    st.divider()

    if view_section == "📋 Predictions":
        _render_predictions_view(agenda_alignment, subject_alignment, topic_color_fn)

    if view_section in ("Agenda Items", "Both"):
        _render_agenda_alignment(agenda_alignment, topic_color_fn, "sequential", filter_iou)
        st.divider()

    if view_section in ("Subjects", "Both"):
        _render_subject_alignment(subject_alignment, topic_color_fn, filter_iou)