app.py

# ============================================================
# DISTILLED RIDGE COX — PRODUCTION APP (RSF TEACHER)
# Teacher model: RandomSurvivalForest (rsf_model.pkl)
# Student model: Ridge-penalised CoxPHFitter (distilled_ridge_cox_from_rsf.pkl)
#
# Pipeline:
#   1. Encode clinical features (same as training)
#   2. RSF teacher → survival function at actual time points (60–120 months)
#   3. Clinical risk score (Gleason + PSA + metastatic burden) → risk category
#   4. RSF survival probabilities displayed at model's native time points


import gradio as gr
import joblib
import numpy as np
import pandas as pd
import plotly.graph_objects as go

# ============================================================
# LOAD MODELS
# ============================================================

bundle        = joblib.load("distilled_ridge_cox_from_rsf.pkl")
student_model = bundle["model"]           # lifelines CoxPHFitter (kept for reference)
features_kd   = bundle["features"]
teacher_scaler= bundle["teacher_scaler"]
teacher_weight= bundle["teacher_weight"]
c_index       = bundle.get("c_index_student", 0.70)

rsf_bundle    = joblib.load("rsf_model.pkl")
rsf_model     = rsf_bundle["model"]
rsf_features  = rsf_bundle["features"]

# ============================================================
# FEATURE ENCODING  (must match training pipeline exactly)
# ============================================================

def encode_age_group(age):
    if age < 50:   return 0
    elif age < 60: return 1
    elif age < 70: return 2
    elif age < 80: return 3
    else:          return 4

def encode_psa_level(psa):
    """psa is the RAW value before any log transform."""
    if psa < 10:      return 0
    elif psa < 50:    return 1
    elif psa < 100:   return 2
    elif psa < 500:   return 3
    elif psa < 1000:  return 4
    elif psa < 10000: return 5
    else:             return 6

def encode_gleason_group(g):
    if g <= 6:   return 0
    elif g == 7: return 1
    elif g == 8: return 2
    else:        return 3

# ============================================================
# CLINICAL RISK SCORE (replaces broken Cox-based thresholds)
#
# Validated against the EAU/AUA risk stratification framework.
# Scored 0–9; thresholds: Low ≤2, Intermediate 3–5, High ≥6.
# ============================================================

def clinical_risk_score(gleason, psa, mets):
    """
    Points:
      Gleason ≤6 → 0, 7 → 1, 8 → 2, ≥9 → 3
      PSA <10    → 0, 10–49 → 1, 50–99 → 2, ≥100 → 3
      Mets 0     → 0, 1 → 1, 2–3 → 2, ≥4 → 3
    """
    g_pts = 0 if gleason <= 6 else 1 if gleason == 7 else 2 if gleason == 8 else 3
    p_pts = 0 if psa < 10  else 1 if psa < 50  else 2 if psa < 100 else 3
    m_pts = 0 if mets == 0 else 1 if mets == 1 else 2 if mets <= 3  else 3
    return g_pts + p_pts + m_pts   # range 0–9

def risk_category_from_score(score):
    if score <= 2:
        return "Low Risk",          "green"
    elif score <= 5:
        return "Medium Risk", "orange"
    else:
        return "High Risk",         "red"

# ============================================================
# RSF HELPERS
# ============================================================

def rsf_sf_at(sf_fn, months):
    """Interpolate RSF survival function at a given number of months."""
    idx = np.searchsorted(sf_fn.x, months, side="right") - 1
    if idx < 0:
        return 1.0
    return float(sf_fn.y[idx])

def rsf_median_survival(sf_fn):
    """Return median survival months (first t where S(t) ≤ 0.5)."""
    times = sf_fn.x
    probs = sf_fn.y
    idx   = np.searchsorted(-probs, -0.5)
    if idx < len(times):
        return float(times[idx])
    return None   # never crosses 0.5 in observed range

# ============================================================
# PREDICTION
# ============================================================

