streamlit_onestopqa.py

"""
Streamlit explorer for OneStopQA dataset.
"""
from __future__ import annotations

import html
import json
import re
from pathlib import Path
from typing import Iterable, List, Sequence, Tuple

import streamlit as st

Span = Tuple[int, int]


@st.cache_data(show_spinner=False)
def load_onestop_qa() -> List[dict]:
    """Load the OneStopQA JSON payload once per session."""
    data_path = Path(__file__).parent / "onestop_qa.json"
    with data_path.open("r", encoding="utf-8") as f:
        payload = json.load(f)
    return payload.get("data", [])


def _as_pairs(spans: Sequence[Sequence[int]] | Sequence[int] | None) -> List[Span]:
    """Expand span definitions into (start, end) pairs.

    Input can be either:
    - list[list[int]] (usual case: one item per question)
    - list[int] for a single question
    Each item is [start, end] or [start1, end1, start2, end2].
    Invalid or empty spans are ignored.
    """
    pairs: List[Span] = []
    if spans is None:
        return pairs

    # If spans is already a flat list of ints (single question), wrap it.
    if spans and isinstance(spans[0], int):  # type: ignore[index]
        spans = [spans]  # type: ignore[assignment]

    for span in spans:  # type: ignore[iteration-over-annotation]
        if not isinstance(span, Sequence) or isinstance(span, (str, bytes)):
            continue
        if len(span) == 2:
            pairs.append((int(span[0]), int(span[1])))
        elif len(span) == 4:
            pairs.append((int(span[0]), int(span[1])))
            pairs.append((int(span[2]), int(span[3])))
    return pairs


def _normalize_spans(spans: Iterable[Span], text_length: int, label: str, color: str):
    """Clamp spans to the text length and attach styling metadata."""
    normalized = []
    for start, end in spans:
        lo, hi = sorted((max(0, start), min(text_length, end)))
        if lo >= hi:
            continue
        normalized.append({
            "start": lo,
            "end": hi,
            "label": label,
            "color": color,
        })
    return sorted(normalized, key=lambda x: (x["start"], -(x["end"] - x["start"])))


def _get_span_colors_for_duplicates(a_spans_list: List[Sequence[Sequence[int]] | None], text: str) -> dict[Tuple[int, int], str]:
    """Identify duplicate critical spans across questions and return color mapping.
    
    For spans that appear in 2+ questions (exactly or with overlaps):
    - Overlapping regions: #FF8C00 (orange)
    - Non-overlapping regions of overlapped spans: #FFA500 (lighter orange)
    
    For spans in only 1 question: #20B2AA (teal)
    """
    # Convert all spans to normalized character spans with question tracking
    all_spans_by_question: list[list[dict]] = []
    
    for q_idx, a_spans in enumerate(a_spans_list):
        if a_spans is None:
            all_spans_by_question.append([])
            continue
        word_spans = _as_pairs(a_spans)
        char_spans = _word_spans_to_char_spans(word_spans, text)
        normalized = _normalize_spans(char_spans, len(text), "critical", "#d4f5a3")
        all_spans_by_question.append(normalized)
    
    # Build a character-level map of which spans cover each position
    char_to_spans: dict[int, list[Tuple[int, int, int]]] = {}  # char_pos -> [(q_idx, start, end), ...]
    
    for q_idx, spans in enumerate(all_spans_by_question):
        for span in spans:
            start, end = span["start"], span["end"]
            for char_pos in range(start, end):
                if char_pos not in char_to_spans:
                    char_to_spans[char_pos] = []
                char_to_spans[char_pos].append((q_idx, start, end))
    
    # Assign colors based on character coverage
    span_colors: dict[Tuple[int, int], str] = {}
    
    for q_idx, spans in enumerate(all_spans_by_question):
        for span in spans:
            start, end = span["start"], span["end"]
            key = (start, end)
            
            # Check if any character in this span is covered by spans from other questions
            is_overlapped = False
            for char_pos in range(start, end):
                other_questions = set(
                    s[0] for s in char_to_spans.get(char_pos, []) if s[0] != q_idx
                )
                if other_questions:
                    is_overlapped = True
                    break
            
            if is_overlapped:
                # This span overlaps with spans from other questions
                span_colors[key] = "#E65100"  # Dark orange
            else:
                # Check if this span is unique to this question or duplicated exactly
                question_count = len(set(
                    s[0] for s in char_to_spans.get(start, []) 
                    if (s[1], s[2]) == (start, end)
                ))
                
                if question_count >= 2:
                    span_colors[key] = "#E65100"  # Dark orange for exact duplicates
                else:
                    span_colors[key] = "#00838F"  # Dark teal for unique
    
    # Now handle partial overlaps by checking span boundaries
    for q_idx, spans in enumerate(all_spans_by_question):
        for span in spans:
            start, end = span["start"], span["end"]
            key = (start, end)
            
            # Check for other spans that partially overlap but don't cover the whole span
            for other_q_idx, other_spans in enumerate(all_spans_by_question):
                if other_q_idx == q_idx:
                    continue
                for other_span in other_spans:
                    other_start, other_end = other_span["start"], other_span["end"]
                    # Check for any overlap
                    if not (end <= other_start or start >= other_end):
                        # There's an overlap - mark this span as orange
                        span_colors[key] = "#E65100"
                        break
    
    return span_colors


