src/streamlit_app.py

import os
import requests
import streamlit as st
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# ───────────────────────── CONFIG ─────────────────────────
MURIL_REPO = "yashsha7/mindlens-muril"   # your pushed model
OPENROUTER_KEY = os.environ.get("OPENAI_API_KEY", "")

st.set_page_config(page_title="MindLens", page_icon="🧠", layout="wide")

NEG = {"empty","hopeless","worthless","exhausted","lonely","burden","cry",
       "dark","helpless","numb","tired","alone","lost","overwhelmed","depressed"}
POS = {"happy","grateful","excited","love","great","good","joy","proud",
       "calm","peaceful","hopeful","motivated","energized","positive"}


# ───────────────────────── CACHED LOADERS ─────────────────────────
@st.cache_resource(show_spinner="Loading MuRIL model (first load only)...")
def load_muril():
    tok = AutoTokenizer.from_pretrained(MURIL_REPO)
    model = AutoModelForSequenceClassification.from_pretrained(MURIL_REPO)
    model.eval()
    return tok, model


@st.cache_resource(show_spinner="Loading spaCy NER...")
def load_spacy():
    import spacy
    try:
        return spacy.load("en_core_web_sm")
    except OSError:
        from spacy.cli import download
        download("en_core_web_sm")
        return spacy.load("en_core_web_sm")


@st.cache_resource(show_spinner="Building FAISS knowledge base (first load only)...")
def load_vectorstore():
    from langchain_community.vectorstores import FAISS
    from langchain_community.embeddings import HuggingFaceEmbeddings
    from langchain.text_splitter import RecursiveCharacterTextSplitter
    from langchain.schema import Document
    from datasets import load_dataset

    raw = load_dataset("hugginglearners/reddit-depression-cleaned", split="train")
    df = raw.to_pandas()[["clean_text"]].dropna().head(500)
    posts = df["clean_text"].astype(str).tolist()

    split_docs = RecursiveCharacterTextSplitter(
        chunk_size=300, chunk_overlap=30
    ).split_documents([Document(page_content=p) for p in posts])

    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
    return FAISS.from_documents(split_docs, embeddings)


tokenizer_bert, muril_model = load_muril()


# ───────────────────────── CORE FUNCTIONS ─────────────────────────
def predict(text):
    enc = tokenizer_bert(str(text), return_tensors="pt", truncation=True,
                          max_length=128, padding=True)
    with torch.no_grad():
        probs = torch.softmax(muril_model(**enc).logits, dim=1)[0]
    pred = torch.argmax(probs).item()
    label = "Depression" if pred == 1 else "Normal"
    conf = round(probs[pred].item() * 100, 1)
    dep_prob = round(probs[1].item() * 100, 1)
    return label, conf, dep_prob


def sentiment(text):
    t = set(text.lower().split())
    ns, ps = len(t & NEG), len(t & POS)
    if ns > ps:
        return f"Negative({ns})"
    if ps > ns:
        return f"Positive({ps})"
    return "Neutral"


def detect_lang(text):
    return "Hindi" if any('\u0900' <= c <= '\u097F' for c in text) else "English"


def gpt_chat(messages, max_tokens=400):
    if not OPENROUTER_KEY:
        return "(No API key set — add OPENAI_API_KEY in Space secrets.)"
    r = requests.post(
        "https://openrouter.ai/api/v1/chat/completions",
        headers={
            "Authorization": f"Bearer {OPENROUTER_KEY}",
            "Content-Type": "application/json",
            "HTTP-Referer": "https://huggingface.co",
            "X-Title": "MindLens",
        },
        json={"model": "openai/gpt-4o-mini", "messages": messages,
              "temperature": 0.7, "max_tokens": max_tokens},
    )
    try:
        return r.json()["choices"][0]["message"]["content"]
    except Exception:
        return f"(Error calling model: {r.text[:200]})"


# ───────────────────────── SIDEBAR NAV ─────────────────────────
st.sidebar.title("🧠 MindLens")
st.sidebar.caption("Mental Health NLP · v1.0")
page = st.sidebar.radio("Navigate", [
    "Depression Detector", "Bulk Analyze", "AI Wellness Chatbot",
    "Agentic AI Pipeline", "RAG Knowledge Chat"
])
st.sidebar.markdown("---")
st.sidebar.caption("Models: MuRIL · FLAN-T5 · BiLSTM")

# ───────────────────────── PAGE: DETECTOR ─────────────────────────
if page == "Depression Detector":
    st.title("Depression Detector")
    st.caption("Paste a social media post in English or Hindi — analyzed instantly.")
    text = st.text_area("Post text", height=100)
    if st.button("Analyze", type="primary") and text.strip():
        label, conf, dep_prob = predict(text)
        lang = detect_lang(text)
        sent = sentiment(text)
        col1, col2 = st.columns(2)
        col1.metric("Result", label)
        col2.metric("Language", lang)
        st.progress(conf / 100, text=f"Confidence Score: {conf}%")
        st.write(f"**Sentiment:** {sent}")
        if label == "Depression":
            st.warning(
                "I'm really sorry you're feeling this way. It might help to talk to "
                "someone you trust, or try journaling your feelings — it can be a "
                "powerful way to gain some clarity."
            )

