"""
Heart Disease Risk Predictor — Streamlit chatbot
Deployed on Hugging Face Spaces.

Architecture:
  1. User fills in their health info via the sidebar form.
  2. On submit, the app deterministically computes:
       - 10-year CHD risk + tier (from the trained joblib model)
       - NHANES percentile rankings
       - Top 2-3 counterfactual scenarios
       - PubMed citations for each flagged risk factor
  3. Those results are formatted into a structured context block and handed
     to GPT-5-mini, which explains the results conversationally and answers
     follow-up questions using the same context for the rest of the session.

The LLM never invents predictions — it only narrates and explains the
deterministic results computed in Python.
"""

from dataclasses import dataclass
from pathlib import Path
from typing import Dict, List

import joblib
import openai
import pandas as pd
import streamlit as st
from sklearn.pipeline import Pipeline


# --- Re-declare model classes/constants for joblib unpickling -------------
# These mirror section 2 of heart_risk_modeling.ipynb. They MUST live in
# __main__ for joblib.load() to find TrainedModel.

REQUIRED_FEATURES: List[str] = [
    "age", "male", "cigsPerDay", "sysBP", "totChol", "BMI", "diabetes",
]

MODIFIABLE_TARGETS: Dict[str, Dict] = {
    "cigsPerDay": {"vocab": "smoking",      "target":   0.0, "label": "daily cigarette count"},
    "BMI":        {"vocab": "bmi",          "target":  24.0, "label": "BMI"},
    "sysBP":      {"vocab": "hypertension", "target": 120.0, "label": "systolic blood pressure"},
    "totChol":    {"vocab": "cholesterol",  "target": 180.0, "label": "total cholesterol"},
}

TIER_THRESHOLDS = {"low_max": 0.05, "medium_max": 0.20}


@dataclass
class TrainedModel:
    pipeline: Pipeline
    feature_names: List[str]
    train_size: int
    notes: str = ""

    def _to_frame(self, input_dict):
        return pd.DataFrame([{f: input_dict[f] for f in self.feature_names}])

    def predict_proba(self, input_dict):
        return float(self.pipeline.predict_proba(self._to_frame(input_dict))[0, 1])


def risk_tier(p):
    if p < TIER_THRESHOLDS["low_max"]: return "Low"
    if p < TIER_THRESHOLDS["medium_max"]: return "Medium"
    return "High"


def counterfactual_scenarios(model, user_input, top_n=3):
    baseline = model.predict_proba(user_input)
    out = []
    for feat, meta in MODIFIABLE_TARGETS.items():
        original = float(user_input[feat])
        target = float(meta["target"])
        if original <= target:
            continue
        modified = dict(user_input)
        modified[feat] = target
        new_risk = model.predict_proba(modified)
        out.append({
            "feature": feat, "vocab": meta["vocab"], "label": meta["label"],
            "original_value": original, "new_value": target,
            "baseline_risk": baseline, "new_risk": new_risk,
            "delta": baseline - new_risk,
        })
    out.sort(key=lambda s: s["delta"], reverse=True)
    return out[:top_n]


# --- Cached resource loaders ---------------------------------------------
@st.cache_resource
def load_model():
    return joblib.load("artifacts/risk_model.joblib")


@st.cache_data
def load_reference_data():
    nhanes_summary = pd.read_csv("data/nhanes_reference_summary.csv")
    nhanes_raw = pd.read_csv("data/nhanes_clean.csv")  # available for richer features
    pubmed = pd.read_csv("data/pubmed_clean.csv")
    return nhanes_summary, nhanes_raw, pubmed


def percentile_rank(value, summary_row):
    """Approximate ranking from the 5-point summary."""
    pcts = [10, 25, 50, 75, 90]
    cuts = [summary_row[f"p{p}"] for p in pcts]
    if any(pd.isna(c) for c in cuts):
        return None
    if value <= cuts[0]:
        return f"below the 10th percentile"
    for i in range(len(cuts) - 1):
        if cuts[i] <= value <= cuts[i + 1]:
            return f"between the {pcts[i]}th and {pcts[i+1]}th percentile"
    return f"above the 90th percentile"


