import joblib
import numpy as np
import pandas as pd
import streamlit as st

# =========================
# CONFIG / CONSTANTS
# =========================

# These thresholds decide the UI label.
# You can change them anytime.
TH_LOW = 0.20   # below this => Low Risk
TH_HIGH = 0.40  # above this => High Risk
# between TH_LOW and TH_HIGH => Moderate Risk

MODEL_PATH = "best_model_v2_calibrated.joblib"

FEATURE_ORDER = [
    "Age",
    "Gender",
    "BMI",
    "BP_SYS",
    "BP_DIA",
    "Phys_Activity_Days",
    "DPQ_Score",
    "Smoking_Indicator",
    "Alcohol_Feature",
    "Diabetes_Indicator",
    "Cycle",
]

# =========================
# HELPERS
# =========================

def label_from_probability(prob: float) -> str:
    """3-level risk label based on thresholds."""
    if prob < TH_LOW:
        return "Low Risk"
    elif prob < TH_HIGH:
        return "Moderate Risk"
    else:
        return "High Risk"

def get_texts(lang: str):
    """Simple bilingual text dictionary."""
    if lang == "বাংলা":
        return {
            "title": "SleepGuardAI – Sleep Risk Screening (NHANES-based)",
            "subtitle": "ফর্ম পূরণ করুন → ঝুঁকির স্কোর দেখুন। (এটি মেডিকেল ডায়াগনসিস নয়)",
            "predict_btn": "Predict / ফলাফল দেখুন",
            "result": "ফলাফল",
            "prob": "ঝুঁকি সম্ভাবনা",
            "thresholds": "থ্রেশহোল্ড",
            "warning": "⚠️ এটি একটি স্ক্রিনিং টুল, মেডিকেল ডায়াগনসিস নয়। লক্ষণ থাকলে ডাক্তার দেখান।",
            "low": "কম ঝুঁকি",
            "mod": "মাঝারি ঝুঁকি",
            "high": "উচ্চ ঝুঁকি",
        }
    else:
        return {
            "title": "SleepGuardAI – Sleep Risk Screening (NHANES-based)",
            "subtitle": "Fill the form → get a risk score. (This is NOT a medical diagnosis)",
            "predict_btn": "Predict / ফলাফল দেখুন",
            "result": "Result",
            "prob": "Risk probability",
            "thresholds": "Thresholds",
            "warning": "⚠️ This is a screening tool, not a medical diagnosis. If you have symptoms, consult a clinician.",
            "low": "Low Risk",
            "mod": "Moderate Risk",
            "high": "High Risk",
        }

@st.cache_resource
def load_model():
    return joblib.load(MODEL_PATH)

def build_feature_row(
    age, gender, bmi, bp_sys, bp_dia, phys_days, dpq, smoking, alcohol, diabetes, cycle
) -> pd.DataFrame:
    """Build one-row DataFrame in the exact feature order used by training."""
    row = {
        "Age": float(age),
        "Gender": str(gender),  # keep as string (your pipeline encodes it)
        "BMI": float(bmi),
        "BP_SYS": float(bp_sys),
        "BP_DIA": float(bp_dia),
        "Phys_Activity_Days": float(phys_days),
        "DPQ_Score": float(dpq),
        "Smoking_Indicator": int(smoking),
        "Alcohol_Feature": float(alcohol),
        "Diabetes_Indicator": int(diabetes),
        "Cycle": str(cycle),
    }
    df = pd.DataFrame([row])[FEATURE_ORDER]
    return df

# =========================
# UI
# =========================

st.set_page_config(page_title="SleepGuardAI", layout="wide")

lang = st.radio("Language / ভাষা", ["English", "বাংলা"], horizontal=True)
T = get_texts(lang)

st.title(T["title"])
st.caption(T["subtitle"])

# Load model
try:
    bundle = load_model()
except FileNotFoundError:
    st.error(
        f"Model file not found: '{MODEL_PATH}'. Put it in the same folder as app.py."
    )
    st.stop()

# We support two possible save formats:
# 1) bundle is a dict with {"model":..., "calibrator":..., ...}
# 2) bundle is directly a sklearn Pipeline/Calibrated model with predict_proba
model = bundle.get("model") if isinstance(bundle, dict) and "model" in bundle else bundle
calibrator = bundle.get("calibrator") if isinstance(bundle, dict) and "calibrator" in bundle else None

colA, colB = st.columns(2)

with colA:
    age = st.slider("Age / বয়স", 10, 90, 30)
    gender = st.selectbox("Gender", ["Male", "Female"])

    bmi = st.slider("BMI", 10.0, 50.0, 25.0, 0.1)
    bp_sys = st.slider("Systolic BP", 80, 220, 120)
    bp_dia = st.slider("Diastolic BP", 40, 140, 80)

    phys_days = st.slider("Physical Activity Days/Week", 0, 7, 3)

with colB:
    st.markdown("### Optional (improves accuracy if known)")
    dpq = st.slider("DPQ Depression Score (0–27)", 0, 27, 5)
    smoking = st.selectbox("Smoking (1=Yes, 2=No)", [2, 1], index=0)
    diabetes = st.selectbox("Diabetes (1=Yes, 2=No)", [2, 1], index=0)
    alcohol = st.slider("Alcohol feature (proxy)", 0.0, 30.0, 2.0, 0.5)
    cycle = st.selectbox("NHANES Cycle", ["G", "H", "I", "J"])

st.divider()

if st.button(T["predict_btn"], use_container_width=True):
    X = build_feature_row(age, gender, bmi, bp_sys, bp_dia, phys_days, dpq, smoking, alcohol, diabetes, cycle)

    # Get probability
    # If you saved a separate calibrator, run model -> probs -> calibrator
    # Otherwise assume model already outputs calibrated proba.
    try:
        raw_proba = model.predict_proba(X)[:, 1]
    except Exception as e:
        st.error(f"Model could not compute predict_proba. Error: {e}")
        st.stop()

    if calibrator is not None:
        # calibrator usually expects 2D array
        prob = float(calibrator.predict(raw_proba.reshape(-1, 1))[0])
    else:
        prob = float(raw_proba[0])

    label = label_from_probability(prob)

    # Display with bilingual mapping
    if label == "Low Risk":
        pretty_label = T["low"]
        st.success(f"{T['result']}: {pretty_label}")
    elif label == "Moderate Risk":
        pretty_label = T["mod"]
        st.warning(f"{T['result']}: {pretty_label}")
    else:
        pretty_label = T["high"]
        st.error(f"{T['result']}: {pretty_label}")

    st.write(f"**{T['prob']}:** {prob:.3f}")
    st.write(f"**{T['thresholds']}:** Low<{TH_LOW:.2f} | Moderate {TH_LOW:.2f}–{TH_HIGH:.2f} | High≥{TH_HIGH:.2f}")
    st.caption(T["warning"])

    # Debug view if you want to show professor:
    with st.expander("Show input feature row (for debugging/paper)"):
        st.dataframe(X)