# src/calibration_utils.py
"""
Calibration assessment + bootstrap confidence intervals for binary
classifiers and Cox survival models.
 
Designed to slot into the existing manuscript pipeline:
  - call signatures mirror inference_utils.compute_metrics()
  - returns dicts that downstream code can merge into existing metric dicts
  - matplotlib figures use the same style as the existing ROC/PR plots
"""
 
from __future__ import annotations
 
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.metrics import roc_auc_score, brier_score_loss
from sklearn.calibration import calibration_curve
from scipy.special import logit
import statsmodels.api as sm
from lifelines.utils import concordance_index
 
 
# ----------------------------------------------------------------------
# Helpers
# ----------------------------------------------------------------------
 
def _clean_binary_inputs(y_true, y_prob):
    """Mirror the NaN-handling pattern used in inference_utils.compute_metrics()."""
    y_true = pd.to_numeric(pd.Series(y_true).astype(str).str.strip(), errors="coerce")
    y_prob = pd.to_numeric(pd.Series(y_prob), errors="coerce")
    valid = y_true.notna() & y_prob.notna()
    y_true = y_true.loc[valid].astype(int).to_numpy()
    y_prob = y_prob.loc[valid].astype(float).to_numpy()
    return y_true, y_prob
 
 
def _safe_logit(p, eps=1e-6):
    return logit(np.clip(p, eps, 1.0 - eps))
 
 
# ----------------------------------------------------------------------
# Bootstrap CI helpers
# ----------------------------------------------------------------------
 
def bootstrap_auroc_ci(y_true, y_prob, n_bootstraps=1000, seed=42):
    """
    Returns (auroc_point, ci_low, ci_high).
    Stratified bootstrap (preserves event prevalence).
    """
    y_true, y_prob = _clean_binary_inputs(y_true, y_prob)
    if len(y_true) == 0 or len(np.unique(y_true)) < 2:
        return (np.nan, np.nan, np.nan)
 
    point = float(roc_auc_score(y_true, y_prob))
    pos_idx = np.where(y_true == 1)[0]
    neg_idx = np.where(y_true == 0)[0]
    rng = np.random.default_rng(seed)
 
    boot = []
    for _ in range(n_bootstraps):
        pos_b = rng.choice(pos_idx, size=len(pos_idx), replace=True)
        neg_b = rng.choice(neg_idx, size=len(neg_idx), replace=True)
        idx = np.concatenate([pos_b, neg_b])
        try:
            boot.append(roc_auc_score(y_true[idx], y_prob[idx]))
        except ValueError:
            continue
 
    if len(boot) == 0:
        return (point, np.nan, np.nan)
    lo, hi = np.percentile(boot, [2.5, 97.5])
    return (point, float(lo), float(hi))
 
 
def bootstrap_c_index_ci(durations, events, risk_scores, n_bootstraps=1000, seed=42):
    """
    Bootstrap C-index for a Cox-style risk score (higher score = higher risk).
    Returns (c_point, ci_low, ci_high).
    """
    durations = np.asarray(durations, dtype=float).ravel()
    events = np.asarray(events, dtype=int).ravel()
    risk_scores = np.asarray(risk_scores, dtype=float).ravel()
 
    valid = ~(np.isnan(durations) | np.isnan(risk_scores)) & (durations > 0)
    durations = durations[valid]
    events = events[valid]
    risk_scores = risk_scores[valid]
 
    if len(durations) < 10 or events.sum() < 5:
        return (np.nan, np.nan, np.nan)
 
    try:
        point = float(concordance_index(durations, -risk_scores, events))
    except Exception:
        return (np.nan, np.nan, np.nan)
 
    rng = np.random.default_rng(seed)
    n = len(durations)
    boot = []
    for _ in range(n_bootstraps):
        idx = rng.choice(n, size=n, replace=True)
        if np.asarray(events)[idx].sum() < 2:
            continue
        try:
            boot.append(concordance_index(durations[idx], -risk_scores[idx], events[idx]))
        except Exception:
            continue
 
    if len(boot) == 0:
        return (point, np.nan, np.nan)
    lo, hi = np.percentile(boot, [2.5, 97.5])
    return (point, float(lo), float(hi))
 
 
# ----------------------------------------------------------------------
# Core calibration stats
# ----------------------------------------------------------------------
 
