# sjt_answers_viewer.py
# Gradio viewer for case_study_answers.json
# - Shows one question at a time
# - Dropdown 1: filter by **Name**
# - Dropdown 2: filter by **Selected Trait** (HEXACO slice)
# - Underlines & highlights the actually selected option + trait name
# - Always orders options consistently (HEXACO)
# - Top "Summary" bar shows proportion of selections by trait (under current Name filter)

import json
from pathlib import Path
from typing import List, Dict, Any, Optional, Tuple
import random

import gradio as gr
import matplotlib.pyplot as plt

# ---------- Constants ----------

DATA_PATH = Path("case_study_answers.json")

HEXACO_ORDER = [
    "hh",
    "emotionality",
    "extraversion",
    "agreeableness",
    "conscientiousness",
    "openness",
]

TRAIT_LABELS = {
    "hh": "Honesty–Humility",
    "emotionality": "Emotionality",
    "extraversion": "Extraversion",
    "agreeableness": "Agreeableness",
    "conscientiousness": "Conscientiousness",
    "openness": "Openness",
}

# ---------- Icons ----------

ICONS = {
    "header": "📝",
    "question": "❓",
    "options": "✅",
    "summary": "📊",
    "progress": "⏭️",
    "metadata": "🔖",
    "filters": "🎛️",
}

# ---------- Data Loading & Normalization ----------

def load_json(path: Path) -> Any:
    if not path.exists():
        return []
    with path.open("r", encoding="utf-8") as f:
        return json.load(f)

def _safe_get_question_block(item: Dict[str, Any]) -> Tuple[str, Dict[str, str], Optional[str]]:
    """
    Extract (question_text, options_map, selected_trait) from a raw item.
    Heuristics:
      - Selected trait is at top-level key 'option'.
      - Question text/options may be under item['question'] with nested 'corrected_sjt' or 'original_sjt'.
      - Options are expected as keys like '<trait>_option' where trait ∈ HEXACO_ORDER.
    """
    selected = item.get("option")

    q = item.get("question", {}) or {}
    block = q.get("corrected_sjt") or q.get("original_sjt") or {}

    question_text = ""
    options: Dict[str, str] = {}

    if isinstance(block, dict):
        question_text = block.get("question") or q.get("question") or ""
        for trait in HEXACO_ORDER:
            k = f"{trait}_option"
            if k in block:
                options[trait] = str(block[k]).strip()
            elif k in q:
                options[trait] = str(q[k]).strip()
    else:
        # sometimes block is a plain string
        question_text = str(block) if block else str(q.get("question", ""))

    # Fallback: look for options directly on item if missing
    if not options and isinstance(q, dict):
        for trait in HEXACO_ORDER:
            k = f"{trait}_option"
            if k in q:
                options[trait] = str(q[k]).strip()

    return question_text.strip(), options, selected

def flatten_entries(raw: Any) -> List[Dict[str, Any]]:
    """
    Returns a list of entries with keys:
      - name (str)
      - question (str)
      - options (dict[trait->text])
      - selected (trait str)
    """
    out: List[Dict[str, Any]] = []

    def handle_item(obj: Dict[str, Any], default_name: str):
        q_text, opts, sel = _safe_get_question_block(obj)
        # Prefer name from object if present; else inherit from container
        nm = (obj.get("name") or default_name or "Unknown").strip() or "Unknown"
        if q_text and opts and sel:
            out.append({"name": nm, "question": q_text, "options": opts, "selected": sel})

    if isinstance(raw, list):
        for x in raw:
            if isinstance(x, dict):
                handle_item(x, "Unknown")
    elif isinstance(raw, dict):
        # Could be {persona_name: [items]} or {persona_name: {...}} etc.
        for k, v in raw.items():
            default_name = str(k)
            if isinstance(v, list):
                for x in v:
                    if isinstance(x, dict):
                        handle_item(x, default_name)
            elif isinstance(v, dict):
                handle_item(v, default_name)
    return out

DATA_RAW = load_json(DATA_PATH)
DATA: List[Dict[str, Any]] = flatten_entries(DATA_RAW)