# ───────────────────────── PAGE: BULK ANALYZE ─────────────────────────
elif page == "Bulk Analyze":
    st.title("Bulk Analyze")
    st.caption("Paste multiple posts (one per line) — analyze all at once. English & Hindi supported.")
    bulk_text = st.text_area("Posts (one per line)", height=200)
    if st.button("Analyze All", type="primary") and bulk_text.strip():
        lines = [l for l in bulk_text.split("\n") if l.strip()]
        rows = []
        for line in lines:
            label, conf, _ = predict(line)
            rows.append({
                "Post": line[:60] + ("..." if len(line) > 60 else ""),
                "Label": label,
                "Confidence": f"{conf}%",
                "Sentiment": sentiment(line),
                "Language": detect_lang(line),
            })
        import pandas as pd
        res_df = pd.DataFrame(rows)
        st.dataframe(res_df, use_container_width=True)

        dep_count = sum(1 for r in rows if r["Label"] == "Depression")
        norm_count = len(rows) - dep_count
        avg_conf = sum(float(r["Confidence"].strip("%")) for r in rows) / len(rows)

        c1, c2, c3, c4 = st.columns(4)
        c1.metric("Posts Analyzed", len(rows))
        c2.metric("Depression", dep_count)
        c3.metric("Normal", norm_count)
        c4.metric("Avg Confidence", f"{avg_conf:.1f}%")

        st.bar_chart(res_df["Label"].value_counts())
        st.download_button("Download CSV", res_df.to_csv(index=False), "results.csv")

# ───────────────────────── PAGE: CHATBOT ─────────────────────────
elif page == "AI Wellness Chatbot":
    st.title("Wellness Chatbot")
    st.caption("Bilingual (EN + HI) · Depression detection on every message · GPT-4o-mini responses")

    if "history" not in st.session_state:
        st.session_state.history = []

    for msg in st.session_state.history:
        with st.chat_message(msg["role"]):
            st.write(msg["content"])

    user_msg = st.chat_input("How are you feeling today?")
    if user_msg:
        label, conf, _ = predict(user_msg)
        sent = sentiment(user_msg)
        lang = detect_lang(user_msg)
        det = f"[{label} | {conf}% | {sent} | {lang}]"

        st.session_state.history.append({"role": "user", "content": user_msg})
        with st.chat_message("user"):
            st.write(user_msg)

        sys_prompt = (
            "You are a compassionate mental health support assistant. "
            "If depression detected: validate feelings, suggest one coping strategy. "
            "If normal: respond warmly. Keep it short and human. Respond in user's language."
        )
        msgs = [{"role": "system", "content": sys_prompt},
                {"role": "system", "content": det}] + st.session_state.history
        reply = gpt_chat(msgs)

        st.session_state.history.append({"role": "assistant", "content": reply})
        with st.chat_message("assistant"):
            st.caption(det)
            st.write(reply)

    if st.button("Clear conversation"):
        st.session_state.history = []
        st.rerun()

# ───────────────────────── PAGE: AGENTIC PIPELINE ─────────────────────────
elif page == "Agentic AI Pipeline":
    st.title("Agentic AI Pipeline")
    st.caption("5-step autonomous analysis: Detect → NER → Sentiment → Similar Posts → Clinical Report")

    post = st.text_area("Social media post to analyze", height=100)
    if st.button("Run Pipeline", type="primary") and post.strip():
        nlp_spacy = load_spacy()
        vectorstore = load_vectorstore()

        with st.status("Running pipeline...", expanded=True) as status:
            label, conf, dep_prob = predict(post)
            st.write(f"**1. Detection:** {label} · Confidence: {conf}% · Dep prob: {dep_prob}%")

            doc = nlp_spacy(post[:400])
            ents = [(e.text, e.label_) for e in doc.ents]
            st.write(f"**2. NER:** {ents if ents else 'No named entities found.'}")

            sent = sentiment(post)
            st.write(f"**3. Sentiment:** {sent}")

            similar = vectorstore.similarity_search(post, k=3)
            st.write("**4. Similar posts retrieved:**")
            for s in similar:
                st.caption("- " + s.page_content[:100] + "...")

            report = gpt_chat([{
                "role": "user",
                "content": (
                    f"Clinical NLP report for:\nPost: {post[:200]}\n"
                    f"Condition: {label} ({conf}%)\nSentiment: {sent}\nEntities: {ents}\n"
                    "Write: 1.Summary 2.Risk indicators 3.Linguistic patterns 4.Recommendations"
                )
            }])
            st.write("**5. Clinical Report:**")
            st.write(report)
            status.update(label="Pipeline complete", state="complete")

        c1, c2, c3 = st.columns(3)
        c1.metric("Condition", label[:4])
        c2.metric("Confidence", f"{conf}%")
        c3.metric("Entities", len(ents))

# ───────────────────────── PAGE: RAG CHAT ─────────────────────────
elif page == "RAG Knowledge Chat":
    st.title("RAG Knowledge Chat")
    st.caption("Ask questions about mental health patterns in the Reddit dataset. FAISS + GPT-4o-mini.")

    query = st.text_input("Ask about the dataset...")
    if st.button("Ask") and query.strip():
        vectorstore = load_vectorstore()
        docs = vectorstore.similarity_search(query, k=4)
        context = "\n\n".join(d.page_content for d in docs)

        answer = gpt_chat([{
            "role": "user",
            "content": f"Context from Reddit posts:\n{context}\n\nQuestion: {query}\nAnswer based on the context above:"
        }])
        st.write(answer)
        with st.expander(f"Retrieved {len(docs)} chunks · FAISS · all-MiniLM-L6-v2 embeddings"):
            for i, d in enumerate(docs, 1):
                st.caption(f"Source {i}: {d.page_content[:150]}...")