def compute_calibration_stats(y_true, y_prob, n_bins=10):
    """
    Compute calibration intercept, slope, decile-level points, and Brier score.
 
    Returns a dict with the same flat-key style as compute_metrics():
        {
            'N': int, 'Events': int, 'Prevalence': float,
            'Brier': float, 'CalibrationInTheLarge': float,
            'CalibrationIntercept': float,
            'CalibrationIntercept_CI_low': float,
            'CalibrationIntercept_CI_high': float,
            'CalibrationSlope': float,
            'CalibrationSlope_CI_low': float,
            'CalibrationSlope_CI_high': float,
            'BinProbTrue': np.ndarray,   # observed event rate per bin
            'BinProbPred': np.ndarray,   # mean predicted prob per bin
            'BinCounts':   np.ndarray,   # N per bin
        }
    """
    y_true, y_prob = _clean_binary_inputs(y_true, y_prob)
 
    out = {
        "N": int(len(y_true)),
        "Events": int(y_true.sum()) if len(y_true) else 0,
        "Prevalence": float(y_true.mean()) if len(y_true) else np.nan,
        "Brier": np.nan,
        "CalibrationInTheLarge": np.nan,
        "CalibrationIntercept": np.nan,
        "CalibrationIntercept_CI_low": np.nan,
        "CalibrationIntercept_CI_high": np.nan,
        "CalibrationSlope": np.nan,
        "CalibrationSlope_CI_low": np.nan,
        "CalibrationSlope_CI_high": np.nan,
        "BinProbTrue": np.array([]),
        "BinProbPred": np.array([]),
        "BinCounts": np.array([]),
    }
 
    if len(y_true) < 20 or len(np.unique(y_true)) < 2:
        return out
 
    # Brier
    out["Brier"] = float(brier_score_loss(y_true, np.clip(y_prob, 1e-15, 1 - 1e-15)))
 
    # Calibration-in-the-large
    out["CalibrationInTheLarge"] = float(y_true.mean() - y_prob.mean())
 
    # Decile points
    try:
        prob_true, prob_pred = calibration_curve(y_true, y_prob, n_bins=n_bins, strategy="quantile")
        # Bin counts via quantile cut on predicted probabilities
        bin_edges = np.quantile(y_prob, np.linspace(0, 1, n_bins + 1))
        bin_edges[0] -= 1e-9
        bin_edges[-1] += 1e-9
        bin_idx = np.digitize(y_prob, bin_edges, right=True) - 1
        bin_idx = np.clip(bin_idx, 0, n_bins - 1)
        bin_counts = np.bincount(bin_idx, minlength=n_bins)[: len(prob_pred)]
 
        out["BinProbTrue"] = prob_true
        out["BinProbPred"] = prob_pred
        out["BinCounts"] = bin_counts
    except Exception:
        pass
 
    # Calibration intercept (offset) and slope via logistic recalibration:
    #   logit(p_obs) = intercept + slope * logit(p_pred)
    try:
        logits = _safe_logit(y_prob)
        X = sm.add_constant(logits, has_constant='add')   # force constant column
        # statsmodels expects float64 numpy arrays
        y_arr = np.asarray(y_true, dtype=np.float64).ravel()
        X_arr = np.asarray(X, dtype=np.float64)

        # Diagnostic: print shapes
        print(f"[calibration] X shape: {X_arr.shape}, y shape: {y_arr.shape}")
        print(f"[calibration] X has constant? first col unique: {np.unique(X_arr[:, 0])[:5]}")
        print(f"[calibration] logit range: {logits.min():.3f} to {logits.max():.3f}")

        result = sm.Logit(y_arr, X_arr).fit(disp=0, method='bfgs', maxiter=200)
        print(f"[calibration] params: {result.params}")
        print(f"[calibration] conf_int:\n{result.conf_int()}")

        intercept = float(result.params[0])
        slope = float(result.params[1])
        ci_arr = np.asarray(result.conf_int())

        out["CalibrationIntercept"] = intercept
        out["CalibrationIntercept_CI_low"] = float(ci_arr[0, 0])
        out["CalibrationIntercept_CI_high"] = float(ci_arr[0, 1])
        out["CalibrationSlope"] = slope
        out["CalibrationSlope_CI_low"] = float(ci_arr[1, 0])
        out["CalibrationSlope_CI_high"] = float(ci_arr[1, 1])
    except Exception as e:
        print(f"[calibration] FIT FAILED: {type(e).__name__}: {e}")
        import traceback
        traceback.print_exc()
 
    return out
 
 
# ----------------------------------------------------------------------
# Plotting
# ----------------------------------------------------------------------
 