# Unique names for dropdown
def all_names(entries: List[Dict[str, Any]]) -> List[str]:
    seen = []
    for e in entries:
        nm = e.get("name", "Unknown") or "Unknown"
        if nm not in seen:
            seen.append(nm)
    return sorted(seen)

NAME_FILTERS = ["All"] + all_names(DATA)
TRAIT_FILTERS = ["All"] + HEXACO_ORDER

# ---------- Filtering & Navigation ----------

def get_filtered_indices(entries: List[Dict[str, Any]], name_filt: str, trait_filt: str) -> List[int]:
    idxs = list(range(len(entries)))
    if name_filt != "All":
        idxs = [i for i in idxs if entries[i].get("name") == name_filt]
    if trait_filt != "All":
        idxs = [i for i in idxs if entries[i].get("selected") == trait_filt]
    return idxs

def clamp_index(i: int, n: int) -> int:
    return 0 if n == 0 else (i % n)

# ---------- Summary ----------

def compute_summary(entries: List[Dict[str, Any]]):
    total = len(entries)
    counts = {t: 0 for t in HEXACO_ORDER}
    for e in entries:
        sel = e.get("selected")
        if sel in counts:
            counts[sel] += 1
    labels = [TRAIT_LABELS[t] for t in HEXACO_ORDER]
    props = [counts[t] / total if total else 0.0 for t in HEXACO_ORDER]
    return labels, props, counts, total

def summary_plot(entries: List[Dict[str, Any]]):
    # Returns Markdown with proportions per trait under the current Name filter
    labels, props, counts, total = compute_summary(entries)
    lines = ["## 📊 Summary (Name filter)", f"**Total:** {total}"]
    for label, p in zip(labels, props):
        lines.append(f"- {label}: {p:.2f}")
    return "\n".join(lines)

# ---------- Rendering ----------

def md_question(entry: Dict[str, Any]) -> str:
    q = entry.get("question", "")
    name = entry.get("name", "—")
    return f"## {ICONS['question']} Question\n**Name:** {name}\n\n{q if q else '—'}"

def md_options(entry: Dict[str, Any]) -> str:
    opts: Dict[str, str] = entry.get("options", {})
    selected = entry.get("selected")
    lines = []
    for i, trait in enumerate(HEXACO_ORDER, start=1):
        if trait not in opts:
            continue
        label = TRAIT_LABELS[trait]
        text = opts[trait]
        if trait == selected:
            # underline + highlight both the label and the text
            line = (
                f"{i}. <u><mark><strong>{label}</strong>:</mark></u> "
                f"<u><mark>{text}</mark></u>"
            )
        else:
            line = f"{i}. **{label}:** {text}"
        lines.append(line)
    body = "\n\n".join(lines) if lines else "—"
    return f"## {ICONS['options']} Options (HEXACO order)\n{body}"

def md_metadata(entry: Dict[str, Any], idx: int, total_in_filter: int) -> str:
    sel = entry.get("selected", "—")
    sel_disp = TRAIT_LABELS.get(sel, sel)
    nm = entry.get("name", "—")
    return (
        f"## {ICONS['metadata']} Metadata\n"
        f"**Name:** {nm}  \n"
        f"**Selected Option (Trait):** {sel_disp}  \n"
        f"**Position in Filter:** {idx + 1} / {total_in_filter}"
    )

def md_progress(idx: int, total: int) -> str:
    return f"## {ICONS['progress']} Progress\n**{idx + 1} / {total}**"

def render(entries: List[Dict[str, Any]], name_filt: str, trait_filt: str, pos: int):
    # For summary, use "name-only" filter to show that persona's distribution
    name_only_indices = [i for i, e in enumerate(entries) if (name_filt == "All" or e.get("name") == name_filt)]
    name_only_slice = [entries[i] for i in name_only_indices]

    # For the main view selection, apply both filters
    indices = get_filtered_indices(entries, name_filt, trait_filt)
    n = len(indices)

    if n == 0:
        return (
            summary_plot(name_only_slice),
            f"## {ICONS['question']} Question\n_No questions for filters **Name={name_filt}**, **Trait={trait_filt}**._",
            f"## {ICONS['options']} Options\n—",
            f"## {ICONS['metadata']} Metadata\n—",
            f"## {ICONS['progress']} Progress\n0 / 0",
            0,  # expose pos back
        )

    pos = clamp_index(pos, n)
    entry = entries[indices[pos]]
    return (
        summary_plot(name_only_slice),
        md_question(entry),
        md_options(entry),
        md_metadata(entry, pos, n),
        md_progress(pos, n),
        pos,  # expose pos back
    )

