import gradio as gr
import joblib
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

# Load model and scaler
model = joblib.load('rsf_tuned.pkl')
scaler = joblib.load('scaler.pkl')

# --- ORDINAL MAPPINGS ---

wealth_map = {
    "poorest": 0, "poorer": 1, "middle": 2, "richer": 3, "richest": 4
}

mat_edu_map = {
    "no education": 0, "primary": 1, "secondary": 2, "higher": 3
}

# --- FEATURE NAMES (same as training) ---
numeric_cols = [
    'birth_order', 'mat_age', 'mat_height_cm', 'mat_weight_kg',
    'wealth_ordinal', 'mat_edu_ordinal', 'multiple_birth_binary'
]

categorical_cols = [
    'child_sex', 'religion', 'water_improved', 'sanitation_improved',
    'partner_edu', 'region', 'residence'
]

cat_feature_names = [
    'child_sex_female', 'child_sex_male',
    'religion_catholic', 'religion_islam', 'religion_Not asked/Missing',
    'religion_other', 'religion_other christian', 'religion_traditionalist',
    'water_improved_Improved', 'water_improved_Not asked/Missing', 'water_improved_Unimproved',
    'sanitation_improved_Improved', 'sanitation_improved_Not asked/Missing', 'sanitation_improved_Unimproved',
    'partner_edu_don\'t know', 'partner_edu_higher', 'partner_edu_no education',
    'partner_edu_Not asked/Missing', 'partner_edu_primary', 'partner_edu_secondary',
    'region_North Central', 'region_North East', 'region_North West',
    'region_South East', 'region_South South', 'region_South West',
    'residence_rural', 'residence_urban'
]

def create_inputs():
    return [
        gr.Number(label="Birth Order", value=1, minimum=1, maximum=20),
        gr.Number(label="Maternal Age (years)", value=25, minimum=12, maximum=49),
        gr.Number(label="Maternal Height (cm)", value=160, minimum=100, maximum=200),
        gr.Number(label="Maternal Weight (kg)", value=60, minimum=30, maximum=120),
        gr.Dropdown(
            choices=[("Poorest", "poorest"), ("Poorer", "poorer"), ("Middle", "middle"),
                     ("Richer", "richer"), ("Richest", "richest")],
            label="Wealth Index",
            value="middle"
        ),
        gr.Dropdown(
            choices=[("No Education", "no education"), ("Primary", "primary"),
                     ("Secondary", "secondary"), ("Higher", "higher")],
            label="Maternal Education",
            value="primary"
        ),
        gr.Dropdown(choices=[("No", 0), ("Yes", 1)], label="Multiple Birth", value=0),
        gr.Dropdown(choices=["male", "female"], label="Child Sex", value="male"),
        gr.Dropdown(
            choices=["islam", "catholic", "other christian", "traditionalist", "other"],
            label="Religion",
            value="islam"
        ),
        gr.Dropdown(
            choices=["Improved", "Unimproved"],
            label="Water Source",
            value="Improved"
        ),
        gr.Dropdown(
            choices=["Improved", "Unimproved"],
            label="Sanitation",
            value="Improved"
        ),
        gr.Dropdown(
            choices=["no education", "primary", "secondary", "higher"],
            label="Partner Education",
            value="primary"
        ),
        gr.Dropdown(
            choices=["North West", "North East", "North Central", "South West", "South East", "South South"],
            label="Region",
            value="North West"
        ),
        gr.Dropdown(choices=["rural", "urban"], label="Residence", value="rural")
    ]

import spaces

