app.py

import os
import re
from datetime import datetime

import streamlit as st
import pandas as pd
import numpy as np
import altair as alt

from rules_engine import compute_bmi, evaluate_risks, rules_risk_assessment
from data_store import init_db, upsert_entry, get_user_history, to_csv_bytes

# ----------------------
# App Config & Theming
# ----------------------
st.set_page_config(
    page_title="NeuraVia – Symptom Tracker & Risk Dashboard",
    page_icon="🧑‍⚕️",
    layout="wide",
)

PRIMARY_CONDITIONS = [
    "Type 2 Diabetes Risk",
    "Hypertension Risk",
    "Depression/Mood Concern",
    "Migraine Risk",
    "Sleep Apnea Risk",
    "Anemia Risk",
]

# Map human label → DB column
def to_db_key(label: str) -> str:
    return "risk_" + re.sub(r"[^0-9A-Za-z]+", "_", label).strip("_")

@st.cache_resource(show_spinner=False)
def get_hf_pipeline():
    try:
        return load_zero_shot_classifier()
    except Exception:
        return None

# ----------------------
# Sidebar – Navigation
# ----------------------
with st.sidebar:
    st.title("🧑‍⚕️NeuraVia – Symptom Tracker")
    st.caption("Detect • Connect • Personalize")
    page = st.radio("Go to", ["Home", "Start", "Input", "Dashboard", "History"], index=0)
    st.divider()
    demo_mode = st.toggle("Use Demo Data", value=False, help="Prefill a demo input")
    st.divider()
    st.info(
        "Educational use only. Not a medical device. Consult a clinician for advice.",
        icon="⚠️",
    )

def explain_rules(rules_explain: dict) -> str:
    if not rules_explain:
        return "No rule-based explanation available."
    lines = []
    for key, value in rules_explain.items():
        lines.append(f"- **{key}**: {value}")
    return "\n".join(lines)

# ----------------------
# Global init
# ----------------------
init_db()

st.markdown(
    """
<style>
.card {background:#fff;border-radius:16px;padding:16px;box-shadow:0 8px 24px rgba(0,0,0,0.07)}
.metric{font-size:22px;font-weight:800}
.section-title{font-size:26px;font-weight:800;margin-bottom:8px}
.hero {text-align:center; padding:50px 20px;}
.hero h1 {font-size:48px; font-weight:900; margin-bottom:20px;}
.hero h3 {font-size:22px; font-weight:400; color:#444; margin-bottom:30px;}
.hero .tagline {font-size:18px; font-weight:500; color:#666;}
</style>
""",
    unsafe_allow_html=True,
)

# ----------------------
# Landing Page
# ----------------------
if page == "Home":
    st.markdown(
        """
        <div class="hero">
            <h1>🧑‍⚕️ NeuraVia</h1>
            <h3>Symptom Tracker & Risk Dashboard</h3>
            <p class="tagline">Capture data → interpretable risk flags → trend charts.<br>
            Built for <b>NeuraViaHacks</b>.</p>
        </div>
        """,
        unsafe_allow_html=True,
    )

    st.markdown("<div class='card'>", unsafe_allow_html=True)
    st.markdown("### 🌟 Why NeuraVia?")
    st.write(
        "- Track your health symptoms, vitals, and trends.\n"
        "- Get explainable risk scores using rules + AI.\n"
        "- Store securely and export anytime.\n"
        "- Beginner-friendly, interpretable, and practical."
    )
    st.markdown("</div>", unsafe_allow_html=True)

    st.info("👉 Use the sidebar to navigate to **Start, Input, Dashboard, or History**.")

# ----------------------
# Start Page
# ----------------------
if page == "Start":
    c1, c2 = st.columns([1.2, 1])
    with c1:
        st.markdown("### What this does")
        st.write(
            "- Collect demographics, vitals, and symptoms.\n"
            "- Score conditions via explainable rules + optional HF zero-shot.\n"
            "- Store locally in SQLite and visualize trends.\n"
            "- Export CSV for sharing."
        )
        st.markdown("### How to demo")
        st.write(
            "1) Open **Input** and submit (or toggle Demo).\n"
            "2) Open **Dashboard** for cards & charts.\n"
            "3) Use **History** to review and export CSV."
        )
    with c2:
        st.markdown("<div class='card'>", unsafe_allow_html=True)
        st.markdown("**Beginner-friendly • Interpretable • Real-world relevance**")
        st.caption("Detect · Connect · Personalize")
        st.markdown("</div>", unsafe_allow_html=True)

