Supervivencia: ¿cuándo se va el cliente?

"""#2 Análisis de supervivencia: ¿cuándo se va el cliente?

Pasamos de "¿se va?" a "¿cuándo se va?" con Kaplan-Meier y Cox Proportional Hazards.
Datos: IBM Telco Customer Churn (`tenure` = meses de permanencia, `Churn` = evento).
"""
import numpy as np
import pandas as pd
from lifelines import CoxPHFitter, KaplanMeierFitter, WeibullAFTFitter
from lifelines.statistics import multivariate_logrank_test, proportional_hazard_test
from lifelines.utils import concordance_index, k_fold_cross_validation
from sklearn.model_selection import train_test_split

from common import SEED, bars, download, export, line, table
from p01_churn import URL
from validacion import model_card


def load() -> pd.DataFrame:
    df = pd.read_csv(download(URL, "telco_churn.csv"))
    df["evento"] = (df["Churn"] == "Yes").astype(int)
    df["tenure"] = df["tenure"].clip(lower=0.5)  # clientes con 0 meses: medio mes
    return df


def km_points(kmf: KaplanMeierFitter, max_t=72):
    sf = kmf.survival_function_.iloc[:, 0]
    return [[float(t), round(float(v) * 100, 2)] for t, v in sf.items() if t <= max_t]


def km_band(kmf: KaplanMeierFitter, max_t=72):
    ci = kmf.confidence_interval_survival_function_
    return [[float(t), round(float(lo) * 100, 2), round(float(hi) * 100, 2)]
            for t, (lo, hi) in zip(ci.index, ci.values) if t <= max_t]


def s_en(curva, t):
    return next(v for x, v in reversed(curva) if x <= t)


VARS = ["Contract", "PaymentMethod", "InternetService", "MonthlyCharges", "SeniorCitizen", "Partner",
        "Dependents", "PaperlessBilling"]


def design(df: pd.DataFrame) -> pd.DataFrame:
    X = pd.get_dummies(df[VARS], drop_first=True, dtype=float)
    X["MonthlyCharges"] = (X["MonthlyCharges"] - 65) / 30          # escala interpretable: +30 USD
    X["tenure"], X["evento"] = df["tenure"].values, df["evento"].values
    return X