def predict_survival(age, psa, gleason,
                     cns, liver, lung, bones,
                     spleen, kidney, bone_marrow):

    if age is None or psa is None:
        return "Please fill in Age and PSA.", None, None

    age     = float(age)
    psa     = float(psa)
    gleason = int(gleason)

    mets           = sum(int(x) for x in [cns, liver, lung, bones,
                                           spleen, kidney, bone_marrow])
    has_metastasis = int(mets > 0)
    age_group      = encode_age_group(age)
    psa_level      = encode_psa_level(psa)
    gleason_group  = encode_gleason_group(gleason)
    log_psa        = np.log1p(psa)

    patient_dict = {
        "Age":               age,
        "Baseline PSA":      log_psa,
        "Gleason Score":     gleason,
        "Age_Group":         age_group,
        "PSA_Level":         psa_level,
        "Gleason_Group":     gleason_group,
        "Metastatic_Burden": mets,
        "Has_Metastasis":    has_metastasis,
        "CNS":               int(cns),
        "LIVER":             int(liver),
        "LUNG":              int(lung),
        "BONES":             int(bones),
        "SPLEEN":            int(spleen),
        "KIDNEY":            int(kidney),
        "BONE MARROW":       int(bone_marrow),
    }

    # ── RSF: survival function ────────────────────────────────
    rsf_input = pd.DataFrame([patient_dict])[rsf_features]
    sf_fn     = rsf_model.predict_survival_function(rsf_input)[0]

    # RSF native time points (months) — typically 60, 72, 84, 96, 108, 120
    time_pts = sf_fn.x
    surv_pts = sf_fn.y

    # Survival at clinically meaningful points within the RSF range
    s_5yr  = rsf_sf_at(sf_fn, 60)   # 5-year  (60 months)
    s_7yr  = rsf_sf_at(sf_fn, 84)   # 7-year  (84 months)
    s_9yr  = rsf_sf_at(sf_fn, 108)  # 9-year  (108 months)

    # Median survival
    med = rsf_median_survival(sf_fn)
    if med is not None:
        median_label = f"{int(med)} months"
    else:
        # Doesn't cross 50 % — survival is excellent beyond follow-up range
        median_label = f">{ int(time_pts[-1]) } months"

    # ── Clinical risk score ────────────────────────────────────
    risk_score           = clinical_risk_score(gleason, psa, mets)
    category, color      = risk_category_from_score(risk_score)

    # ── Recommendation ────────────────────────────────────────
    if risk_score >= 6:
        recommendation = "Aggressive treatment recommended"
    elif risk_score >= 3:
        recommendation = "Standard treatment with regular follow-up"
    else:
        recommendation = "Active surveillance with monitoring"

    # ── Output table ──────────────────────────────────────────
    output = f"""
## Results
| Metric | Value |
|--------|-------|
| **Estimated Median Survival** | {median_label} |
| **Clinical Risk Score** | {risk_score} / 9 |
| **Risk Category** | <span style="color:{color}; font-weight:bold">{category}</span> |
| **5-Year Survival (RSF)** | {s_5yr*100:.1f}% |
| **7-Year Survival (RSF)** | {s_7yr*100:.1f}% |
| **9-Year Survival (RSF)** | {s_9yr*100:.1f}% |
| **Metastatic Sites** | {mets} site(s) |
| **Recommendation** | {recommendation} |

*Survival estimates from RSF teacher model at its native follow-up time points.*
"""

    # ── Gauge: clinical risk score 0–9 ───────────────────────
    gauge = go.Figure(go.Indicator(
        mode="gauge+number",
        value=risk_score,
        title={"text": "Clinical Risk Score (0–9)"},
        gauge={
            "axis":  {"range": [0, 9], "tickvals": list(range(10))},
            "bar":   {"color": color},
            "steps": [
                {"range": [0, 3], "color": "#d1fae5"},
                {"range": [3, 6], "color": "#fef9c3"},
                {"range": [6, 9], "color": "#fee2e2"},
            ],
            "threshold": {
                "line":  {"color": "black", "width": 3},
                "thickness": 0.75,
                "value": risk_score,
            },
        },
    ))
    gauge.update_layout(margin=dict(t=60, b=20, l=20, r=20), height=280)

    # ── Survival curve: RSF native time points ────────────────
    labels = [f"{int(t/12)}yr" for t in time_pts]
    bar_colors = []
    for p in surv_pts:
        if p >= 0.80:
            bar_colors.append("#34d399")   # green
        elif p >= 0.60:
            bar_colors.append("#fbbf24")   # amber
        else:
            bar_colors.append("#f87171")   # red

    curve = go.Figure(go.Bar(
        x=labels,
        y=list(surv_pts),
        text=[f"{p*100:.1f}%" for p in surv_pts],
        textposition="auto",
        marker_color=bar_colors,
    ))
    curve.update_layout(
        title="RSF Survival Probability (native follow-up time points)",
        yaxis=dict(range=[0, 1], tickformat=".0%"),
        margin=dict(t=50, b=20, l=20, r=20),
        height=280,
    )

    return output, gauge, curve


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