# ---------- Gradio App ----------

with gr.Blocks(title="SJT Answers Viewer") as demo:
    gr.Markdown("# SJT Answers Viewer")
    gr.Markdown(
        f"{ICONS['filters']} **Filters:** Choose a Name and a HEXACO Selected-Trait slice.\n\n"
        f"{ICONS['summary']} **Summary:** Bar shows the trait-selection proportions under the current **Name** filter.\n\n"
        "Options are consistently ordered by HEXACO. The actual selected option is underlined and highlighted."
    )

    with gr.Row():
        name_dd = gr.Dropdown(choices=NAME_FILTERS, value="All", label="Filter by Name", interactive=True)
        trait_dd = gr.Dropdown(choices=TRAIT_FILTERS, value="All", label="Filter by Selected Trait", interactive=True)
        st_pos = gr.State(0)

    with gr.Row():
        prev_btn = gr.Button("Previous")
        next_btn = gr.Button("Next")
        rand_btn = gr.Button("Random")

    # Outputs
    summary_out = gr.Markdown(label="Selections Summary (Name filter)")
    question_out = gr.Markdown()
    options_out = gr.Markdown()
    metadata_out = gr.Markdown()
    progress_out = gr.Markdown()

    # ----- Callbacks -----
    def on_filters_change(name_filt: str, trait_filt: str):
        return [*render(DATA, name_filt, trait_filt, 0), 0]

    def on_prev(name_filt: str, trait_filt: str, pos: int):
        indices = get_filtered_indices(DATA, name_filt, trait_filt)
        if not indices:
            return [*render(DATA, name_filt, trait_filt, pos), pos]
        pos = clamp_index(pos - 1, len(indices))
        return [*render(DATA, name_filt, trait_filt, pos), pos]

    def on_next(name_filt: str, trait_filt: str, pos: int):
        indices = get_filtered_indices(DATA, name_filt, trait_filt)
        if not indices:
            return [*render(DATA, name_filt, trait_filt, pos), pos]
        pos = clamp_index(pos + 1, len(indices))
        return [*render(DATA, name_filt, trait_filt, pos), pos]

    def on_rand(name_filt: str, trait_filt: str, pos: int):
        indices = get_filtered_indices(DATA, name_filt, trait_filt)
        if not indices:
            return [*render(DATA, name_filt, trait_filt, pos), pos]
        pos = random.randrange(len(indices))
        return [*render(DATA, name_filt, trait_filt, pos), pos]

    name_dd.change(
        on_filters_change,
        inputs=[name_dd, trait_dd],
        outputs=[summary_out, question_out, options_out, metadata_out, progress_out, st_pos, st_pos],
    )
    trait_dd.change(
        on_filters_change,
        inputs=[name_dd, trait_dd],
        outputs=[summary_out, question_out, options_out, metadata_out, progress_out, st_pos, st_pos],
    )

    prev_btn.click(
        on_prev,
        inputs=[name_dd, trait_dd, st_pos],
        outputs=[summary_out, question_out, options_out, metadata_out, progress_out, st_pos, st_pos],
    )
    next_btn.click(
        on_next,
        inputs=[name_dd, trait_dd, st_pos],
        outputs=[summary_out, question_out, options_out, metadata_out, progress_out, st_pos, st_pos],
    )
    rand_btn.click(
        on_rand,
        inputs=[name_dd, trait_dd, st_pos],
        outputs=[summary_out, question_out, options_out, metadata_out, progress_out, st_pos, st_pos],
    )

    # initial load
    demo.load(
        lambda: [*render(DATA, "All", "All", 0), 0],
        inputs=None,
        outputs=[summary_out, question_out, options_out, metadata_out, progress_out, st_pos, st_pos],
    )

if __name__ == "__main__":
    demo.launch()