def _word_spans_to_char_spans(word_spans: List[Span], text: str) -> List[Span]:
    """Convert word-index spans (zero-based) to character spans."""
    tokens = list(re.finditer(r"\S+", text))
    if not tokens:
        return []

    char_spans: List[Span] = []
    last_idx = len(tokens) - 1

    for start_word, end_word in word_spans:
        lo_idx = min(start_word, end_word)
        hi_idx = max(start_word, end_word)
        lo_idx = max(0, lo_idx)
        hi_idx = min(last_idx, hi_idx)
        if lo_idx > hi_idx:
            continue
        char_start = tokens[lo_idx].start()
        char_end = tokens[hi_idx].end()
        char_spans.append((char_start, char_end))

    return char_spans


def highlight_text(text: str, a_spans: Sequence[Sequence[int]] | None, d_spans: Sequence[Sequence[int]] | None, span_colors: dict[Tuple[int, int], str] | None = None) -> str:
    """Return HTML with critical and distractor spans highlighted.
    
    Args:
        text: The text to highlight
        a_spans: Critical spans (word indices)
        d_spans: Distractor spans (word indices)
        span_colors: Optional dict mapping (char_start, char_end) to color hex codes
    """
    crit_words = _as_pairs(a_spans)
    dist_words = _as_pairs(d_spans)
    crit_char_spans = _word_spans_to_char_spans(crit_words, text)
    dist_char_spans = _word_spans_to_char_spans(dist_words, text)
    
    # Use custom colors if provided, otherwise use defaults
    if span_colors:
        crit = _normalize_spans(crit_char_spans, len(text), "critical", "#2E7D32")
        dist = _normalize_spans(dist_char_spans, len(text), "distractor", "#D32F2F")
        # Override colors from the span_colors dict
        for span in crit + dist:
            key = (span["start"], span["end"])
            if key in span_colors:
                span["color"] = span_colors[key]
        ordered = sorted(crit + dist, key=lambda x: (x["start"], x["label"]))
    else:
        crit = _normalize_spans(crit_char_spans, len(text), "critical", "#2E7D32")
        dist = _normalize_spans(dist_char_spans, len(text), "distractor", "#D32F2F")
        ordered = sorted(crit + dist, key=lambda x: (x["start"], x["label"]))

    parts: List[str] = []
    cursor = 0
    for span in ordered:
        start, end = span["start"], span["end"]
        # If spans overlap, trim the new span so we do not double paint.
        if start < cursor:
            start = cursor
        if start >= end:
            continue
        if cursor < start:
            parts.append(html.escape(text[cursor:start]))
        span_text = html.escape(text[start:end])
        parts.append(
            f"<span style='color:{span['color']}; font-weight:bold;' title='{span['label']}'>"
            f"{span_text}</span>"
        )
        cursor = end
    if cursor < len(text):
        parts.append(html.escape(text[cursor:]))
    return "".join(parts)