def plot_calibration_curve(y_true, y_prob, title="Calibration", n_bins=10,
                           ax=None, return_stats=False):
    """
    Calibration plot with:
      - decile points (observed vs predicted), marker size proportional to N per bin
      - diagonal reference line (perfect calibration)
      - histogram of predicted probabilities along the bottom
      - intercept and slope (with 95% CIs) annotated in the title
      - Brier score annotated in the legend
 
    If ax is provided, draws into that axis. Otherwise creates a new fig/ax.
    Returns the matplotlib figure (and stats dict if return_stats=True).
    """
    stats = compute_calibration_stats(y_true, y_prob, n_bins=n_bins)
 
    if ax is None:
        fig, ax = plt.subplots(figsize=(6.5, 6.0))
    else:
        fig = ax.figure
 
    # Diagonal reference
    ax.plot([0, 1], [0, 1], linestyle="--", color="0.4", label="Perfect calibration")
 
    # Decile points with marker size proportional to bin count
    prob_true = stats["BinProbTrue"]
    prob_pred = stats["BinProbPred"]
    bin_counts = stats["BinCounts"]
 
    if len(prob_true) and len(prob_pred):
        if len(bin_counts) == len(prob_pred) and bin_counts.sum() > 0:
            sizes = 40 + 200 * (bin_counts / bin_counts.max())
        else:
            sizes = np.full(len(prob_pred), 80.0)
 
        ax.scatter(prob_pred, prob_true, s=sizes, color="#1F77B4",
                   edgecolor="white", linewidth=0.8, zorder=3, label="Model")
        ax.plot(prob_pred, prob_true, color="#1F77B4", alpha=0.5, zorder=2)
 
    # Annotation
    icpt = stats["CalibrationIntercept"]
    icpt_lo = stats["CalibrationIntercept_CI_low"]
    icpt_hi = stats["CalibrationIntercept_CI_high"]
    slope = stats["CalibrationSlope"]
    slope_lo = stats["CalibrationSlope_CI_low"]
    slope_hi = stats["CalibrationSlope_CI_high"]
    brier = stats["Brier"]
 
    def fmt(v):
        return "NA" if (v is None or np.isnan(v)) else f"{v:.2f}"
 
    full_title = (
        f"{title}\n"
        f"Intercept = {fmt(icpt)} ({fmt(icpt_lo)} to {fmt(icpt_hi)})    "
        f"Slope = {fmt(slope)} ({fmt(slope_lo)} to {fmt(slope_hi)})    "
        f"Brier = {fmt(brier)}"
    )
    ax.set_title(full_title, fontsize=10)
    ax.set_xlabel("Mean predicted probability")
    ax.set_ylabel("Observed event rate")
    ax.set_xlim(-0.02, 1.02)
    ax.set_ylim(-0.02, 1.02)
    ax.grid(alpha=0.3, linestyle=":")
    ax.legend(loc="upper left", fontsize=9, framealpha=0.85)
 
    # Histogram of predicted probabilities at the bottom (twin axis)
    y_true_arr, y_prob_arr = _clean_binary_inputs(y_true, y_prob)
    if len(y_prob_arr) > 0:
        ax2 = ax.twinx()
        ax2.hist(y_prob_arr, bins=30, range=(0, 1), color="0.7", alpha=0.5,
                 edgecolor="white", linewidth=0.3)
        ax2.set_ylabel("Count (predicted prob)", fontsize=9, color="0.5")
        ax2.tick_params(axis="y", labelsize=8, colors="0.5")
        # Keep histogram in the bottom third
        hist_max = ax2.get_ylim()[1]
        ax2.set_ylim(0, hist_max * 3)
        ax2.set_zorder(0)
        ax.set_zorder(1)
        ax.patch.set_alpha(0)
 
    fig.tight_layout()
    if return_stats:
        return fig, stats
    return fig
 
 
def plot_multi_cohort_calibration(cohort_results, ncols=3, figsize=None):
    """
    Composite calibration figure across cohorts. Use this for the
    supplementary figure referenced in the manuscript revision roadmap.
 
    cohort_results: dict of {cohort_name: (y_true, y_prob)}
    """
    n = len(cohort_results)
    if n == 0:
        return None
    nrows = int(np.ceil(n / ncols))
    if figsize is None:
        figsize = (5 * ncols, 4.5 * nrows)
    fig, axes = plt.subplots(nrows, ncols, figsize=figsize, squeeze=False)
    for i, (cohort_name, (y_true, y_prob)) in enumerate(cohort_results.items()):
        r, c = divmod(i, ncols)
        plot_calibration_curve(y_true, y_prob, title=cohort_name, ax=axes[r][c])
    # Hide unused panels
    for j in range(n, nrows * ncols):
        r, c = divmod(j, ncols)
        axes[r][c].axis("off")
    fig.tight_layout()
    return fig