"""
Samsung Health Chart Intent — Streamlit Test App
=================================================
Uses the two-stage pipeline from test.py:
  1. OOD detector  — rejects non-health queries
  2. DistilBERT    — chart vs no_chart

Usage:
    streamlit run app.py
"""

import json
import pandas as pd
from pathlib import Path

import streamlit as st

# Import the full pipeline from test.py (same folder)
from test import build_pipeline, full_predict, OOD_PERCENTILE, CONF_THRESHOLD

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

MODEL_DIR     = "./chart_intent_model"
TRAINING_DATA = "./samsung_health_intent.csv"

# ─────────────────────────────────────────────────────────
# Load pipeline — cached so it only runs once
# ─────────────────────────────────────────────────────────

@st.cache_resource
def load_pipeline(model_dir: str, training_data: str):
    return build_pipeline()

def load_history(model_dir: str):
    path = Path(model_dir) / "history.json"
    if path.exists():
        with open(path) as f:
            return json.load(f)
    return None

# ─────────────────────────────────────────────────────────
# Result rendering helpers
# ─────────────────────────────────────────────────────────

def render_single_result(result: dict, threshold: float):
    intent = result["intent"]
    st.divider()

    if intent == "out_of_domain":
        st.error(
            f"🚫 **OUT OF DOMAIN** — not recognised as a health query\n\n"
            f"OOD distance: `{result['distance']:.4f}` (exceeds threshold)"
        )
        return

    is_chart  = intent == "chart"
    uncertain = not result["confident"] or result.get("confidence", 1.0) < threshold
    icon      = "📊" if is_chart else "💬"
    label_txt = "CHART" if is_chart else "NO CHART"

    if uncertain:
        st.warning(
            f"⚠️ **Uncertain** — confidence `{result['confidence']:.1%}` "
            f"is below threshold (`{threshold:.0%}`)"
        )

    col_res, col_gauge = st.columns(2)

    with col_res:
        st.markdown(f"### {icon} Prediction: `{label_txt}`")
        st.markdown(f"**Confidence:** {result['confidence']:.1%}")
        st.markdown(f"**OOD distance:** {result['distance']:.4f}")
        st.markdown(f"**Query:** _{result['text']}_")

    with col_gauge:
        st.markdown("**Probability breakdown**")
        st.markdown(f"📊 Chart &nbsp;&nbsp;&nbsp; `{result['prob_chart']:.1%}`")
        st.progress(result["prob_chart"])
        st.markdown(f"💬 No chart &nbsp; `{result['prob_no_chart']:.1%}`")
        st.progress(result["prob_no_chart"])


def result_to_row(result: dict, threshold: float, gt: int = None) -> dict:
    """Convert a full_predict result dict into a table row."""
    intent = result["intent"]

    if intent == "out_of_domain":
        pred_str  = "🚫 out_of_domain"
        conf_str  = "—"
        uncertain = ""
        p_chart   = "—"
    else:
        is_chart  = intent == "chart"
        pred_str  = "📊 chart" if is_chart else "💬 no_chart"
        conf_str  = f"{result['confidence']:.1%}"
        uncertain = "⚠️" if result.get("confidence", 1.0) < threshold else ""
        p_chart   = f"{result['prob_chart']:.3f}"

    row = {
        "Query":      result["text"],
        "Prediction": pred_str,
        "Confidence": conf_str,
        "P(chart)":   p_chart,
        "OOD dist":   f"{result['distance']:.4f}",
        "Uncertain":  uncertain,
    }

    if gt is not None:
        if intent == "out_of_domain":
            correct = False
        else:
            correct = result["label"] == gt
        row["Ground truth"] = "📊 chart" if gt == 1 else "💬 no_chart"
        row["Correct"]      = "✅" if correct else "❌"

    return row

# ─────────────────────────────────────────────────────────
# Page setup
# ─────────────────────────────────────────────────────────

st.set_page_config(
    page_title="Samsung Health — Chart Intent Tester",
    page_icon="📊",
    layout="wide",
)

st.title("📊 Samsung Health Chart Intent Classifier")
st.caption("Fine-tuned DistilBERT + OOD detector · Test your health chatbot queries")