@spaces.GPU
def predict_stunting(birth_order, mat_age, mat_height_cm, mat_weight_kg,
                     wealth_label, mat_edu_label, multiple_birth_binary,
                     child_sex, religion, water_improved, sanitation_improved,
                     partner_edu, region, residence):

    print("🔍 Starting prediction...")

    wealth_ordinal = wealth_map[wealth_label]
    mat_edu_ordinal = mat_edu_map[mat_edu_label]

    # Build raw input (14 features)
    input_data = pd.DataFrame([[
        birth_order, mat_age, mat_height_cm, mat_weight_kg,
        wealth_ordinal, mat_edu_ordinal, multiple_birth_binary,
        child_sex, religion, water_improved, sanitation_improved,
        partner_edu, region, residence
    ]], columns=numeric_cols + categorical_cols)

    print(f"📊 input_data shape: {input_data.shape}")

    # One-hot encoding
    one_hot_data = pd.DataFrame(0, index=input_data.index, columns=cat_feature_names)
    for col in categorical_cols:
        val = input_data[col].iloc[0]
        if pd.isna(val) or val in ["Not asked/Missing", "don't know"]:
            col_name = f"{col}_Not asked/Missing"
        else:
            col_name = f"{col}_{val}"
        if col_name in one_hot_data.columns:
            one_hot_data[col_name] = 1

    # Combine → 35 features
    X_sksurv = pd.concat([input_data[numeric_cols], one_hot_data], axis=1)
    print(f"📊 X_sksurv shape: {X_sksurv.shape}")  # Should be (1, 35)
    print(f"📊 Scaler expects: {scaler.n_features_in_} features")

    # The scaler was fit with specific feature names/order recorded in
    # scaler.feature_names_in_. Our one-hot + numeric concat above does not
    # guarantee that same order, and scaler.transform() raises a ValueError
    # if the column order doesn't match what was seen during fit. Reorder
    # to match before transforming.
    if hasattr(scaler, "feature_names_in_"):
        X_sksurv = X_sksurv[scaler.feature_names_in_]

    # Scale ALL 35 features
    X_scaled = pd.DataFrame(
        scaler.transform(X_sksurv),
        columns=X_sksurv.columns,
        index=X_sksurv.index
    )

    # Ensure column order matches what the model was trained on, if available
    if hasattr(model, "feature_names_in_"):
        X_scaled = X_scaled[model.feature_names_in_]

    # Predict
    risk = model.predict(X_scaled)[0]
    print(f"✅ Risk: {risk:.4f}")

    # Survival function for this sample
    surv_funcs = model.predict_survival_function(X_scaled)
    fn = surv_funcs[0]

    # Requested time points (months)
    times = np.array([6, 12, 18, 24, 36, 48])

    # The step function is only defined within the domain of the training
    # follow-up times. Querying outside that range raises a ValueError, so
    # clip the requested times into the valid domain before evaluating.
    domain_min, domain_max = fn.domain
    times_clipped = np.clip(times, domain_min, domain_max)

    surv_probs = fn(times_clipped)
    stunting_risks = 1 - surv_probs
    print(f"✅ Stunting risks: {stunting_risks}")

    # Plot (x-axis still shows the originally requested ages)
    fig, ax = plt.subplots(figsize=(8, 4))
    ax.plot(times, stunting_risks, marker='o', linestyle='-', color='steelblue', linewidth=2, markersize=8)
    ax.axhline(y=0.20, color='red', linestyle='--', label='20% Risk Threshold')
    ax.set_xlabel('Age (months)')
    ax.set_ylabel('Stunting Risk')
    ax.set_title('Stunting Risk by Age')
    ax.set_ylim(0, 1)
    ax.grid(True, linestyle=':', alpha=0.7)
    ax.legend()
    ax.set_xticks(times)

    # --- Build a detailed, human-readable explanation ---

    RISK_THRESHOLD = 0.20

    # Find first age (from requested times) where predicted stunting risk
    # crosses the 20% threshold, if it does within the given horizon.
    crossing_age = None
    for age, r in zip(times, stunting_risks):
        if r >= RISK_THRESHOLD:
            crossing_age = int(age)
            break

    if crossing_age is not None:
        threshold_text = (
            f"Predicted stunting risk is estimated to reach the **{RISK_THRESHOLD:.0%} threshold "
            f"by around {crossing_age} months** of age."
        )
    else:
        threshold_text = (
            f"Predicted stunting risk **does not reach the {RISK_THRESHOLD:.0%} threshold** "
            f"within the 48-month horizon shown."
        )

    # Risk at each timepoint, formatted as a simple bullet list
    age_risk_lines = "\n".join(
        f"- **{int(age)} months:** {r:.1%} estimated risk"
        for age, r in zip(times, stunting_risks)
    )

    # Highlight the strongest known risk/protective factors present in this input
    contributing_factors = []
    if wealth_ordinal <= 1:
        contributing_factors.append("lower household wealth quintile")
    if mat_edu_ordinal == 0:
        contributing_factors.append("mother has no formal education")
    if water_improved == "Unimproved":
        contributing_factors.append("unimproved water source")
    if sanitation_improved == "Unimproved":
        contributing_factors.append("unimproved sanitation")
    if multiple_birth_binary == 1:
        contributing_factors.append("multiple birth (twins/triplets)")
    if mat_age < 18 or mat_age > 40:
        contributing_factors.append("maternal age outside the 18–40 range")
    if birth_order >= 5:
        contributing_factors.append("higher birth order (5th child or later)")

    if contributing_factors:
        factors_text = "Factors in this input associated with **elevated** stunting risk in the literature:\n" + \
            "\n".join(f"- {f}" for f in contributing_factors)
    else:
        factors_text = "No major elevated-risk factors were flagged among the inputs provided."

    summary_md = f"""
## Stunting Risk Prediction

### Predicted stunting risk over time
{age_risk_lines}

{threshold_text}

### Factors noted for this case
{factors_text}

---
⚠️ **Reminder:** This is a research prototype with no external validation. Do not use for 
clinical or policy decisions. Predicted probabilities and risk scores should be interpreted 
as exploratory model output only.
"""

    return summary_md, fig

inputs = create_inputs()
demo = gr.Interface(
    fn=predict_stunting,
    inputs=inputs,
    outputs=[gr.Markdown(label="Risk Score"), gr.Plot(label="Stunting Risk Curve")],
    title="Stunting Risk Predictor",
    description="""
    ⚠️ **DISCLAIMER**: This is a **research prototype** developed using DHS data. 
    **External validation has not been performed**. This tool is for 
    **educational and exploratory purposes only** and is **NOT** intended 
    for clinical or policy decision-making.
    
    Predicts stunting risk at ages 6, 12, 18, 24, 36, and 48 months using a 
    Random Survival Forest model.
    """,
    examples=[
        [1, 25, 160, 60, "middle", "primary", 0, "male", "islam", "Improved", "Improved", "primary", "North West", "rural"],
        [4, 18, 150, 55, "poorest", "no education", 0, "female", "catholic", "Unimproved", "Unimproved", "no education", "South East", "rural"],
        [2, 35, 170, 70, "richest", "higher", 1, "male", "other christian", "Improved", "Improved", "secondary", "South West", "urban"]
    ],
    #allow_flagging="never"
)

if __name__ == "__main__":
    print("🚀 Launching app...")
    demo.launch(server_name="0.0.0.0", server_port=7860)