risk_model.py

"""
risk_model.py — CognitivePulse

Trains and serves an XGBoost classifier on the Alzheimer's dataset to produce:
  - A calibrated risk score (0-100) for any input patient profile
  - SHAP values for per-feature contribution to that score
  - Global feature importance for the population-level dashboard

SHAP note: shap.TreeExplainer receives model.get_booster() (the raw Booster
object) rather than the sklearn XGBClassifier wrapper. This avoids a bug in
XGBoost 2.x where base_score is serialised as scientific notation
(e.g. '[5.003238E-1]') which SHAP's string-to-float parser cannot parse,
raising: ValueError: could not convert string to float: '[5.003238E-1]'

On GPU (Colab T4): pass device="cuda" to train_model() to use XGBoost's
GPU-accelerated histogram method.
"""

from __future__ import annotations

import pickle
from pathlib import Path

import numpy as np
import pandas as pd
import xgboost as xgb
import shap
from sklearn.model_selection import StratifiedKFold, cross_val_score

from data_loader import FEATURE_COLS, TARGET_COL, FEATURE_META

MODEL_PATH         = Path(__file__).parent / "data" / "model.json"
EXPLAINER_PATH     = Path(__file__).parent / "data" / "shap_explainer.pkl"

XGB_PARAMS = {
    "n_estimators":     300,
    "max_depth":        5,
    "learning_rate":    0.05,
    "subsample":        0.8,
    "colsample_bytree": 0.8,
    "min_child_weight": 3,
    "gamma":            0.1,
    "reg_alpha":        0.1,
    "reg_lambda":       1.0,
    "scale_pos_weight": 1.83,
    "base_score":       0.5,
    "eval_metric":      "auc",
    "random_state":     42,
}


def train_model(df: pd.DataFrame, device: str = "cpu") -> tuple:
    """
    Trains XGBoost + SHAP TreeExplainer.
    Returns (model, explainer, cv_results_dict).
    """
    X = df[FEATURE_COLS].values.astype("float32")
    y = df[TARGET_COL].values

    params = XGB_PARAMS.copy()
    if device == "cuda":
        # XGBoost 1.7.x uses gpu_hist + gpu_id; 2.x uses device + hist
        import xgboost as _xgb
        xgb_major = int(_xgb.__version__.split(".")[0])
        if xgb_major >= 2:
            params["device"]      = "cuda"
            params["tree_method"] = "hist"
        else:
            params["tree_method"] = "gpu_hist"
            params["gpu_id"]      = 0

    model = xgb.XGBClassifier(**params)

    # 5-fold stratified cross-validation
    cv        = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
    cv_auc    = cross_val_score(model, X, y, cv=cv, scoring="roc_auc")
    cv_f1     = cross_val_score(model, X, y, cv=cv, scoring="f1")

    # Final fit on full dataset
    model.fit(X, y, verbose=False)

    # ------------------------------------------------------------------ #
    # FIX: pass the raw Booster to TreeExplainer, NOT the sklearn wrapper.
    # XGBoost 2.x serialises base_score as '[5.003238E-1]' inside the
    # sklearn wrapper's JSON, which SHAP cannot parse. The raw Booster
    # serialises the same value as a plain float and parses correctly.
    # ------------------------------------------------------------------ #
    explainer = shap.TreeExplainer(model.get_booster())

    cv_results = {
        "auc_mean": round(float(cv_auc.mean()), 4),
        "auc_std":  round(float(cv_auc.std()),  4),
        "f1_mean":  round(float(cv_f1.mean()),  4),
        "f1_std":   round(float(cv_f1.std()),   4),
    }

    # Cache to disk — use pickle for the sklearn wrapper (avoids XGBoost 1.7.x
    # _estimator_type bug in save_model()) and separately cache the booster JSON
    # so the explainer can reload without re-training.
    MODEL_PATH.parent.mkdir(exist_ok=True)
    with open(MODEL_PATH.with_suffix(".pkl"), "wb") as f:
        pickle.dump(model, f)
    with open(EXPLAINER_PATH, "wb") as f:
        pickle.dump(explainer, f)

    return model, explainer, cv_results