def get_citations(risk_factor, pubmed_df, n=2):
    matches = pubmed_df[
        pubmed_df["matched_risk_factors"].astype(str).str.contains(
            risk_factor, case=False, na=False
        )
    ]
    if matches.empty:
        return matches
    high = matches[matches["is_high_evidence"] == True]
    pool = high if len(high) >= n else matches
    return pool.sample(min(n, len(pool)))[["title", "citation", "url"]]


def build_context_block(user_input, trained, nhanes_summary, pubmed):
    """Compute everything deterministically; format as a structured block
    for the LLM. The LLM never re-runs these calculations."""
    risk = trained.predict_proba(user_input)
    tier = risk_tier(risk)

    feature_to_nhanes = {
        "age": "age_years", "BMI": "bmi", "sysBP": "sys_bp", "totChol": "total_chol",
    }
    percentile_lines = []
    for feat, nv in feature_to_nhanes.items():
        row = nhanes_summary[nhanes_summary["nhanes_variable"] == nv]
        if not row.empty:
            r = percentile_rank(user_input[feat], row.iloc[0])
            if r:
                percentile_lines.append(f"  - {feat} = {user_input[feat]}: {r} of U.S. adults aged 30–65")

    if user_input["cigsPerDay"] > 0:
        smoking_row = nhanes_summary[nhanes_summary["nhanes_variable"] == "is_current_smoker"]
        if not smoking_row.empty:
            pct = smoking_row.iloc[0]["mean"] * 100
            percentile_lines.append(f"  - smoking: user is a current smoker; {pct:.1f}% of U.S. adults currently smoke")
    if user_input["diabetes"] == 1:
        d_row = nhanes_summary[nhanes_summary["nhanes_variable"] == "has_diabetes"]
        if not d_row.empty:
            pct = d_row.iloc[0]["mean"] * 100
            percentile_lines.append(f"  - diabetes: user has diabetes; {pct:.1f}% of U.S. adults have diabetes")

    scenarios = counterfactual_scenarios(trained, user_input, top_n=3)
    cf_lines = []
    for s in scenarios:
        cf_lines.append(
            f"  - {s['label']}: from {s['original_value']:g} to {s['new_value']:g} "
            f"-> risk {s['baseline_risk']*100:.0f}% -> {s['new_risk']*100:.0f}% "
            f"(delta -{s['delta']*100:.1f} pp)"
        )

    citation_lines = []
    for s in scenarios[:3]:
        cites = get_citations(s["vocab"], pubmed, n=1)
        for _, row in cites.iterrows():
            citation_lines.append(
                f"  - [{s['vocab']}] {row['title']} | {row['citation']} | {row['url']}"
            )

    block = f"""=== USER PROFILE ===
age={user_input['age']}, sex={'male' if user_input['male']==1 else 'female'}, BMI={user_input['BMI']}, sysBP={user_input['sysBP']}, totChol={user_input['totChol']}, cigsPerDay={user_input['cigsPerDay']}, diabetes={'yes' if user_input['diabetes']==1 else 'no'}

=== MODEL OUTPUT ===
10-year CHD risk: {risk*100:.1f}%
Tier: {tier} (Low <5%, Medium 5-20%, High >20%)

=== NHANES PERCENTILE CONTEXT ===
{chr(10).join(percentile_lines) if percentile_lines else '  (no percentile data available)'}

=== TOP COUNTERFACTUAL SCENARIOS (model-predicted, holds other variables constant) ===
{chr(10).join(cf_lines) if cf_lines else '  (user is already at or below all modifiable targets)'}

=== PUBMED CITATIONS FOR FLAGGED RISK FACTORS ===
{chr(10).join(citation_lines) if citation_lines else '  (no flagged factors)'}
"""
    return block, risk, tier, scenarios


# --- UI -------------------------------------------------------------------
st.set_page_config(page_title="Heart Risk Predictor", layout="wide")
st.title("Heart Disease Risk Predictor")
st.caption("Educational tool — not medical advice")

