import streamlit as st
import torch
import numpy as np
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# ── Config ────────────────────────────────────────────────────────────────────
HF_MODEL_ID = "ishaan1402/cbt-thought-pattern-classifier"
THRESHOLD   = 0.15
MAX_LENGTH  = 256

PATTERN_CLASSES = [
    "Catastrophizing",
    "Discounting the positive",
    "Labeling and mislabeling",
    "Mental filtering",
    "Jumping to conclusions: mind reading",
    "Jumping to conclusions: Fortune-telling",
    "Overgeneralization",
    "Personalization",
    "Black-and-white or polarized thinking / All or nothing thinking",
    "Should statements",
    "None"
]

PATTERN_DESCRIPTIONS = {
    "Catastrophizing":              "Giving greater weight to the worst possible outcome.",
    "Discounting the positive":     "Rejecting positive experiences by insisting they don't count.",
    "Labeling and mislabeling":     "Attributing actions to character rather than situation.",
    "Mental filtering":             "Dwelling only on the negative details of a situation.",
    "Jumping to conclusions: mind reading": "Inferring negative thoughts from someone's behaviour.",
    "Jumping to conclusions: Fortune-telling": "Predicting negative outcomes of events.",
    "Overgeneralization":           "Making faulty generalisations from insufficient evidence.",
    "Personalization":              "Assigning disproportionate personal blame to oneself.",
    "Black-and-white or polarized thinking / All or nothing thinking":
                                    "Viewing things as either all good or all bad with no middle ground.",
    "Should statements":            "Demanding particular behaviours regardless of realistic circumstances.",
    "None":                         "No unhelpful thought pattern detected.",
}

# ── Model loading (cached so it only runs once per session) ───────────────────
@st.cache_resource
def load_model():
    tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_ID)
    model     = AutoModelForSequenceClassification.from_pretrained(HF_MODEL_ID)
    model.eval()
    return tokenizer, model

# ── Inference ─────────────────────────────────────────────────────────────────
def classify(thought: str, persona: str, tokenizer, model) -> dict:
    input_text = f"Persona: {persona} | Thought: {thought}" if persona.strip() \
                 else f"Persona:  | Thought: {thought}"

    inputs = tokenizer(
        input_text,
        return_tensors="pt",
        max_length=MAX_LENGTH,
        truncation=True,
        padding="max_length"
    )

    with torch.no_grad():
        logits = model(**inputs).logits

    probs      = torch.softmax(logits, dim=-1).squeeze().numpy()
    top_idx    = int(np.argmax(probs))
    top_label  = PATTERN_CLASSES[top_idx]
    confidence = float(probs[top_idx])

    is_unhelpful = (top_label != "None") and (confidence >= THRESHOLD)

    return {
        "is_unhelpful":      is_unhelpful,
        "predicted_pattern": top_label if is_unhelpful else "None",
        "confidence":        confidence,
        "distribution":      {PATTERN_CLASSES[i]: float(p) for i, p in enumerate(probs)},
    }

# ── UI ────────────────────────────────────────────────────────────────────────
st.set_page_config(page_title="CBT Thought Classifier", page_icon="🧠", layout="centered")

st.title("CBT Thought Pattern Classifier")
st.caption(
    "Enter a thought below. The model will identify which unhelpful cognitive pattern "
    "it exhibits, if any, based on Cognitive Behavioral Therapy (CBT) research."
)

st.divider()

persona = st.text_input(
    "Persona (optional)",
    placeholder="e.g. I am a college student. I love playing the guitar.",
    help="Adding context about the person improves classification accuracy."
)

thought = st.text_area(
    "Thought",
    placeholder="e.g. I failed this test, I'm going to fail my entire degree.",
    height=120
)

classify_btn = st.button("Classify", type="primary", use_container_width=True)

if classify_btn:
    if not thought.strip():
        st.warning("Please enter a thought to classify.")
    else:
        with st.spinner("Loading model...") if "tokenizer" not in st.session_state else st.spinner("Classifying..."):
            tokenizer, model = load_model()
            result = classify(thought, persona, tokenizer, model)

        st.divider()

        pattern     = result["predicted_pattern"]
        confidence  = result["confidence"]
        unhelpful   = result["is_unhelpful"]

        # ── Result banner ──
        if unhelpful:
            st.error(f"**Unhelpful pattern detected:** {pattern}")
        else:
            st.success("**No unhelpful pattern detected** — this thought looks okay.")

        # ── Pattern description ──
        if pattern in PATTERN_DESCRIPTIONS:
            st.info(f"📖 **{pattern}:** {PATTERN_DESCRIPTIONS[pattern]}")

        # ── Confidence ──
        st.metric("Model confidence", f"{confidence:.1%}")

        # ── Full distribution ──
        with st.expander("See full probability distribution"):
            dist = result["distribution"]
            sorted_dist = sorted(dist.items(), key=lambda x: x[1], reverse=True)
            for label, prob in sorted_dist:
                bar_pct = int(prob * 100)
                # Highlight the predicted class
                label_display = f"**{label}**" if label == pattern else label
                col1, col2 = st.columns([3, 1])
                with col1:
                    st.markdown(label_display)
                    st.progress(bar_pct)
                with col2:
                    st.markdown(f"`{prob:.3f}`")

st.divider()
st.caption("Model trained on the [PATTERNREFRAME dataset](https://github.com/facebookresearch/ParlAI/tree/main/projects/reframe_thoughts) · Built with RoBERTa-large")