# ----------------------
# Input Page
# ----------------------
if page == "Input":
    st.markdown("<div class='section-title'>Patient Data Input</div>", unsafe_allow_html=True)

    if demo_mode:
        default = dict(
            user_id="demo_user",
            age=42,
            sex="Male",
            height_cm=175,
            weight_kg=92,
            sbp=142,
            dbp=92,
            hr=88,
            spo2=96,
            glucose=118,
            symptoms=["Frequent urination", "Headache", "Daytime sleepiness", "Loud snoring"],
            free_text="Often thirsty and tired; loud snoring; morning headaches.",
        )
    else:
        default = dict(
            user_id="",
            age=30,
            sex="Female",
            height_cm=165,
            weight_kg=65,
            sbp=120,
            dbp=80,
            hr=72,
            spo2=98,
            glucose=95,
            symptoms=[],
            free_text="",
        )

    with st.form("input_form"):
        st.markdown("**Profile**")
        c1, c2, c3 = st.columns(3)
        user_id = c1.text_input("User ID (unique)", value=default["user_id"], help="Used to group your timeline")
        age = c2.number_input("Age", 1, 120, int(default["age"]))
        sex = c3.selectbox(
            "Sex",
            ["Female", "Male", "Other", "Prefer not to say"],
            index=1 if default["sex"] == "Male" else 0,
        )

        st.markdown("**Vitals**")
        c1, c2, c3, c4 = st.columns(4)
        height_cm = c1.number_input("Height (cm)", 80, 230, int(default["height_cm"]))
        weight_kg = c2.number_input("Weight (kg)", 20, 250, int(default["weight_kg"]))
        sbp = c3.number_input("Systolic BP (SBP)", 70, 240, int(default["sbp"]))
        dbp = c4.number_input("Diastolic BP (DBP)", 40, 140, int(default["dbp"]))

        c1, c2, c3 = st.columns(3)
        hr = c1.number_input("Heart Rate (bpm)", 30, 200, int(default["hr"]))
        spo2 = c2.number_input("SpO₂ (%)", 70, 100, int(default["spo2"]))
        glucose = c3.number_input("Glucose (mg/dL)", 50, 400, int(default["glucose"]))

        st.markdown("**Symptoms**")
        SYMPTOMS_LIST = [
            "Frequent urination", "Excessive thirst", "Unintended weight loss", "Blurred vision",
            "Headache", "Nausea", "Sensitivity to light", "Sleep problems", "Daytime sleepiness",
            "Loud snoring", "Pauses in breathing during sleep", "Shortness of breath",
            "Fatigue", "Dizziness", "Pale skin", "Cold hands/feet",
            "Low mood", "Loss of interest", "Anxiety", "Irritability",
        ]
        symptoms = st.multiselect("Select common symptoms (optional)", SYMPTOMS_LIST, default=default["symptoms"])
        free_text = st.text_area("Describe your symptoms (optional)", value=default["free_text"], height=110)

        submitted = st.form_submit_button("Compute Risk", use_container_width=True)

    if submitted:
        if not user_id:
            st.error("Please enter a User ID.")
            st.stop()

        bmi = compute_bmi(height_cm, weight_kg)
        rules_scores, rules_explain = rules_risk_assessment(
            age=age,
            sex=sex,
            sbp=sbp,
            dbp=dbp,
            hr=hr,
            spo2=spo2,
            glucose=glucose,
            bmi=bmi,
            symptoms=symptoms,
        )

        # Optional HF zero-shot scoring
        hf_scores = {}
        pipe = get_hf_pipeline()
        if pipe is not None:
            text = "; ".join(symptoms) + ("; " + free_text if free_text else "")
            hf_scores = zero_shot_score(pipe, text=text, labels=PRIMARY_CONDITIONS)

        # Combine scores
        combined = {}
        for cond in PRIMARY_CONDITIONS:
            r = float(rules_scores.get(cond, 0.0))
            h = hf_scores.get(cond, np.nan)
            combined[cond] = r if (isinstance(h, float) and np.isnan(h)) else (0.65 * r + 0.35 * float(h))

        entry = {
            "user_id": user_id,
            "ts": datetime.utcnow().isoformat(),
            "age": age,
            "sex": sex,
            "height_cm": height_cm,
            "weight_kg": weight_kg,
            "bmi": round(bmi, 2),
            "sbp": sbp,
            "dbp": dbp,
            "hr": hr,
            "spo2": spo2,
            "glucose": glucose,
            "symptoms": ", ".join(symptoms),
            "free_text": free_text,
            **{ to_db_key(cond): round(score, 3) for cond, score in combined.items() },
        }

        upsert_entry(entry)

        st.success("Risk computed and saved. Open the **Dashboard** tab.")
        with st.expander("See rule explanations"):
            st.markdown(explain_rules(rules_explain))