# ─────────────────────────────────────────────────────────
# Sidebar
# ─────────────────────────────────────────────────────────

with st.sidebar:
    st.header("⚙️ Settings")

    model_dir     = st.text_input("Model directory", value=MODEL_DIR)
    training_data = st.text_input("Training CSV",    value=TRAINING_DATA)

    st.divider()

    if not Path(model_dir).exists():
        st.error(f"Model not found at `{model_dir}`\n\nRun `finetune.py` first.")
        st.stop()
    if not Path(training_data).exists():
        st.error(f"Training data not found at `{training_data}`\n\nNeeded to fit the OOD detector.")
        st.stop()

    with st.spinner("Loading model + OOD detector..."):
        try:
            ood, clf = load_pipeline(model_dir, training_data)
        except Exception as e:
            st.error(f"Failed to load pipeline:\n\n{e}")
            st.stop()

    st.success("Pipeline ready")

    import torch
    if torch.cuda.is_available():
        device_name = "CUDA GPU"
    elif torch.backends.mps.is_available():
        device_name = "Apple Silicon MPS"
    else:
        device_name = "CPU"
    st.markdown(f"**Device:** {device_name}")

    config_path = Path(model_dir) / "config.json"
    if config_path.exists():
        with open(config_path) as f:
            cfg = json.load(f)
        st.markdown(f"**Architecture:** {cfg.get('model_type', 'unknown')}")

    st.divider()

    st.subheader("Thresholds")
    threshold = st.slider(
        "Confidence threshold",
        min_value=0.50, max_value=0.99,
        value=float(CONF_THRESHOLD), step=0.01,
        help="Classifier predictions below this are flagged ⚠️ uncertain",
    )
    ood_pct = st.slider(
        "OOD percentile",
        min_value=80, max_value=99,
        value=int(OOD_PERCENTILE), step=1,
        help="Higher = more permissive OOD gate. Change takes effect on restart.",
    )
    if ood_pct != OOD_PERCENTILE:
        st.info("OOD percentile change takes effect on next app restart.")

    st.divider()

    history = load_history(model_dir)
    if history:
        st.subheader("Training history")
        best = max(history, key=lambda x: x["val"]["f1_chart"])
        st.metric("Best val F1",       f"{best['val']['f1_chart']:.4f}", f"epoch {best['epoch']}")
        st.metric("Best val accuracy", f"{best['val']['accuracy']:.4f}")
        chart_df = pd.DataFrame({
            "val accuracy":   [h["val"]["accuracy"]   for h in history],
            "train accuracy": [h["train"]["accuracy"] for h in history],
        })
        st.line_chart(chart_df, height=150)

# ─────────────────────────────────────────────────────────
# Tabs
# ─────────────────────────────────────────────────────────

tab1, tab2 = st.tabs(["🔍 Single query", "📋 Batch test"])