def load_or_train(df: pd.DataFrame, force_retrain: bool = False) -> tuple:
    """Loads cached model/explainer if available; trains from scratch otherwise."""
    model_pkl = MODEL_PATH.with_suffix(".pkl")
    if not force_retrain and model_pkl.exists() and EXPLAINER_PATH.exists():
        with open(model_pkl, "rb") as f:
            model = pickle.load(f)
        with open(EXPLAINER_PATH, "rb") as f:
            explainer = pickle.load(f)
        return model, explainer, None
    return train_model(df)


def predict_patient(model, explainer, patient: dict) -> dict:
    """
    Returns risk_score (0-100), risk_band, and per-feature SHAP contributions
    for a single patient supplied as a feature-name -> value dict.
    """
    row = pd.DataFrame([{f: patient.get(f, 0) for f in FEATURE_COLS}])
    X   = row[FEATURE_COLS].values.astype("float32")

    prob       = float(model.predict_proba(X)[0, 1])
    risk_score = round(prob * 100, 1)

    # float32 required when explainer wraps a raw Booster
    shap_vals     = explainer.shap_values(X)[0]
    contributions = {
        FEATURE_COLS[i]: round(float(shap_vals[i]), 4)
        for i in range(len(FEATURE_COLS))
    }

    return {
        "risk_score":        risk_score,
        "risk_probability":  round(prob, 4),
        "risk_band":         _risk_band(prob),
        "shap_contributions": contributions,
        "top_drivers":       _top_drivers(contributions),
    }


def _risk_band(prob: float) -> str:
    if prob < 0.25: return "low"
    if prob < 0.50: return "moderate"
    if prob < 0.75: return "elevated"
    return "high"


def _top_drivers(contributions: dict, n: int = 5) -> list:
    sorted_feats = sorted(contributions.items(), key=lambda x: abs(x[1]), reverse=True)
    return [
        {
            "feature":    feat,
            "label":      FEATURE_META.get(feat, {}).get("label", feat),
            "shap_value": val,
            "direction":  "increases risk" if val > 0 else "decreases risk",
            "modifiable": FEATURE_META.get(feat, {}).get("modifiable", False),
        }
        for feat, val in sorted_feats[:n]
    ]


def global_feature_importance(model) -> pd.DataFrame:
    importance = model.get_booster().get_score(importance_type="gain")
    return pd.DataFrame(
        [(FEATURE_META.get(k, {}).get("label", k), round(v, 2))
         for k, v in importance.items()],
        columns=["Feature", "Importance"],
    ).sort_values("Importance", ascending=False).reset_index(drop=True)


if __name__ == "__main__":
    from data_loader import load_dataset
    df, source = load_dataset()
    print(f"Training on {source} data ({len(df)} rows)...")
    model, explainer, cv = train_model(df)
    print(f"CV AUC : {cv['auc_mean']:.4f} ± {cv['auc_std']:.4f}")
    print(f"CV F1  : {cv['f1_mean']:.4f} ± {cv['f1_std']:.4f}")

    sample = {
        "Age": 68, "Gender": 0, "Ethnicity": 0, "EducationLevel": 2,
        "BMI": 29.5, "Smoking": 0, "AlcoholConsumption": 5.0,
        "PhysicalActivity": 2.5, "DietQuality": 5.0, "SleepQuality": 6.0,
        "FamilyHistoryAlzheimers": 1, "CardiovascularDisease": 1,
        "Diabetes": 0, "Depression": 0, "HeadInjury": 0, "Hypertension": 1,
        "SystolicBP": 148, "DiastolicBP": 88, "CholesterolTotal": 240,
        "CholesterolLDL": 158, "CholesterolHDL": 45, "CholesterolTriglycerides": 185,
        "MMSE": 25, "FunctionalAssessment": 7.0, "MemoryComplaints": 1,
        "BehavioralProblems": 0, "ADL": 7.5, "Confusion": 0,
        "Disorientation": 0, "PersonalityChanges": 0,
        "DifficultyCompletingTasks": 0, "Forgetfulness": 1,
    }
    import json
    print(json.dumps(predict_patient(model, explainer, sample), indent=2))