# ----------------------
# Dashboard Page
# ----------------------
if page == "Dashboard":
    st.markdown("<div class='section-title'>Risk Dashboard</div>", unsafe_allow_html=True)
    uid = st.text_input("Enter User ID to view", value="demo_user" if demo_mode else "")
    if not uid:
        st.info("Enter your User ID to see assessments.")
        st.stop()

    df = get_user_history(uid)
    if df.empty:
        st.warning("No records found. Submit an entry in **Input**.")
        st.stop()

    df["ts"] = pd.to_datetime(df["ts"], errors="coerce")
    df = df.sort_values("ts")
    latest = df.iloc[-1]

    # Top risk cards
    st.markdown("#### Latest Risk Flags")
    cols = st.columns(3)

    def badge(v: float) -> str:
        if v >= 0.75:
            return "🔴 High"
        if v >= 0.5:
            return "🟠 Moderate"
        if v >= 0.25:
            return "🟡 Mild"
        return "🟢 Low"

    metrics = []
    for label in PRIMARY_CONDITIONS:
        colname = to_db_key(label)
        if colname in df.columns:
            try:
                val = float(latest[colname])
            except Exception:
                val = 0.0
            pretty = label.replace(" Risk", "")
            metrics.append((pretty, val))

    for i, (name, val) in enumerate(metrics):
        with cols[i % 3]:
            st.markdown("<div class='card'>", unsafe_allow_html=True)
            st.markdown(f"**{name}**")
            st.markdown(f"<div class='metric'>{val:.2f}</div>", unsafe_allow_html=True)
            st.caption(badge(val))
            st.markdown("</div>", unsafe_allow_html=True)

    # Trend charts
    st.markdown("#### Trends Over Time")
    vitals = ["bmi", "sbp", "dbp", "hr", "spo2", "glucose"]
    display_to_col = {**{v: v for v in vitals}, **{f"risk: {lbl}": to_db_key(lbl) for lbl in PRIMARY_CONDITIONS}}

    pick = st.multiselect(
        "Choose metrics to plot",
        options=vitals + [f"risk: {lbl}" for lbl in PRIMARY_CONDITIONS],
        default=["bmi", "sbp", "dbp"],
    )

    for disp in pick:
        colname = display_to_col.get(disp, disp)
        if colname not in df.columns:
            continue
        chart = (
            alt.Chart(df).mark_line(point=True).encode(
                x=alt.X("ts:T", title="Time"),
                y=alt.Y(f"{colname}:Q", title=disp),
                tooltip=["ts:T", colname],
            ).properties(height=220)
        )
        st.altair_chart(chart, use_container_width=True)

# ----------------------
# History Page
# ----------------------
if page == "History":
    st.markdown("<div class='section-title'>History & Export</div>", unsafe_allow_html=True)
    uid = st.text_input("User ID", value="demo_user" if demo_mode else "")
    if not uid:
        st.info("Enter a User ID to load history.")
        st.stop()

    df = get_user_history(uid)
    if df.empty:
        st.warning("No records yet. Submit data in **Input**.")
    else:
        st.dataframe(df, use_container_width=True)
        csv_bytes = to_csv_bytes(df)
        st.download_button(
            "Download CSV",
            data=csv_bytes,
            file_name=f"{uid}_history.csv",
            mime="text/csv",
            use_container_width=True,
        )

st.markdown("---")
st.caption("NeuraViaHacks submission by Qaisar A • Built with Streamlit • Not medical advice.")