def main():
    st.set_page_config(page_title="OneStopQA Explorer", layout="wide")
    st.title("OneStopQA Explorer")
    st.caption("View passages, questions, and highlight critical/distractor spans.")

    dataset = sorted(load_onestop_qa(), key=lambda x: x.get("article_id", ""))
    if not dataset:
        st.error("Could not load OneStopQA dataset.")
        return

    article_titles = [f"{item['title']} ({item['article_id']})" for item in dataset]
    selected_article_label = st.selectbox("Article", article_titles, index=0)
    selected_article_idx = article_titles.index(selected_article_label)
    article = dataset[selected_article_idx]

    # Reset navigation state when article changes
    if "last_article_idx" not in st.session_state:
        st.session_state.last_article_idx = selected_article_idx
    if selected_article_idx != st.session_state.last_article_idx:
        st.session_state.para_idx = 0
        st.session_state.question_idx = 0
        st.session_state.last_article_idx = selected_article_idx

    paragraphs = article.get("paragraphs", [])
    if not paragraphs:
        st.warning("No paragraphs found for this article.")
        return

    # Initialize session state for navigation
    if "para_idx" not in st.session_state:
        st.session_state.para_idx = 0
    if "question_idx" not in st.session_state:
        st.session_state.question_idx = 0
    if "view_mode" not in st.session_state:
        st.session_state.view_mode = "By Question"

    # View mode selector
    view_mode = st.radio(
        "View Mode",
        ["By Question", "All Critical Spans"],
        horizontal=True,
        key="view_mode_radio"
    )
    st.session_state.view_mode = view_mode

    # Paragraph navigation buttons
    para_col1, para_col3, para_col2 = st.columns([0.5, 0.5, 3])
    with para_col1:
        if st.button("← Prev", key="prev_para_btn", shortcut="A"):
            st.session_state.para_idx = max(0, st.session_state.para_idx - 1)
            st.session_state.question_idx = 0
    with para_col2:
        if st.button("Next →", key="next_para_btn", shortcut="D"):
            st.session_state.para_idx = min(len(paragraphs) - 1, st.session_state.para_idx + 1)
            st.session_state.question_idx = 0
    with para_col3:
        st.write(f"**Paragraph {st.session_state.para_idx + 1} / {len(paragraphs)}**")

    para_idx = st.session_state.para_idx
    para = paragraphs[para_idx]
    # st.caption(f"Paragraph {para.get('paragraph_id', para_idx + 1)}")

    qas = para.get("qas", [])
    if not qas:
        st.warning("No questions for this paragraph.")
        return

    # Question navigation buttons (only show in "By Question" mode)
    if st.session_state.view_mode == "By Question":
        q_col1, q_col3, q_col2 = st.columns([0.5, 0.5, 3])
        with q_col1:
            if st.button("← Prev", key="prev_q_btn", shortcut="W"):
                st.session_state.question_idx = max(0, st.session_state.question_idx - 1)
        with q_col2:
            if st.button("Next →", key="next_q_btn", shortcut="S"):
                st.session_state.question_idx = min(len(qas) - 1, st.session_state.question_idx + 1)
        with q_col3:
            q_ind = qas[st.session_state.question_idx].get('q_ind', st.session_state.question_idx)
            st.write(f"**Question {st.session_state.question_idx + 1} / {len(qas)}**")

    question_idx = st.session_state.question_idx
    # st.caption(f"Question {qas[question_idx].get('q_ind', question_idx + 1)}")

    selected_question = qas[question_idx] if qas else {}

    # st.write("**Title**: " + article.get("title", ""))

    for level in ["Adv", "Int", "Ele"]:
        context_block = para.get(level, {})
        context_text = context_block.get("context", "")
        a_list = context_block.get("a_spans") or []
        d_list = context_block.get("d_spans") or []
        
        level_label = {
            "Adv": "Advanced",
            "Int": "Intermediate",
            "Ele": "Elementary",
        }.get(level, level)
        st.subheader(f"{level_label}")
        
        if context_text:
            if st.session_state.view_mode == "By Question":
                # Show only the current question's critical and distractor spans
                a_spans = a_list[question_idx] if question_idx < len(a_list) else []
                d_spans = d_list[question_idx] if question_idx < len(d_list) else []
                html_block = highlight_text(context_text, a_spans, d_spans)
            else:  # All Critical Spans mode
                # Collect all critical spans from all questions, no distractors
                all_a_spans = []
                for q_idx in range(len(qas)):
                    if q_idx < len(a_list) and a_list[q_idx]:
                        all_a_spans.append(a_list[q_idx])
                    else:
                        all_a_spans.append(None)  # Preserve question order with None for missing spans
                # Get color mapping for duplicates vs unique spans
                span_colors = _get_span_colors_for_duplicates(all_a_spans, context_text)
                html_block = highlight_text(context_text, all_a_spans, None, span_colors)
            st.markdown(html_block, unsafe_allow_html=True)
        else:
            st.info("No context available for this level.")

    # st.subhe  ader("")
    if st.session_state.view_mode == "By Question":
        st.write("**Question**: " + selected_question.get("question", ""))
        if selected_question.get("answers"):
            st.markdown("**Answer choices**")
            for idx, ans in enumerate(selected_question["answers"]):
                st.write(f"{idx + 1}. {ans}")
    else:  # All Critical Spans mode
        st.subheader("Questions")
        for q_idx, qa in enumerate(qas):
            q_ind = qa.get('q_ind', q_idx)
            st.write(f"**Question {q_ind+1}**: {qa.get('question', '')}")
            if qa.get("answers"):
                for idx, ans in enumerate(qa["answers"]):
                    st.write(f"  {idx + 1}. {ans}")

    
    # Legend
    if st.session_state.view_mode == "By Question":
        st.markdown(
            "<span style='color:#2E7D32; font-weight:bold;'>critical span</span> "
            "<span style='color:#D32F2F; font-weight:bold;'>distractor span</span>",
            unsafe_allow_html=True,
        )
    else:
        st.markdown(
            "<span style='color:#E65100; font-weight:bold;'>duplicate critical span</span> "
            "<span style='color:#00838F; font-weight:bold;'>unique critical span</span>",
            unsafe_allow_html=True,
        )
    # Keyboard shortcuts
    st.markdown(
        "<style>body { font-size: 12px; } p { margin: 0; }</style>",
        unsafe_allow_html=True,
    )
if __name__ == "__main__":
    main()