# ═══════════════════════════════════════════
# TAB 1: Single query
# ═══════════════════════════════════════════
with tab1:
    st.subheader("Test a single query")

    query = st.text_input(
        "Type a health query",
        placeholder="e.g. show me my heart rate chart for this week",
        key="single_query",
    )

    col_btn, _ = st.columns([1, 5])
    with col_btn:
        run = st.button("Classify", type="primary", use_container_width=True)

    if run and query.strip():
        result = full_predict(query.strip(), ood, clf)
        render_single_result(result, threshold)
    elif run:
        st.warning("Please enter a query first.")

    st.divider()
    st.subheader("Quick examples")

    examples_chart = [
        "plot my heart rate over the last 7 days",
        "show me a chart of my weekly step count",
        "graph my sleep patterns for the past month",
        "visualize my calorie burn trend",
        "draw a chart comparing my weekday vs weekend steps",
    ]
    examples_no_chart = [
        "what is my resting heart rate",
        "how many steps did I take today",
        "did I hit my step goal yesterday",
        "what was my sleep score last night",
        "what is my current weight",
    ]
    examples_ood = [
        "your name",
        "tell me a joke",
        "what time is it",
        "hello",
        "what is the weather today",
    ]

    col_ex1, col_ex2, col_ex3 = st.columns(3)

    with col_ex1:
        st.markdown("**📊 Chart queries**")
        for ex in examples_chart:
            if st.button(ex, key=f"ex1_{ex}", use_container_width=True):
                r  = full_predict(ex, ood, clf)
                ok = r["intent"] == "chart"
                st.markdown(f"{'✅' if ok else '❌'} `chart={r.get('prob_chart', 0):.2f}` dist=`{r['distance']:.3f}`")

    with col_ex2:
        st.markdown("**💬 No-chart queries**")
        for ex in examples_no_chart:
            if st.button(ex, key=f"ex2_{ex}", use_container_width=True):
                r  = full_predict(ex, ood, clf)
                ok = r["intent"] == "no_chart"
                st.markdown(f"{'✅' if ok else '❌'} `no_chart={r.get('prob_no_chart', 0):.2f}` dist=`{r['distance']:.3f}`")

    with col_ex3:
        st.markdown("**🚫 OOD queries**")
        for ex in examples_ood:
            if st.button(ex, key=f"ex3_{ex}", use_container_width=True):
                r  = full_predict(ex, ood, clf)
                ok = r["intent"] == "out_of_domain"
                st.markdown(f"{'✅' if ok else '❌'} `{r['intent']}` dist=`{r['distance']:.3f}`")

# ═══════════════════════════════════════════
# TAB 2: Batch test
# ═══════════════════════════════════════════
with tab2:
    st.subheader("Test multiple queries at once")
    st.caption(
        "One query per line. Optionally add `,0` or `,1` for ground truth. "
        "OOD queries are flagged automatically."
    )

    default_batch = """\
plot my heart rate over the last 7 days,1
what is my resting heart rate,0
show me a chart of my weekly step count,1
how many calories did I burn this week,0
graph my sleep patterns for the past month,1
did I hit my step goal yesterday,0
visualize my calorie burn trend,1
what was my sleep score last night,0
draw a recovery score chart for the last 2 weeks,1
what is my current body fat percentage,0
your name
tell me a joke
what is the weather today"""

    batch_input = st.text_area(
        "Queries",
        value=default_batch,
        height=280,
        placeholder="One query per line. Add ,0 or ,1 for ground truth.",
    )

    if st.button("Run batch", type="primary"):
        lines = [l.strip() for l in batch_input.strip().splitlines() if l.strip()]

        rows        = []
        correct     = 0
        has_labels  = False
        n_ood       = 0
        n_uncertain = 0

        for line in lines:
            parts = line.rsplit(",", 1)
            text  = parts[0].strip()
            gt    = None
            if len(parts) == 2 and parts[1].strip() in ("0", "1"):
                gt = int(parts[1].strip())
                has_labels = True

            result = full_predict(text, ood, clf)
            row    = result_to_row(result, threshold, gt)
            rows.append(row)

            if result["intent"] == "out_of_domain":
                n_ood += 1
            elif result.get("confidence", 1.0) < threshold:
                n_uncertain += 1

            if gt is not None and row.get("Correct") == "✅":
                correct += 1

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

        st.divider()
        total   = len(rows)
        n_chart = sum(1 for r in rows if "📊 chart" in r["Prediction"])

        col_m1, col_m2, col_m3, col_m4, col_m5 = st.columns(5)
        col_m1.metric("Total",        total)
        col_m2.metric("Chart",        f"{n_chart} ({n_chart/total:.0%})")
        col_m3.metric("OOD rejected", f"{n_ood} ({n_ood/total:.0%})")
        col_m4.metric("Uncertain",    f"{n_uncertain} ({n_uncertain/total:.0%})")
        if has_labels:
            labelled = sum(1 for r in rows if "Correct" in r)
            col_m5.metric("Accuracy", f"{correct/labelled:.1%}", f"{correct}/{labelled}")

        csv = results_df.to_csv(index=False)
        st.download_button(
            "⬇️ Download results CSV",
            data=csv,
            file_name="batch_results.csv",
            mime="text/csv",
        )