# NEW:
import os
openai.api_key = os.getenv("OPENAI_API_KEY")
grounding = os.getenv("GROUNDING")

if not openai.api_key:
    st.error("OPENAI_API_KEY not found. Add it under Settings → Variables and secrets.")
    st.stop()
if not grounding:
    st.error("GROUNDING not found. Add it under Settings → Variables and secrets.")
    st.stop()

# Load artifacts
trained = load_model()
nhanes_summary, nhanes_raw, pubmed = load_reference_data()

if "openai_model" not in st.session_state:
    st.session_state["openai_model"] = "gpt-5-mini"
if "messages" not in st.session_state:
    st.session_state.messages = []
if "context_block" not in st.session_state:
    st.session_state.context_block = None

# --- Sidebar input form --------------------------------------------------
with st.sidebar:
    st.header("Your information")
    with st.form("inputs"):
        age = st.number_input("Age (years)", 30, 90, 50)
        sex = st.selectbox("Sex", ["Female", "Male"])
        is_smoker = st.checkbox("Currently smoke?")
        cigs = st.number_input("Cigarettes/day (if smoker)", 0, 60, 0)
        sysBP = st.number_input("Systolic BP (mmHg)", 70, 250, 120)
        totChol = st.number_input("Total cholesterol (mg/dL)", 80, 500, 200)
        BMI = st.number_input("BMI", 12.0, 70.0, 25.0, step=0.1)
        diabetes = st.checkbox("Diabetes?")
        submit = st.form_submit_button("Compute my risk")

    if submit:
        user_input = {
            "age": age, "male": 1 if sex == "Male" else 0,
            "cigsPerDay": cigs if is_smoker else 0,
            "sysBP": sysBP, "totChol": totChol, "BMI": BMI,
            "diabetes": 1 if diabetes else 0,
        }
        block, risk, tier, scenarios = build_context_block(
            user_input, trained, nhanes_summary, pubmed
        )
        st.session_state.context_block = block
        # Reset the conversation so the bot starts fresh on the new profile
        st.session_state.messages = []
        # Seed the conversation with an opening user message that asks the
        # bot to explain the results
        st.session_state.messages.append({
            "role": "user",
            "content": "Please walk me through my results.",
        })

# --- Main panel: chat ----------------------------------------------------
if st.session_state.context_block is None:
    st.info("Fill in your information in the sidebar and click 'Compute my risk' to begin.")
else:
    # Show the deterministic numbers prominently — they're the source of truth
    with st.expander("Computed results (deterministic — what the chatbot sees)", expanded=False):
        st.code(st.session_state.context_block, language="text")

    # Replay chat history
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])

    # Handle the seeded opening turn (if no assistant response yet)
    needs_response = (
        len(st.session_state.messages) > 0
        and st.session_state.messages[-1]["role"] == "user"
    )

    # Capture follow-up input
    user_followup = st.chat_input("Ask a follow-up question about your results")
    if user_followup:
        st.session_state.messages.append({"role": "user", "content": user_followup})
        with st.chat_message("user"):
            st.markdown(user_followup)
        needs_response = True

    if needs_response:
        with st.chat_message("assistant"):
            message_placeholder = st.empty()
            full_response = ""

            # Inject the context block into the system instructions for every turn.
            # This keeps the bot grounded in the numbers across follow-ups.
            instructions_with_context = (
                grounding
                + "\n\n=== DETERMINISTIC RESULTS FOR THIS USER (DO NOT MODIFY OR RECALCULATE) ===\n"
                + st.session_state.context_block
            )

            response = openai.responses.create(
                model=st.session_state["openai_model"],
                instructions=instructions_with_context,
                input=[
                    {"role": m["role"], "content": m["content"]}
                    for m in st.session_state.messages
                ],
                stream=True,
            )

            for event in response:
                if event.type == "response.output_text.delta":
                    full_response += event.delta.replace("\\$", "$").replace("$", "\\$")
                    message_placeholder.markdown(full_response + "▌")

            message_placeholder.markdown(full_response)

        st.session_state.messages.append({"role": "assistant", "content": full_response})