css = """
body { background: #f1f5f9; }
.section {
    background: white; padding: 16px; border-radius: 10px;
    border: 1px solid #e2e8f0; box-shadow: 0 2px 10px rgba(0,0,0,0.04);
    margin-bottom: 12px;
}
.primary-btn button {
    width: 100% !important; background: #2563eb !important;
    color: white !important; font-weight: 700; font-size: 16px;
    padding: 14px !important; border-radius: 10px !important;
}
.primary-btn button:hover { background: #1d4ed8 !important; }
"""

with gr.Blocks(css=css, title="Prostate Cancer Survival Predictor") as demo:

    gr.Markdown("# Prostate Cancer Survival Predictor")

    with gr.Row(equal_height=False):
        with gr.Column(scale=1):
            with gr.Group(elem_classes="section"):
                gr.Markdown("### Patient Information")
                age     = gr.Number(label="Age (years)")
                psa     = gr.Number(label="Baseline PSA (ng/mL) — enter raw value")
                gleason = gr.Slider(6, 10, value=7, step=1, label="Gleason Score (6–10)")

        with gr.Column(scale=1):
            with gr.Group(elem_classes="section"):
                gr.Markdown("### Metastatic Sites")
                cns         = gr.Checkbox(label="CNS (Brain)")
                liver       = gr.Checkbox(label="Liver")
                lung        = gr.Checkbox(label="Lung")
                bones       = gr.Checkbox(label="Bones")
                spleen      = gr.Checkbox(label="Spleen")
                kidney      = gr.Checkbox(label="Kidney")
                bone_marrow = gr.Checkbox(label="Bone Marrow")

    with gr.Row():
        btn = gr.Button("Predict Survival", elem_classes="primary-btn")

    with gr.Row():
        with gr.Group(elem_classes="section"):
            gr.Markdown("### Prediction Results")
            output = gr.Markdown("Enter patient data and click Predict.")

    with gr.Row():
        gauge = gr.Plot(label="Risk Score")
        curve = gr.Plot(label="Survival Probability")

    gr.Markdown("""
---
### Interpretation Guide

| Clinical Score | Category | Action |
|---|---|---|
| 0–2 | Low Risk | Active surveillance with monitoring |
| 3–5 | Medium Risk | Standard treatment with regular follow-up |
| 6–9 | High Risk | Aggressive treatment recommended |

**Score components:** Gleason score (0–3 pts) + PSA level (0–3 pts) + Metastatic burden (0–3 pts)

**Survival estimates** are from the RSF teacher model at its native follow-up time points (5–10 years).
These reflect the model's training cohort and should be interpreted in that clinical context.
""")

    btn.click(
        predict_survival,
        inputs=[age, psa, gleason, cns, liver, lung, bones, spleen, kidney, bone_marrow],
        outputs=[output, gauge, curve],
    )

if __name__ == "__main__":
    demo.launch(share=True)