def run():
    df = load()
    tr, te = train_test_split(df, test_size=0.3, stratify=df["evento"], random_state=SEED)

    # --- Kaplan-Meier por contrato con IC 95 % ---
    curvas, bandas, medianas = {}, {}, {}
    nombres = {"Month-to-month": "Mes a mes", "One year": "Un año", "Two year": "Dos años"}
    for contrato, g in df.groupby("Contract"):
        kmf = KaplanMeierFitter().fit(g["tenure"], g["evento"], label=contrato)
        curvas[nombres[contrato]] = km_points(kmf)
        bandas[nombres[contrato]] = km_band(kmf)
        medianas[contrato] = kmf.median_survival_time_
    lr = multivariate_logrank_test(df["tenure"], df["Contract"], df["evento"])

    # --- Cox PH y verificación del supuesto de riesgos proporcionales (residuos de Schoenfeld) ---
    X_tr, X_te = design(tr), design(te)
    cph = CoxPHFitter(penalizer=0.01).fit(X_tr, "tenure", "evento")
    ph = proportional_hazard_test(cph, X_tr, time_transform="rank").summary
    viola = ph[ph["p"] < 0.01].index.tolist()

    # Cambio de enfoque: estratificar por contrato (cada contrato tiene su propia curva base)
    X_tr_s = X_tr.drop(columns=[c for c in X_tr if c.startswith("Contract")]).assign(Contract=tr["Contract"].values)
    X_te_s = X_te.drop(columns=[c for c in X_te if c.startswith("Contract")]).assign(Contract=te["Contract"].values)
    cph_s = CoxPHFitter(penalizer=0.01).fit(X_tr_s, "tenure", "evento", strata=["Contract"])
    # Alternativa paramétrica: Weibull AFT (permite extrapolar más allá de 72 meses)
    aft = WeibullAFTFitter(penalizer=0.01).fit(X_tr, "tenure", "evento")

    # C-index con una métrica común a los tres modelos: tiempo medio restringido de supervivencia
    # (RMST a 72 meses). Ojo: con estratos lifelines reordena las filas, por eso se re-indexa.
    tiempos = np.arange(1, 73)
    def rmst(modelo, X):
        return modelo.predict_survival_function(X, times=tiempos).T.reindex(X.index).sum(axis=1).values
    def c_test(modelo, X):
        return concordance_index(X["tenure"], rmst(modelo, X), X["evento"])
    c_cox, c_strat, c_aft = c_test(cph, X_te), c_test(cph_s, X_te_s), c_test(aft, X_te)
    # Validación cruzada 5-fold de la concordancia (sobre entrenamiento)
    cv_c = k_fold_cross_validation(CoxPHFitter(penalizer=0.01), X_tr, "tenure", "evento", k=5,
                                   scoring_method="concordance_index", seed=SEED)

    # --- Calibración a 12 meses: S(12) predicha vs. Kaplan-Meier observada por quintil de riesgo ---
    s12_pred = cph_s.predict_survival_function(X_te_s, times=[12]).T.reindex(X_te_s.index)[12].values
    q = pd.qcut(s12_pred, 5, labels=False, duplicates="drop")
    calib = []
    for k in sorted(np.unique(q)):
        m = q == k
        km = KaplanMeierFitter().fit(te["tenure"].values[m], te["evento"].values[m])
        calib.append([round(float(s12_pred[m].mean() * 100), 2), round(float(km.predict(12) * 100), 2)])
    err_cal = float(np.mean([abs(a - b) for a, b in calib]))

    # --- Explicabilidad: perfiles tipo ---
    perfiles = {
        "Mes a mes · fibra · cheque electrónico": dict(Contract="Month-to-month", InternetService="Fiber optic",
                                                        PaymentMethod="Electronic check"),
        "Mes a mes · DSL · tarjeta automática": dict(Contract="Month-to-month", InternetService="DSL",
                                                      PaymentMethod="Credit card (automatic)"),
        "Un año · fibra · transferencia automática": dict(Contract="One year", InternetService="Fiber optic",
                                                           PaymentMethod="Bank transfer (automatic)"),
    }
    base = df.iloc[[0]].copy()
    base[["MonthlyCharges", "SeniorCitizen", "Partner", "Dependents", "PaperlessBilling"]] = [70, 0, "No", "No", "Yes"]
    curvas_perfil = {}
    for nombre, cambios in perfiles.items():
        fila = base.copy()
        for k, v in cambios.items():
            fila[k] = v
        Xp = design(pd.concat([fila, df]).head(len(df) + 1)).iloc[[0]].drop(columns=[c for c in X_tr if c.startswith("Contract")])
        Xp["Contract"] = cambios["Contract"]
        sf = cph_s.predict_survival_function(Xp, times=np.arange(0, 73, 3)).iloc[:, 0]
        curvas_perfil[nombre] = [[float(t), round(float(v) * 100, 2)] for t, v in sf.items()]

    resumen = cph_s.summary
    hr = resumen[["coef", "exp(coef)", "exp(coef) lower 95%", "exp(coef) upper 95%", "p"]].sort_values("coef")
    hr_tbl = hr.reset_index().rename(columns={"covariate": "Variable", "coef": "Coef (log-HR)", "exp(coef)": "HR",
                                              "exp(coef) lower 95%": "HR IC 2.5 %", "exp(coef) upper 95%": "HR IC 97.5 %"})
    for c in ["Coef (log-HR)", "HR", "HR IC 2.5 %", "HR IC 97.5 %"]:
        hr_tbl[c] = hr_tbl[c].round(3)
    hr_tbl["p"] = hr_tbl["p"].map(lambda v: "<0.001" if v < 0.001 else f"{v:.3f}")
    ph_tbl = ph.reset_index().rename(columns={"index": "Variable", "test_statistic": "Estadístico", "p": "p"})
    ph_tbl = ph_tbl[["Variable", "Estadístico", "p"]].round(4)
    ph_tbl["¿Cumple PH?"] = np.where(ph_tbl["p"] < 0.01, "No", "Sí")
    val_tbl = pd.DataFrame({
        "Modelo": ["Cox PH", "Cox estratificado por contrato", "Weibull AFT"],
        "C-index prueba": [round(c_cox, 3), round(c_strat, 3), round(c_aft, 3)],
        "Notas": [f"CV 5-fold: {np.mean(cv_c):.3f} ± {np.std(cv_c):.3f}; viola PH en {len(viola)} variables",
                  "Absorbe la violación PH de Contract (curva base por contrato)",
                  f"Paramétrico; AIC {aft.AIC_:,.0f}"],
    })

    s12_mm, s12_2y = s_en(curvas["Mes a mes"], 12), s_en(curvas["Dos años"], 12)
    hr_s = cph_s.summary["exp(coef)"]
    result = {
        "title": "Supervivencia: ¿cuándo se va el cliente?",
        "summary": [
            f"Un cliente mes a mes tiene {s12_mm:.0f} % de probabilidad de seguir activo al mes 12 (2 años: "
            f"{s12_2y:.0f} %); vida mediana {medianas['Month-to-month']:.0f} meses. Log-rank χ² = {lr.test_statistic:,.0f}, p < 0.001.",
            f"La prueba de Schoenfeld rechaza riesgos proporcionales (p < 0.01) en {len(viola)} variables: "
            f"{', '.join(viola)}. El contrato de plazo impide salir antes de su vencimiento, así que su efecto "
            "cambia con el tiempo. Cambio de enfoque: Cox estratificado por contrato.",
            f"Concordancia en prueba (sobre RMST a 72 meses): Cox {c_cox:.3f}, estratificado {c_strat:.3f}, Weibull AFT {c_aft:.3f}; "
            f"CV 5-fold {np.mean(cv_c):.3f} ± {np.std(cv_c):.3f}. Calibración a 12 meses: error medio {err_cal:.1f} pp por quintil.",
            f"Palancas (HR, modelo estratificado): fibra óptica ×{hr_s['InternetService_Fiber optic']:.2f}, cheque "
            f"electrónico ×{hr_s['PaymentMethod_Electronic check']:.2f}, tener pareja ×{hr_s['Partner_Yes']:.2f}.",
        ],
        "metrics": [
            {"label": "C-index estratificado", "value": f"{c_strat:.3f}"},
            {"label": "C-index CV 5-fold", "value": f"{np.mean(cv_c):.3f} ± {np.std(cv_c):.3f}"},
            {"label": "Error calibración S(12)", "value": f"{err_cal:.1f} pp"},
            {"label": "Mediana mes a mes", "value": f"{medianas['Month-to-month']:.0f} meses"},
        ],
        "charts": [
            line("km", "Kaplan-Meier por tipo de contrato (IC 95 %)", curvas, bands=bandas,
                 x_label="Meses de permanencia", y_label="% de clientes que siguen activos", fmt="pct"),
            line("perfiles", "Curvas predichas por perfil de cliente (Cox estratificado)", curvas_perfil,
                 x_label="Meses", y_label="% que sigue activo", fmt="pct",
                 subtitle="Mismo cargo ($70) y demografía; cambian contrato, servicio y forma de pago"),
            line("cal12", "Calibración a 12 meses por quintil de riesgo",
                 {"Modelo": calib, "Referencia": [[min(c[0] for c in calib), min(c[0] for c in calib)], [100, 100]]},
                 x_label="S(12) predicha %", y_label="S(12) observada (Kaplan-Meier) %", fmt="pct"),
            bars("loghr", "Efectos del Cox estratificado (log-HR): positivo acelera la salida", hr.index,
                 hr["coef"].round(3).values, fmt="dec", subtitle="HR = e^coef · MonthlyCharges por +30 USD"),
            bars("ph", "Prueba de riesgos proporcionales (−log10 p; > 2 viola el supuesto)", ph.index,
                 (-np.log10(ph["p"].clip(lower=1e-300))).round(2).values, fmt="dec"),
        ],
        "tables": [table("Comparación de modelos de supervivencia", val_tbl),
                   table("Hazard ratios con IC 95 % (Cox estratificado)", hr_tbl),
                   table("Supuesto de riesgos proporcionales (Schoenfeld)", ph_tbl)],
        "card": model_card(
            "Kaplan-Meier para describir, Cox estratificado por contrato para explicar y Weibull AFT como "
            "alternativa paramétrica.",
            ["Split 70/30 estratificado por evento; C-index en prueba y CV 5-fold en entrenamiento.",
             "Supuesto de riesgos proporcionales con residuos de Schoenfeld (transformación de rango).",
             "Calibración a 12 meses: S(12) predicha vs. Kaplan-Meier observada por quintil."],
            ["Censura no informativa: los clientes activos no son distintos de los que se irían.",
             "`tenure` mide meses desde el alta; no hay variables que cambien en el tiempo."],
            ["El dataset es una foto: no se observa la historia de cambios de contrato.",
             "Hazard ratios son asociaciones, no efectos causales."]),
    }
    sample = df[["tenure", "evento", "Contract"]].sample(1500, random_state=1)
    export("p02-supervivencia", result, sample)
    return result


if __name__ == "__main__":
    run()