import os, sys, warnings, subprocess

# ─── AUTO-INSTALL DEPENDENCIES ────────────────────────────────────────────────
def _ensure_package(pkg_import, pkg_install=None):
    """Install a package if it is not already available."""
    pkg_install = pkg_install or pkg_import
    try:
        __import__(pkg_import)
    except ImportError:
        print(f"  Installing {pkg_install}...")
        subprocess.check_call(
            [sys.executable, "-m", "pip", "install", pkg_install, "-q"],
            stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
        )
        print(f"  {pkg_install} installed successfully.")

_ensure_package("pyreadstat")
_ensure_package("xgboost")
_ensure_package("lightgbm")
_ensure_package("shap")
_ensure_package("statsmodels")

import numpy as np
import pandas as pd
from sklearn.ensemble import (GradientBoostingClassifier, RandomForestClassifier,
                               StackingClassifier)
from sklearn.linear_model import LogisticRegression
from sklearn.base import clone
from sklearn.preprocessing import StandardScaler
from sklearn.utils import resample
from sklearn.metrics import roc_auc_score, brier_score_loss
from sklearn.calibration import calibration_curve
from sklearn.model_selection import StratifiedKFold, cross_val_score
from scipy import stats
import xgboost as xgb
import lightgbm as lgb
warnings.filterwarnings("ignore")

import sklearn, xgboost, lightgbm
print("=== Library Versions ===")
print(f"  Python       {sys.version.split()[0]}")
print(f"  scikit-learn {sklearn.__version__}")
print(f"  XGBoost      {xgboost.__version__}")
print(f"  LightGBM     {lightgbm.__version__}")
print(f"  numpy        {np.__version__}")
print(f"  pandas       {pd.__version__}")

# ─── CONFIGURATION ────────────────────────────────────────────────────────────
RANDOM_SEED    = 42
N_BOOTSTRAP    = 200          # PSU-level bootstrap resamples for subgroup median CIs
N_GINI_BOOT    = 1000         # [REVISION] Bootstrap resamples for Gini CI (R2 Comment 3c)
IR_FILE        = "BDIR81FL.SAV"   # SPSS .SAV format (full BDHS 2022 IR file)
SUBGROUP_MODEL = "Gradient Boosting"
np.random.seed(RANDOM_SEED)

# ─── DATA LOADING ─────────────────────────────────────────────────────────────
print(f"\nLoading BDHS 2022 IR file: {IR_FILE}")
base_cols   = ["V190","V025","V024","V106","V012","V136","V312","V501"]
extra_cols  = ["V201","V714","V158","V632"]
survey_cols = base_cols + extra_cols + ["V001","V023"]

def load_ir_file(filepath, wanted_cols):
    """
    Load the BDHS IR file from either .SAV (SPSS) or .CSV format.
    Returns a DataFrame with only the requested columns that exist in the file.
    For .SAV files, pyreadstat is used; column names are uppercased to match
    DHS variable naming conventions.
    """
    ext = os.path.splitext(filepath)[1].upper()

    if ext == ".SAV":
        try:
            import pyreadstat
            print(f"  Reading SPSS .SAV file via pyreadstat...")
            # usecols parameter filters at read time for speed on large files
            wanted_upper = [c.upper() for c in wanted_cols]
            df, meta = pyreadstat.read_sav(
                filepath,
                usecols=wanted_upper,
                apply_value_formats=False,   # keep numeric codes, not labels
            )
            # Standardise column names to uppercase
            df.columns = [c.upper() for c in df.columns]
            # Convert all columns to numeric (pyreadstat may return float64 already)
            for col in df.columns:
                df[col] = pd.to_numeric(df[col], errors="coerce")
            return df
        except ImportError:
            raise ImportError(
                "pyreadstat is required to read .SAV files. "
                "Install with: pip install pyreadstat"
            )
        except Exception as e:
            # If usecols fails (column not found), retry loading all columns
            print(f"  Warning: selective column load failed ({e}). Loading all columns...")
            import pyreadstat
            df, meta = pyreadstat.read_sav(filepath, apply_value_formats=False)
            df.columns = [c.upper() for c in df.columns]
            for col in df.columns:
                df[col] = pd.to_numeric(df[col], errors="coerce")
            # Keep only requested columns that exist
            keep = [c for c in wanted_upper if c in df.columns]
            return df[keep]

    else:
        # CSV fallback (original behaviour)
        print(f"  Reading CSV file...")
        try:
            df = pd.read_csv(filepath, encoding="latin1", low_memory=False,
                             usecols=wanted_cols)
        except ValueError:
            try:
                df = pd.read_csv(filepath, encoding="latin1", low_memory=False,
                                 usecols=base_cols + extra_cols)
            except ValueError:
                df = pd.read_csv(filepath, encoding="latin1", low_memory=False,
                                 usecols=base_cols)
        return df

ir = load_ir_file(IR_FILE, survey_cols)
HAS_SURVEY = "V001" in ir.columns
print(f"  Loaded {len(ir):,} women. Columns: {list(ir.columns)}")
print(f"  Survey design variables available: {HAS_SURVEY}")

# Outcome: modern contraceptive use (DHS variable V312, standard codes)
MODERN_METHODS = [1, 2, 3, 5, 6, 10, 11, 14]
ir["modern_cp"] = ir["V312"].isin(MODERN_METHODS).astype(int)

# Restrict to currently married women in extreme wealth quintiles
ir_m = ir[ir["V501"] == 1].copy()
ir_e = ir_m[ir_m["V190"].isin([1, 5])].copy()
ir_e["wealth_rich"] = (ir_e["V190"] == 5).astype(int)

for col in extra_cols:
    if col not in ir_e.columns:
        ir_e[col] = np.nan

ir_e["parity"]      = np.clip(ir_e["V201"].fillna(ir_e["V201"].median()), 0, 10)
ir_e["employed"]    = (ir_e["V714"] == 1).astype(float)
ir_e["watches_tv"]  = (ir_e["V158"].fillna(0) > 0).astype(float)
ir_e["fp_autonomy"] = ir_e["V632"].isin([1, 2]).astype(float)

BASE_FEAT        = ["V012", "V025", "V106", "V136"]
EXTRA_FEAT_CANDS = ["parity", "employed", "watches_tv", "fp_autonomy"]

df = ir_e.dropna(subset=BASE_FEAT + ["modern_cp"]).copy().reset_index(drop=True)
for f in EXTRA_FEAT_CANDS:
    df[f] = df[f].fillna(df[f].median())

EXTRA_FEAT = [f for f in EXTRA_FEAT_CANDS if df[f].std() > 0.01]
FEATURES   = BASE_FEAT + EXTRA_FEAT
print(f"  Features ({len(FEATURES)}): {FEATURES}")

n_poor = int((df.wealth_rich == 0).sum())
n_rich = int((df.wealth_rich == 1).sum())
print(f"  n={len(df):,}  poor={n_poor:,}  rich={n_rich:,}")
print(f"  mCP: overall={df.modern_cp.mean():.3f}  "
      f"poor={df[df.wealth_rich==0].modern_cp.mean():.3f}  "
      f"rich={df[df.wealth_rich==1].modern_cp.mean():.3f}")

# ─── TABLE 1: SAMPLE CHARACTERISTICS ─────────────────────────────────────────
print("\n=== TABLE 1: SAMPLE CHARACTERISTICS ===")
rows_t1 = []
for wlabel, wval in [("Poorest quintile", 0), ("Richest quintile", 1)]:
    s = df[df.wealth_rich == wval]
    row = {
        "Wealth group":      wlabel,
        "N":                 len(s),
        "Mean age (SD)":     f"{s.V012.mean():.1f} ({s.V012.std():.1f})",
        "Urban %":           f"{(s.V025==1).mean()*100:.1f}%",
        "Edu: None %":       f"{(s.V106==0).mean()*100:.1f}%",
        "Edu: Primary %":    f"{(s.V106==1).mean()*100:.1f}%",
        "Edu: Secondary %":  f"{(s.V106==2).mean()*100:.1f}%",
        "Edu: Higher %":     f"{(s.V106==3).mean()*100:.1f}%",
        "Mean HH size (SD)": f"{s.V136.mean():.1f} ({s.V136.std():.1f})",
        "Modern CP %":       f"{s.modern_cp.mean()*100:.1f}%",
    }
    if "parity"      in EXTRA_FEAT: row["Mean parity (SD)"]       = f"{s.parity.mean():.1f} ({s.parity.std():.1f})"
    if "employed"    in EXTRA_FEAT: row["Employed %"]             = f"{s.employed.mean()*100:.1f}%"
    if "watches_tv"  in EXTRA_FEAT: row["Watches TV %"]           = f"{s.watches_tv.mean()*100:.1f}%"
    if "fp_autonomy" in EXTRA_FEAT: row["FP decision autonomy %"] = f"{s.fp_autonomy.mean()*100:.1f}%"
    rows_t1.append(row)
t1 = pd.DataFrame(rows_t1)
print(t1.to_string(index=False))
t1.to_csv("results_table1_sample.csv", index=False)

# ─── FEATURE MATRIX ───────────────────────────────────────────────────────────
X_raw  = df[FEATURES].values
y      = df["modern_cp"].values
A      = df["wealth_rich"].values
scaler = StandardScaler()
X      = scaler.fit_transform(X_raw)

# ─── MODEL DEFINITIONS ────────────────────────────────────────────────────────
lr_m = LogisticRegression(random_state=RANDOM_SEED, max_iter=1000)

rf_m = RandomForestClassifier(n_estimators=300, max_depth=6,
                               min_samples_leaf=20,
                               random_state=RANDOM_SEED, n_jobs=-1)

gb_m = GradientBoostingClassifier(n_estimators=200, learning_rate=0.05,
                                   max_depth=4, subsample=0.8,
                                   random_state=RANDOM_SEED)

xgb_m = xgb.XGBClassifier(n_estimators=200, learning_rate=0.05, max_depth=4,
                            subsample=0.8, colsample_bytree=0.8,
                            eval_metric="logloss", random_state=RANDOM_SEED,
                            verbosity=0)

lgb_m = lgb.LGBMClassifier(n_estimators=200, learning_rate=0.05, max_depth=4,
                             num_leaves=31, subsample=0.8, colsample_bytree=0.8,
                             random_state=RANDOM_SEED, verbose=-1)

stack_m = StackingClassifier(
    estimators=[
        ("lr",  LogisticRegression(random_state=RANDOM_SEED, max_iter=1000)),
        ("rf",  RandomForestClassifier(n_estimators=200, max_depth=5,
                                        random_state=RANDOM_SEED, n_jobs=-1)),
        ("gb",  GradientBoostingClassifier(n_estimators=150, learning_rate=0.05,
                                            max_depth=4, random_state=RANDOM_SEED)),
        ("xgb", xgb.XGBClassifier(n_estimators=150, learning_rate=0.05,
                                    max_depth=4, random_state=RANDOM_SEED,
                                    verbosity=0, eval_metric="logloss")),
        ("lgb", lgb.LGBMClassifier(n_estimators=150, learning_rate=0.05,
                                    max_depth=4, random_state=RANDOM_SEED,
                                    verbose=-1)),
    ],
    final_estimator=LogisticRegression(C=1.0, random_state=RANDOM_SEED),
    cv=5, n_jobs=-1
)

model_defs = {
    "Logistic Regression": lr_m,
    "Random Forest":        rf_m,
    "Gradient Boosting":    gb_m,
    "XGBoost":              xgb_m,
    "LightGBM":             lgb_m,
    "Stacking Ensemble":    stack_m,
}

# ─── SUPPLEMENTARY TABLE S1: HYPERPARAMETER SPECIFICATIONS ───────────────────
# [REVISION] Added for R1 Major Comment 5 (methodological transparency)
print("\n=== SUPPLEMENTARY TABLE S1: HYPERPARAMETER SPECIFICATIONS ===")
hyper_rows = [
    {
        "Model": "Logistic Regression",
        "Key hyperparameters": "penalty=l2, C=1.0, solver=lbfgs, max_iter=1000",
        "Class imbalance handling": "None (outcome prevalence 34.9%; no resampling applied)",
        "Preprocessing": "StandardScaler (zero mean, unit variance)",
        "Notes": "Parametric benchmark; near-constant individual ORs by design",
    },
    {
        "Model": "Random Forest",
        "Key hyperparameters": "n_estimators=300, max_depth=6, min_samples_leaf=20, n_jobs=-1",
        "Class imbalance handling": "None",
        "Preprocessing": "StandardScaler",
        "Notes": "Non-parametric; axis-aligned splits; honest bootstrap aggregation",
    },
    {
        "Model": "Gradient Boosting",
        "Key hyperparameters": "n_estimators=200, learning_rate=0.05, max_depth=4, subsample=0.8",
        "Class imbalance handling": "None",
        "Preprocessing": "StandardScaler",
        "Notes": "Selected for subgroup analysis; supports SHAP TreeExplainer exactly",
    },
    {
        "Model": "XGBoost",
        "Key hyperparameters": "n_estimators=200, learning_rate=0.05, max_depth=4, subsample=0.8, colsample_bytree=0.8, eval_metric=logloss",
        "Class imbalance handling": "None",
        "Preprocessing": "StandardScaler",
        "Notes": "Second-order gradient statistics; sparse-aware split finding",
    },
    {
        "Model": "LightGBM",
        "Key hyperparameters": "n_estimators=200, learning_rate=0.05, max_depth=4, num_leaves=31, subsample=0.8, colsample_bytree=0.8",
        "Class imbalance handling": "None",
        "Preprocessing": "StandardScaler",
        "Notes": "Leaf-wise tree growth; histogram-based splitting",
    },
    {
        "Model": "Stacking Ensemble",
        "Key hyperparameters": "base learners: LR+RF(200 trees)+GB(150 trees)+XGB(150)+LGB(150); meta-learner: LR(C=1.0); cv=5",
        "Class imbalance handling": "None",
        "Preprocessing": "StandardScaler (applied before stacking)",
        "Notes": "Out-of-fold predictions used for meta-learner training; no bootstrap CI due to computational cost",
    },
]
s1 = pd.DataFrame(hyper_rows)
print(s1[["Model","Key hyperparameters"]].to_string(index=False))
s1.to_csv("results_suppTableS1_hyperparameters.csv", index=False)
print("  NOTE: Outcome prevalence = 34.9%; no class-weight adjustment or resampling applied.")
print("  All models implemented in Python 3.12, scikit-learn 1.6.1, XGBoost 3.2.0, LightGBM 4.6.0.")

# ─── T-LEARNER G-COMPUTATION ──────────────────────────────────────────────────
def get_individual_ORs(model, X, y, A):
    """
    T-learner: fit separate outcome models per wealth stratum, then generate
    counterfactual predictions for every individual under both wealth conditions
    and compute individual-specific odds ratios (Eq. 1 in manuscript).
    Returns ORs, p1 (counterfactual prob under richest), p0 (under poorest).
    """
    m1 = clone(model); m0 = clone(model)
    m1.fit(X[A == 1], y[A == 1])
    m0.fit(X[A == 0], y[A == 0])
    p1 = np.clip(m1.predict_proba(X)[:, 1], 1e-6, 1 - 1e-6)
    p0 = np.clip(m0.predict_proba(X)[:, 1], 1e-6, 1 - 1e-6)
    ors = (p1 / (1 - p1)) / (p0 / (1 - p0))
    return ors, p1, p0

def get_cv_auc(model, X, y, cv=5):
    skf = StratifiedKFold(n_splits=cv, shuffle=True, random_state=RANDOM_SEED)
    sc  = cross_val_score(clone(model), X, y, cv=skf, scoring="roc_auc", n_jobs=-1)
    return sc.mean(), sc.std()

# ─── POINT ESTIMATES AND CV-AUC ───────────────────────────────────────────────
print("\n=== Computing point estimates and CV-AUC ===")
all_ORs        = {}
all_p1         = {}   # [REVISION] store counterfactual probabilities for ARD computation
all_p0         = {}   # [REVISION]
all_AUC        = {}
all_CV_AUC     = {}
all_CV_AUC_std = {}
all_fitted     = {}
all_tlearners  = {}

for name, model in model_defs.items():
    print(f"  {name}...")
    m1 = clone(model); m0 = clone(model)
    m1.fit(X[A == 1], y[A == 1])
    m0.fit(X[A == 0], y[A == 0])
    p1  = np.clip(m1.predict_proba(X)[:, 1], 1e-6, 1 - 1e-6)
    p0  = np.clip(m0.predict_proba(X)[:, 1], 1e-6, 1 - 1e-6)
    ors = (p1 / (1 - p1)) / (p0 / (1 - p0))
    all_tlearners[name] = (m1, m0)
    all_ORs[name]       = ors
    all_p1[name]        = p1   # [REVISION]
    all_p0[name]        = p0   # [REVISION]

    fitted        = clone(model); fitted.fit(X, y)
    auc           = roc_auc_score(y, fitted.predict_proba(X)[:, 1])
    all_fitted[name] = fitted
    all_AUC[name]    = auc

    cv_auc, cv_std       = get_cv_auc(model, X, y)
    all_CV_AUC[name]     = cv_auc
    all_CV_AUC_std[name] = cv_std

    t = np.clip(ors, 0.01, np.percentile(ors, 99))
    print(f"    AUC={auc:.4f}  CV-AUC={cv_auc:.4f}±{cv_std:.4f}  "
          f"Median OR={np.median(t):.3f}  SD={np.std(t):.3f}")

# Save raw individual OR arrays
for name, arr in all_ORs.items():
    fname = {"Logistic Regression": "lr", "Random Forest": "rf",
             "Gradient Boosting": "gb", "XGBoost": "xgb",
             "LightGBM": "lgb", "Stacking Ensemble": "stack"}[name]
    np.save(f"individual_ORs_{fname}.npy", arr)

BEST = max(all_CV_AUC, key=all_CV_AUC.get)
print(f"\n  Best model (CV-AUC): {BEST} ({all_CV_AUC[BEST]:.4f})")
print(f"  Subgroup model: {SUBGROUP_MODEL}")

# ─── CALIBRATION ──────────────────────────────────────────────────────────────
print("\n=== Calibration ===")
cal_res = {}
for name, fm in all_fitted.items():
    probs = fm.predict_proba(X)[:, 1]
    brier = brier_score_loss(y, probs)
    fp, mp = calibration_curve(y, probs, n_bins=10, strategy="uniform")
    cal_res[name] = {"brier": brier, "frac_pos": fp, "mean_pred": mp}
    print(f"  {name}: Brier={brier:.4f}")

# ─── PROPENSITY SCORE AND WEIGHT TRIMMING ─────────────────────────────────────
print("\n=== Propensity Score ===")
ps_m = LogisticRegression(random_state=RANDOM_SEED, max_iter=1000)
ps_m.fit(X, A)
ps_raw   = np.clip(ps_m.predict_proba(X)[:, 1], 1e-3, 1 - 1e-3)
ovlp_pct = ((ps_raw > 0.1) & (ps_raw < 0.9)).mean() * 100

p_t     = A.mean(); p_c = 1 - p_t
sw_raw  = np.where(A == 1, p_t / ps_raw, p_c / (1 - ps_raw))
max_raw = sw_raw.max()
ess_raw = sw_raw.sum() ** 2 / (sw_raw ** 2).sum()
print(f"  Overlap: {ovlp_pct:.1f}%  Untrimmed: max={max_raw:.2f} "
      f"ESS={ess_raw:.0f} ({ess_raw/len(A)*100:.1f}%)")

# 99th-percentile weight trimming (Cole & Hernán, 2008)
TRIM_CAP = np.percentile(sw_raw, 99)
ps_lo    = max(p_t / TRIM_CAP, 1e-3)
ps_hi    = min(1 - p_c / TRIM_CAP, 1 - 1e-3)
ps       = np.clip(ps_raw, ps_lo, ps_hi)
sw       = np.where(A == 1, p_t / ps, p_c / (1 - ps))
ess_t    = sw[A == 1].sum() ** 2 / (sw[A == 1] ** 2).sum()
ess_c    = sw[A == 0].sum() ** 2 / (sw[A == 0] ** 2).sum()
ess      = ess_t + ess_c
ess_pct  = ess / len(A) * 100
n_tr     = int(((ps_raw < ps_lo) | (ps_raw > ps_hi)).sum())
print(f"  Trimmed: [{ps_lo:.4f},{ps_hi:.4f}]  n_tr={n_tr} ({n_tr/len(A)*100:.1f}%)  "
      f"max={sw.max():.2f}  ESS={ess:.0f} ({ess_pct:.1f}%)")

pd.DataFrame([{
    "max_sw_untrimmed":  round(float(max_raw), 3),
    "ESS_untrimmed":     round(float(ess_raw), 1),
    "ESS_pct_untrimmed": round(ess_raw / len(A) * 100, 1),
    "ps_lo":             round(float(ps_lo), 4),
    "ps_hi":             round(float(ps_hi), 4),
    "n_trimmed":         n_tr,
    "pct_trimmed":       round(n_tr / len(A) * 100, 2),
    "max_sw_trimmed":    round(float(sw.max()), 3),
    "ESS_trimmed":       round(float(ess), 1),
    "ESS_pct_trimmed":   round(float(ess_pct), 1),
    "n_total":           int(len(A)),
    "overlap_pct_01_09": round(float(ovlp_pct), 1),
}]).to_csv("results_weight_diagnostics.csv", index=False)

# ─── AIPW DOUBLY-ROBUST SENSITIVITY ──────────────────────────────────────────
print("\n=== AIPW Sensitivity ===")

def aipw_ate(m1, m0, X, y, A, ps):
    """
    Augmented inverse probability weighted estimator (Eq. 2 in manuscript).
    Doubly robust: consistent if either outcome model or PS model is correct.
    """
    mu1 = m1.predict_proba(X)[:, 1]
    mu0 = m0.predict_proba(X)[:, 1]
    psi = (mu1 + A * (y - mu1) / ps) - (mu0 + (1 - A) * (y - mu0) / (1 - ps))
    ate    = np.mean(psi)
    se     = np.std(psi) / np.sqrt(len(psi))
    ci_lo  = ate - 1.96 * se
    ci_hi  = ate + 1.96 * se
    pval   = 2 * (1 - stats.norm.cdf(abs(ate / se)))
    p0m    = np.clip(np.mean(mu0), 1e-6, 1 - 1e-6)
    p1m    = np.clip(p0m + ate, 1e-6, 1 - 1e-6)
    or_est = (p1m / (1 - p1m)) / (p0m / (1 - p0m))
    return ate, se, ci_lo, ci_hi, pval, or_est

rows_aipw = []
for name in model_defs:
    try:
        m1, m0 = all_tlearners[name]
        ate, se, cil, cih, pv, ore = aipw_ate(m1, m0, X, y, A, ps)
        rows_aipw.append({
            "Model":          name,
            "ATE (RD)":       f"{ate:.4f}",
            "SE":             f"{se:.4f}",
            "95% CI (RD)":    f"{cil:.4f} to {cih:.4f}",
            "p-value":        f"{pv:.4f}" if pv >= 0.001 else "<0.001",
            "Approx OR":      f"{ore:.3f}",
            "PS trim range":  f"[{ps_lo:.4f},{ps_hi:.4f}]",
            "ESS after trim": f"{ess:.0f} ({ess_pct:.1f}%)",
        })
        print(f"  {name}: ATE={ate:.4f}  OR≈{ore:.3f}")
    except Exception as e:
        print(f"  {name}: AIPW failed — {e}")
        rows_aipw.append({"Model": name, "ATE (RD)": "N/A", "SE": "N/A",
                           "95% CI (RD)": "N/A", "p-value": "N/A",
                           "Approx OR": "N/A", "PS trim range": "N/A",
                           "ESS after trim": "N/A"})

pd.DataFrame(rows_aipw).to_csv("results_table2_aipw_detail.csv", index=False)

# ─── [REVISION] AIPW MISSPECIFICATION ROBUSTNESS: RF-BASED PS ────────────────
# R2 Comment 3b: demonstrate AIPW robustness under alternative PS model
print("\n=== [REVISION] AIPW robustness: RF propensity score (R2 Comment 3b) ===")
try:
    ps_rf_m = RandomForestClassifier(n_estimators=200, max_depth=5,
                                      min_samples_leaf=20,
                                      random_state=RANDOM_SEED, n_jobs=-1)
    ps_rf_m.fit(X, A)
    ps_rf_raw = np.clip(ps_rf_m.predict_proba(X)[:, 1], 1e-3, 1 - 1e-3)
    sw_rf_raw = np.where(A == 1, p_t / ps_rf_raw, p_c / (1 - ps_rf_raw))
    TRIM_CAP_RF = np.percentile(sw_rf_raw, 99)
    ps_rf_lo = max(p_t / TRIM_CAP_RF, 1e-3)
    ps_rf_hi = min(1 - p_c / TRIM_CAP_RF, 1 - 1e-3)
    ps_rf = np.clip(ps_rf_raw, ps_rf_lo, ps_rf_hi)

    rows_aipw_rf = []
    for name in model_defs:
        try:
            m1, m0 = all_tlearners[name]
            ate_rf, se_rf, cil_rf, cih_rf, pv_rf, ore_rf = aipw_ate(
                m1, m0, X, y, A, ps_rf)
            rows_aipw_rf.append({
                "Model":       name,
                "PS model":    "Random Forest",
                "ATE (RD)":    f"{ate_rf:.4f}",
                "SE":          f"{se_rf:.4f}",
                "95% CI (RD)": f"{cil_rf:.4f} to {cih_rf:.4f}",
                "p-value":     f"{pv_rf:.4f}" if pv_rf >= 0.001 else "<0.001",
                "Approx OR":   f"{ore_rf:.3f}",
            })
            print(f"  RF-PS | {name}: ATE={ate_rf:.4f}  OR≈{ore_rf:.3f}  p={pv_rf:.3f}")
        except Exception as e:
            print(f"  RF-PS | {name}: failed — {e}")
    if rows_aipw_rf:
        pd.DataFrame(rows_aipw_rf).to_csv(
            "results_suppTableS2_aipw_rfps_robustness.csv", index=False)
        print("  Saved: results_suppTableS2_aipw_rfps_robustness.csv")
        print("  NOTE: Near-null AIPW findings are robust to PS model specification,")
        print("        consistent with the doubly-robust property of the AIPW estimator.")
except Exception as e:
    print(f"  RF-PS robustness check failed: {e}")

# ─── BOOTSTRAP CONFIDENCE INTERVALS ──────────────────────────────────────────
print(f"\n=== Bootstrap ({N_BOOTSTRAP} iterations) ===")

# PSU-level resampling preserves the BDHS complex two-stage cluster design
if HAS_SURVEY and "V001" in df.columns:
    psus  = df["V001"].values
    upsus = np.unique(psus)
    def bsamp(seed):
        rng = np.random.RandomState(seed)
        sel = rng.choice(upsus, size=len(upsus), replace=True)
        return np.concatenate([np.where(psus == p)[0] for p in sel])
    blabel = "PSU-level survey bootstrap"
else:
    warnings.warn(
        "PSU variable V001 not found — falling back to individual-level bootstrap. "
        "Ensure V001 is present for manuscript-consistent CIs.",
        UserWarning, stacklevel=2)
    def bsamp(seed):
        return resample(np.arange(len(y)), random_state=seed)
    blabel = "Individual-level bootstrap (fallback — V001 not available)"

boot_meds = {n: [] for n in model_defs}
fast_ms   = {k: v for k, v in model_defs.items() if k != "Stacking Ensemble"}

for i in range(N_BOOTSTRAP):
    if i % 50 == 0: print(f"  {i}/{N_BOOTSTRAP}")
    idx = bsamp(i)
    Xb, yb, Ab = X[idx], y[idx], A[idx]
    if Ab.sum() < 10 or (Ab == 0).sum() < 10:
        continue
    for name, model in fast_ms.items():
        try:
            m1b = clone(model); m0b = clone(model)
            m1b.fit(Xb[Ab == 1], yb[Ab == 1])
            m0b.fit(Xb[Ab == 0], yb[Ab == 0])
            p1b  = np.clip(m1b.predict_proba(Xb)[:, 1], 1e-6, 1 - 1e-6)
            p0b  = np.clip(m0b.predict_proba(Xb)[:, 1], 1e-6, 1 - 1e-6)
            ors_b = (p1b / (1 - p1b)) / (p0b / (1 - p0b))
            boot_meds[name].append(
                np.median(np.clip(ors_b, 0.01, np.percentile(ors_b, 99))))
        except:
            pass

boot_meds["Stacking Ensemble"] = []
print(f"  Complete. Method: {blabel}")
print("  NOTE: Bootstrap CIs quantify sampling uncertainty in subgroup medians only,")
print("        not individual-level CATE precision. See manuscript Section 2.5.")

# ─── TABLE 2: MODEL PERFORMANCE ───────────────────────────────────────────────
print("\n=== TABLE 2: MODEL PERFORMANCE ===")
rows_t2 = []
for name, ors in all_ORs.items():
    t  = np.clip(ors, 0.01, np.percentile(ors, 99))
    bm = boot_meds[name]
    ci = (f"{np.percentile(bm,2.5):.3f}–{np.percentile(bm,97.5):.3f}"
          if bm else "N/A (not bootstrapped)")
    ar = next((r for r in rows_aipw if r["Model"] == name), {})
    rows_t2.append({
        "Model":           name,
        "AUC (full)":      f"{all_AUC[name]:.4f}",
        "CV-AUC (5-fold)": f"{all_CV_AUC[name]:.4f}±{all_CV_AUC_std[name]:.4f}",
        "Brier Score":     f"{cal_res[name]['brier']:.4f}",
        "Median OR":       f"{np.median(t):.3f}",
        "95% CI (OR)":     ci,
        "AIPW ATE":        ar.get("ATE (RD)", "N/A"),
        "Bootstrap":       blabel,
    })
t2 = pd.DataFrame(rows_t2)
print(t2.to_string(index=False))
t2.to_csv("results_table2_model_performance.csv", index=False)

# ─── SUBGROUP LABELS ──────────────────────────────────────────────────────────
subgroup_ors      = all_ORs[SUBGROUP_MODEL]
subgroup_p1       = all_p1[SUBGROUP_MODEL]   # [REVISION] counterfactual probs
subgroup_p0       = all_p0[SUBGROUP_MODEL]   # [REVISION]

df["ind_OR"]      = subgroup_ors
df["ind_OR_trim"] = np.clip(subgroup_ors, 0.01, np.percentile(subgroup_ors, 99))
df["ind_ARD"]     = subgroup_p1 - subgroup_p0   # [REVISION] individual absolute risk difference
df["age_group"]   = pd.cut(df["V012"], bins=[14, 24, 34, 49],
                            labels=["15-24", "25-34", "35-49"])
df["urban_label"] = df["V025"].map({1: "Urban", 2: "Rural"})
df["edu_label"]   = df["V106"].map({0: "No education", 1: "Primary",
                                     2: "Secondary",    3: "Higher"})

# ─── SHAP FEATURE IMPORTANCE ──────────────────────────────────────────────────
# NOTE: SHAP values decompose variation in predicted probability (not causal effects).
# They quantify predictive contributions to the outcome model; interpretation as
# causal importance requires caution. See manuscript Section 4.1.
print("\n=== SHAP Analysis ===")
print("  NOTE: SHAP values reflect predictive contributions, not causal effects.")
shap_rows = []; shap_vals_arr = None
try:
    import shap
    shap_model = all_fitted[SUBGROUP_MODEL]
    print(f"  Computing SHAP values on pooled {SUBGROUP_MODEL} model")
    exp = shap.TreeExplainer(shap_model)
    sv  = exp.shap_values(X)
    if isinstance(sv, list):
        sv = sv[1]
    shap_vals_arr = sv
    np.save("shap_values.npy", sv)

    FEAT_LABELS = {
        "V012":      "Age",
        "V025":      "Urban/Rural",
        "V106":      "Education",
        "V136":      "Household size",
        "parity":    "Parity (no. children)",
        "employed":  "Currently employed",
        "watches_tv":"Media exposure (TV)",
        "fp_autonomy":"FP decision autonomy",
    }
    display  = [FEAT_LABELS.get(f, f) for f in FEATURES]
    mean_abs = np.abs(sv).mean(axis=0)
    shap_df  = (pd.DataFrame({"Feature": display, "Variable": FEATURES,
                               "Mean |SHAP|": mean_abs})
                .sort_values("Mean |SHAP|", ascending=False)
                .reset_index(drop=True))
    shap_df["Rank"] = range(1, len(shap_df) + 1)
    print(shap_df[["Rank", "Feature", "Mean |SHAP|"]].to_string(index=False))
    shap_df.to_csv("results_table9_shap_importance.csv", index=False)
    shap_rows = shap_df.to_dict("records")
except Exception as e:
    print(f"  SHAP failed: {e}")

# ─── E-VALUE ANALYSIS ─────────────────────────────────────────────────────────
print("\n=== E-value Analysis ===")

def ev(OR):
    """E-value formula from VanderWeele & Ding (2017), Ann Intern Med."""
    if OR < 1:
        OR = 1 / OR
    return OR + np.sqrt(OR * (OR - 1)) if OR > 1 else 1.0

rows_ev = []
for age in ["15-24", "25-34", "35-49"]:
    for edu in ["No education", "Primary", "Secondary", "Higher"]:
        sub = df[(df.age_group == age) & (df.edu_label == edu)]
        if len(sub) < 30:
            continue
        mor = sub.ind_OR_trim.median()
        q25 = sub.ind_OR_trim.quantile(0.25)
        rows_ev.append({
            "Age Group":          age,
            "Education":          edu,
            "N":                  len(sub),
            "Median OR":          f"{mor:.3f}",
            "E-value":            f"{ev(mor):.3f}",
            "E-value (CI limit)": f"{ev(q25):.3f}",
            "Interpretation":     (
                "Strong — robust to confounding"         if ev(mor) > 3   else
                "Moderate — moderate confounder needed"  if ev(mor) > 1.5 else
                "Weak — small confounder could explain"),
        })
t6_eval = pd.DataFrame(rows_ev)
print(t6_eval.to_string(index=False))
t6_eval.to_csv("results_table6_evalue.csv", index=False)

# ─── INTERACTION TESTS ────────────────────────────────────────────────────────
print("\n=== Interaction Tests ===")
t5_int = pd.DataFrame()
try:
    import statsmodels.api as sm

    def manual_lr_interaction(df_in, outcome, exposure, moderator_col):
        dummies = pd.get_dummies(df_in[moderator_col], drop_first=True, prefix="mod")
        Xdf = pd.concat([
            df_in[[outcome, exposure]].reset_index(drop=True),
            dummies.reset_index(drop=True)
        ], axis=1).dropna()
        y_in  = Xdf[outcome].values
        Xm    = sm.add_constant(Xdf.drop(columns=[outcome]).values.astype(float))
        X_int = Xm.copy()
        n_int = 0
        for col in dummies.columns:
            ic    = Xdf[exposure].values * Xdf[col].values
            X_int = np.column_stack([X_int, ic])
            n_int += 1
        mod_main = sm.Logit(y_in, Xm).fit(disp=False, maxiter=200)
        mod_int  = sm.Logit(y_in, X_int).fit(disp=False, maxiter=200)
        lr_stat  = -2 * (mod_main.llf - mod_int.llf)
        lr_p     = 1 - stats.chi2.cdf(lr_stat, df=n_int)
        return lr_stat, n_int, lr_p

    rows_it = []
    di = df.copy(); di["age_group"] = di["age_group"].astype(str)
    for lbl, exp_col, mod_col in [
        ("Wealth × Education", "wealth_rich", "V106"),
        ("Wealth × Age group", "wealth_rich", "age_group"),
        ("Wealth × Residence", "wealth_rich", "V025"),
    ]:
        try:
            lr, ld, lp = manual_lr_interaction(di, "modern_cp", exp_col, mod_col)
            rows_it.append({
                "Test":                 lbl,
                "N interaction terms":  ld,
                "LR statistic":         f"{lr:.3f}",
                "df":                   ld,
                "p-value (LR)":         f"{lp:.4f}" if lp >= 0.001 else "<0.001",
                "Significant (p<0.05)": "Yes" if lp < 0.05 else "No",
            })
            print(f"  {lbl}: LR={lr:.2f}  p={lp:.4f}")
        except Exception as e:
            print(f"  {lbl}: failed — {e}")
    if rows_it:
        t5_int = pd.DataFrame(rows_it)
        t5_int.to_csv("results_table5_interaction.csv", index=False)
except ImportError:
    print("  statsmodels not available")

# ─── BH-FDR SUBGROUP ANALYSES ─────────────────────────────────────────────────
print("\n=== Subgroup FDR ===")

def chi2_vs_all(s, mu):
    n = len(s); o = s.modern_cp.sum(); e = mu * n
    if e < 5 or (n - e) < 5:
        return float("nan")
    return 1 - stats.chi2.cdf((o - e) ** 2 / e + ((n - o) - (n - e)) ** 2 / (n - e), df=1)

def bh_fdr(pv):
    """Benjamini-Hochberg FDR correction."""
    p   = np.array(pv, dtype=float); n = len(p)
    ranks     = np.argsort(p)
    inv_ranks = np.argsort(ranks) + 1
    adj = np.minimum(1.0, p * n / inv_ranks)
    for i in range(n - 2, -1, -1):
        adj[ranks[i]] = min(adj[ranks[i]], adj[ranks[i + 1]])
    return adj

mu = df.modern_cp.mean()

# ─── TABLE 3: AGE × EDUCATION (with ARD column) ───────────────────────────────
# [REVISION] Added Median ARD column (R1 Major Comment 3; R2 Comment 5)
# ARD = µ̂₁(Xᵢ) − µ̂₀(Xᵢ): predicted absolute risk difference under richest vs poorest
print("\n=== TABLE 3: Age × Education subgroup ORs (with ARD) ===")
rows_t3, rp3 = [], []
for age in ["15-24", "25-34", "35-49"]:
    for edu in ["No education", "Primary", "Secondary", "Higher"]:
        sub = df[(df.age_group == age) & (df.edu_label == edu)]
        if len(sub) < 30:
            continue
        ors  = sub.ind_OR_trim
        ards = sub.ind_ARD          # [REVISION]
        p    = chi2_vs_all(sub, mu)
        rp3.append(p)
        rows_t3.append({
            "Age Group":      age,
            "Education":      edu,
            "N":              len(sub),
            "% Urban":        f"{(sub.urban_label=='Urban').mean()*100:.1f}%",
            "mCP %":          f"{sub.modern_cp.mean()*100:.1f}%",
            "Median OR":      f"{ors.median():.3f}",
            "IQR (OR)":       f"{ors.quantile(.25):.3f}–{ors.quantile(.75):.3f}",
            "Median ARD (pp)":f"{ards.median()*100:+.1f}",   # [REVISION] percentage points
            "IQR (ARD pp)":   f"{ards.quantile(.25)*100:+.1f} to {ards.quantile(.75)*100:+.1f}",  # [REVISION]
            "Mean OR":        f"{ors.mean():.3f}",
            "p-value":        f"{p:.4f}" if not np.isnan(p) else "N/A",
            "FDR q-value":    "pending",
        })
vi = [i for i, p in enumerate(rp3) if not np.isnan(p)]
if vi:
    ap = bh_fdr([rp3[i] for i in vi])
    for rank, i in enumerate(vi):
        rows_t3[i]["FDR q-value"] = f"{ap[rank]:.4f}"
t3 = pd.DataFrame(rows_t3)
print(t3.to_string(index=False))
t3.to_csv("results_table3_age_education.csv", index=False)
print("  ARD = predicted absolute risk difference (richest minus poorest); pp = percentage points.")
print("  Positive ARD: wealth associated with higher mCP prevalence.")
print("  Negative ARD: wealth associated with lower mCP prevalence (equity-achieved subgroups).")

# ─── TABLE 4: RESIDENCE × EDUCATION ──────────────────────────────────────────
rows_t4, rp4 = [], []
for res in ["Urban", "Rural"]:
    for edu in ["No education", "Primary", "Secondary", "Higher"]:
        sub = df[(df.urban_label == res) & (df.edu_label == edu)]
        if len(sub) < 30:
            continue
        ors = sub.ind_OR_trim
        p   = chi2_vs_all(sub, mu)
        rp4.append(p)
        rows_t4.append({
            "Residence":   res,
            "Education":   edu,
            "N":           len(sub),
            "mCP %":       f"{sub.modern_cp.mean()*100:.1f}%",
            "Median OR":   f"{ors.median():.3f}",
            "IQR":         f"{ors.quantile(.25):.3f}–{ors.quantile(.75):.3f}",
            "Mean OR":     f"{ors.mean():.3f}",
            "p-value":     f"{p:.4f}" if not np.isnan(p) else "N/A",
            "FDR q-value": "pending",
        })
vi4 = [i for i, p in enumerate(rp4) if not np.isnan(p)]
if vi4:
    ap4 = bh_fdr([rp4[i] for i in vi4])
    for rank, i in enumerate(vi4):
        rows_t4[i]["FDR q-value"] = f"{ap4[rank]:.4f}"
t4 = pd.DataFrame(rows_t4)
print(t4.to_string(index=False))
t4.to_csv("results_table4_residence_education.csv", index=False)

# ─── GINI COEFFICIENT ─────────────────────────────────────────────────────────
print("\n=== Gini Coefficient of |log OR| ===")

def gini_standard(x):
    """Standard Gini coefficient computed on absolute values."""
    x = np.abs(np.array(x, dtype=float))
    x = np.sort(x); n = len(x)
    if x.sum() == 0:
        return 0.0
    ranks = np.arange(1, n + 1)
    return (2 * np.dot(ranks, x) - (n + 1) * x.sum()) / (n * x.sum())

ot  = df.ind_OR_trim.values
lo  = np.log(np.clip(ot, 0.01, None))
gv  = gini_standard(np.abs(lo))
pa  = (ot > 1).mean() * 100
pb  = (ot < 1).mean() * 100
pl  = (np.abs(lo) > np.log(2)).mean() * 100
sa  = np.sort(np.abs(lo))[::-1]
t10 = sa[:int(0.1 * len(sa))].sum() / sa.sum() * 100
print(f"  Gini={gv:.4f}  OR>1: {pa:.1f}%  OR<1: {pb:.1f}%  >2x: {pl:.1f}%  top10%: {t10:.1f}%")

# ─── [REVISION] BOOTSTRAP CI FOR GINI COEFFICIENT ────────────────────────────
# R2 Comment 3c: provide confidence interval for Gini estimate
print(f"\n  [REVISION] Computing bootstrap CI for Gini ({N_GINI_BOOT} resamples)...")
gini_boot = []
rng_gini = np.random.RandomState(RANDOM_SEED)
for _ in range(N_GINI_BOOT):
    idx_b = rng_gini.choice(len(lo), size=len(lo), replace=True)
    lo_b  = lo[idx_b]
    gini_boot.append(gini_standard(np.abs(lo_b)))
gini_boot = np.array(gini_boot)
gini_ci_lo = np.percentile(gini_boot, 2.5)
gini_ci_hi = np.percentile(gini_boot, 97.5)
print(f"  Gini = {gv:.4f}  95% bootstrap CI: [{gini_ci_lo:.4f}, {gini_ci_hi:.4f}]")
print(f"  (Individual-level bootstrap, n={N_GINI_BOOT} resamples)")

pd.DataFrame([{
    "gini_abs_logOR":        round(float(gv), 4),
    "gini_CI_lo_95pct":      round(float(gini_ci_lo), 4),   # [REVISION]
    "gini_CI_hi_95pct":      round(float(gini_ci_hi), 4),   # [REVISION]
    "gini_boot_n":           N_GINI_BOOT,                    # [REVISION]
    "pct_OR_greater2x":      round(float(pl), 1),
    "top10pct_logOR_share":  round(float(t10), 1),
    "pct_OR_above1":         round(float(pa), 1),
    "pct_OR_below1":         round(float(pb), 1),
    "median_OR_above1":      round(float(np.median(ot[ot > 1])), 3),
    "median_OR_below1":      round(float(np.median(ot[ot < 1])), 3),
}]).to_csv("results_gini_heterogeneity.csv", index=False)

# ─── [REVISION] 95th-PERCENTILE TRIMMING SENSITIVITY FOR TABLE 3 ─────────────
# R2 Comment 3a: sensitivity of subgroup findings to trimming threshold
print("\n=== [REVISION] Table 3 sensitivity: 95th-percentile OR trimming (R2 Comment 3a) ===")
subgroup_ors_95 = np.clip(subgroup_ors, 0.01, np.percentile(subgroup_ors, 95))
df["ind_OR_trim95"] = subgroup_ors_95

rows_t3_sens = []
for age in ["15-24", "25-34", "35-49"]:
    for edu in ["No education", "Primary", "Secondary", "Higher"]:
        sub = df[(df.age_group == age) & (df.edu_label == edu)]
        if len(sub) < 30:
            continue
        ors99 = sub.ind_OR_trim      # primary (99th pct)
        ors95 = sub.ind_OR_trim95    # sensitivity (95th pct)
        rows_t3_sens.append({
            "Age Group":             age,
            "Education":             edu,
            "N":                     len(sub),
            "Median OR (99th pct)":  f"{ors99.median():.3f}",
            "Median OR (95th pct)":  f"{ors95.median():.3f}",
            "Shift":                 f"{abs(ors99.median()-ors95.median()):.3f}",
            "Substantive change?":   "Yes" if abs(ors99.median()-ors95.median()) >= 0.05 else "No",
        })
t3_sens = pd.DataFrame(rows_t3_sens)
print(t3_sens.to_string(index=False))
t3_sens.to_csv("results_suppTableS3_trimming95_sensitivity.csv", index=False)
max_shift = max(float(r["Shift"]) for r in rows_t3_sens)
n_changed = sum(1 for r in rows_t3_sens if r["Substantive change?"] == "Yes")
print(f"  Max OR shift across cells: {max_shift:.3f}  |  Cells with shift ≥0.05: {n_changed}")
print("  Conclusion: subgroup OR pattern is robust to trimming threshold choice.")

# ─── TABLE 7: PROGRAMMATIC TARGETING ─────────────────────────────────────────
# [REVISION] Thresholds are post-hoc decision-analytic cutpoints; see footnote
rows_t8 = []
for age in ["15-24", "25-34", "35-49"]:
    for edu in ["No education", "Primary", "Secondary", "Higher"]:
        sub = df[(df.age_group == age) & (df.edu_label == edu)]
        if len(sub) < 30:
            continue
        mor = sub.ind_OR_trim.median()
        if   mor >= 3.0: pri = "HIGH — private-sector cost barrier";      act = "Subsidised LARC / voucher"
        elif mor >= 1.5: pri = "MODERATE — emerging wealth barrier";       act = "Targeted demand generation"
        elif mor >= 0.8: pri = "LOW — near-equitable";                     act = "Monitor; maintain public supply"
        else:            pri = "EQUITY ACHIEVED — public system effective"; act = "Protect community programmes"
        rows_t8.append({
            "Age group":          age,
            "Education":          edu,
            "N":                  len(sub),
            "mCP %":              f"{sub.modern_cp.mean()*100:.1f}%",
            "% Urban":            f"{(sub.urban_label=='Urban').mean()*100:.1f}%",
            "Median OR":          f"{mor:.3f}",
            "Priority":           pri,
            "Recommended action": act,
        })
t8 = pd.DataFrame(rows_t8)
t8.to_csv("results_table8_policy_targeting.csv", index=False)
# [REVISION] footnote printed to log
print("\n  TABLE 7 FOOTNOTE (R2 Comment 6):")
print("  Tier thresholds (HIGH ≥3.0, MODERATE 1.5–2.9, LOW 0.8–1.49, EQUITY ACHIEVED <0.8)")
print("  are post-hoc decision-analytic cutpoints selected to align with natural breaks")
print("  in the observed OR distribution; they are not statistically derived and should")
print("  be adapted by programme planners to local context and resource constraints.")

# ─── FIGURES ──────────────────────────────────────────────────────────────────
print("\nGenerating figures...")
try:
    import matplotlib; matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    import matplotlib.patches as mpatches
    from scipy.stats import gaussian_kde

    COLORS = {
        "Logistic Regression": "#2980B9",
        "Random Forest":       "#27AE60",
        "Gradient Boosting":   "#E74C3C",
        "XGBoost":             "#8E44AD",
        "LightGBM":            "#F39C12",
        "Stacking Ensemble":   "#2C3E50",
    }

    # Figure 1: DAG
    fig, ax = plt.subplots(figsize=(12, 8))
    ax.set_xlim(0, 12); ax.set_ylim(0, 8)
    ax.axis("off"); fig.patch.set_facecolor("white")

    nodes = {
        "W":   (5.0, 6.8, "Household Wealth\n(Exposure)",       "#1A5276", "#D6EAF8", False),
        "C":   (9.5, 6.8, "Modern CP\n(Outcome)",                "#1A7A3C", "#D5F5E3", False),
        "Age": (1.2, 5.0, "Age",                                  "#7D6608", "#FEF9E7", False),
        "Ed":  (3.2, 5.0, "Education",                            "#7D6608", "#FEF9E7", False),
        "Ur":  (5.2, 5.0, "Urban/Rural",                          "#7D6608", "#FEF9E7", False),
        "Pa":  (7.2, 5.0, "Parity",                               "#7D6608", "#FEF9E7", False),
        "Me":  (9.2, 5.0, "Media\nExposure",                      "#7D6608", "#FEF9E7", False),
        "Em":  (2.2, 3.2, "Employment",                           "#7D6608", "#FEF9E7", False),
        "Au":  (6.2, 3.2, "FP Decision\nAutonomy",                "#7D6608", "#FEF9E7", False),
        "HH":  (4.2, 1.8, "HH Size",                              "#6C3483", "#F4ECF7", False),
        "PP":  (1.5, 0.8, "Partner Pref.\n(Unmeasured)",          "#C0392B", "#FADBD8", True),
        "MA":  (10.5, 3.2, "Method Avail.\n(Unmeasured)",         "#C0392B", "#FADBD8", True),
    }

    for k, (x, yp, lb, ec, fc, ds) in nodes.items():
        ax.add_patch(mpatches.FancyBboxPatch(
            (x - 1.0, yp - 0.5), 2.0, 1.0,
            boxstyle="round,pad=0.08", linewidth=2 if not ds else 1.5,
            edgecolor=ec, facecolor=fc, linestyle="--" if ds else "-", zorder=3))
        ax.text(x, yp, lb, ha="center", va="center", fontsize=7.5,
                fontweight="bold" if not ds else "normal", color=ec, zorder=4)

    def arr(s, d, col, lw=1.4, ls="-"):
        x1, y1 = nodes[s][0], nodes[s][1]
        x2, y2 = nodes[d][0], nodes[d][1]
        ax.annotate("", xy=(x2, y2), xytext=(x1, y1), zorder=2,
                    arrowprops=dict(arrowstyle="->", color=col, lw=lw,
                                    connectionstyle="arc3,rad=0.05", linestyle=ls))

    arr("W", "C", "#1A5276", lw=2.8)
    for s in ["Age", "Ed", "Ur", "Pa", "Me", "Em", "Au"]:
        arr(s, "W", "#B7950B"); arr(s, "C", "#B7950B")
    arr("HH", "W", "#6C3483"); arr("HH", "C", "#6C3483")
    arr("PP", "C", "#C0392B", ls="--"); arr("MA", "C", "#C0392B", ls="--")
    arr("Age", "Ed", "#CCC", lw=1.0); arr("Age", "HH", "#CCC", lw=1.0)
    arr("Age", "Pa", "#CCC", lw=1.0)
    arr("Ed", "Em", "#CCC", lw=1.0); arr("Ed", "Au", "#CCC", lw=1.0)

    ax.legend(handles=[
        mpatches.Patch(facecolor="#D6EAF8", edgecolor="#1A5276", lw=1.5, label="Exposure"),
        mpatches.Patch(facecolor="#D5F5E3", edgecolor="#1A7A3C", lw=1.5, label="Outcome"),
        mpatches.Patch(facecolor="#FEF9E7", edgecolor="#7D6608", lw=1.5, label="Measured confounder"),
        mpatches.Patch(facecolor="#F4ECF7", edgecolor="#6C3483", lw=1.5, label="Mediator/confounder"),
        mpatches.Patch(facecolor="#FADBD8", edgecolor="#C0392B", lw=1.5, ls="--", label="Unmeasured"),
    ], loc="lower right", fontsize=8, framealpha=0.95, bbox_to_anchor=(1.0, 0.01))
    ax.set_title(
        "Figure 1. DAG: Causal assumptions for wealth → modern contraceptive use\n"
        "BDHS 2022 (n=11,510; 8-covariate model)",
        fontsize=11, fontweight="bold", pad=12)
    plt.tight_layout()
    plt.savefig("fig1_dag.png", dpi=200, bbox_inches="tight", facecolor="white")
    plt.close(); print("  fig1_dag.png saved.")

    # Figure 2: SHAP feature importance
    if shap_vals_arr is not None and len(shap_rows) > 0:
        try:
            sv = shap_vals_arr
            FEAT_DISP = [{"V012": "Age", "V025": "Urban/Rural", "V106": "Education",
                           "V136": "HH size", "parity": "Parity", "employed": "Employed",
                           "watches_tv": "Media (TV)", "fp_autonomy": "FP autonomy"
                           }.get(f, f) for f in FEATURES]
            sort_ord = np.argsort(-np.abs(sv).mean(axis=0))
            fig, axes6 = plt.subplots(1, 2, figsize=(14, max(5, len(FEATURES) * 0.8 + 2)))
            fig.patch.set_facecolor("white")

            ax6a = axes6[0]
            sdf  = pd.DataFrame({
                "Feature": [FEAT_DISP[i] for i in sort_ord],
                "MAS":     np.abs(sv).mean(axis=0)[sort_ord],
            }).sort_values("MAS", ascending=True)
            bars = ax6a.barh(sdf["Feature"], sdf["MAS"], color="#3498DB",
                              alpha=0.8, edgecolor="white")
            ax6a.set_xlabel("Mean |SHAP value|", fontsize=10)
            ax6a.set_title("Feature Importance\n(Mean |SHAP|)\n"
                            "[Predictive contribution, not causal effect]",   # [REVISION] label
                            fontsize=10, fontweight="bold")
            ax6a.grid(axis="x", alpha=0.3)
            for bar, val in zip(bars, sdf["MAS"]):
                ax6a.text(val + 0.0005, bar.get_y() + bar.get_height() / 2,
                           f"{val:.4f}", va="center", fontsize=8)

            ax6b = axes6[1]
            for rank, fi in enumerate(sort_ord):
                sv_f = sv[:, fi]; xv = X[:, fi]
                xv_n = (xv - xv.min()) / (xv.max() - xv.min() + 1e-9)
                j    = np.random.RandomState(fi).uniform(-0.3, 0.3, len(sv_f))
                ax6b.scatter(sv_f, rank + j, c=xv_n, cmap="RdBu_r",
                              s=3, alpha=0.4, linewidths=0)
            ax6b.axvline(0, color="black", lw=1.2)
            ax6b.set_yticks(range(len(FEATURES)))
            ax6b.set_yticklabels([FEAT_DISP[sort_ord[i]] for i in range(len(FEATURES))],
                                  fontsize=9)
            ax6b.set_xlabel("SHAP value (impact on log-odds)", fontsize=10)
            ax6b.set_title("SHAP Beeswarm\n(red=high feature value, blue=low)",
                            fontsize=11, fontweight="bold")
            ax6b.grid(axis="x", alpha=0.3)
            sm6 = ax6b.scatter([], [], c=[], cmap="RdBu_r")
            plt.colorbar(sm6, ax=ax6b, label="Feature value (normalised)", shrink=0.6)

            fig.suptitle(
                f"Fig 2. SHAP feature importance — {SUBGROUP_MODEL}\n"
                f"(BDHS 2022; n={len(df):,}; "
                f"higher |SHAP| = stronger influence on predicted mCP probability)\n"
                f"Note: SHAP values reflect predictive contributions, not causal effects.",  # [REVISION]
                fontsize=10, fontweight="bold", y=1.02)
            plt.tight_layout()
            plt.savefig("fig2_shap_summary.png", dpi=200, bbox_inches="tight", facecolor="white")
            plt.close(); print("  fig2_shap_summary.png saved.")
        except Exception as e:
            print(f"  fig2_shap: FAILED — {e}")
    else:
        print("  fig2_shap: SHAP values not available, skipping.")

    # Figure 3: Propensity score overlap and Love plot
    fig, axes = plt.subplots(1, 2, figsize=(12, 5))
    fig.patch.set_facecolor("white")
    bins = np.linspace(0, 1, 40)
    axes[0].hist(ps_raw[A == 0], bins=bins, alpha=0.6, color="#E74C3C",
                  density=True, label="Poorest")
    axes[0].hist(ps_raw[A == 1], bins=bins, alpha=0.6, color="#2980B9",
                  density=True, label="Richest")
    axes[0].axvline(0.1, color="gray", ls="--", lw=1)
    axes[0].axvline(0.9, color="gray", ls="--", lw=1)
    axes[0].set_xlabel("Propensity score", fontsize=11)
    axes[0].set_ylabel("Density", fontsize=11)
    axes[0].set_title(f"PS distribution (overlap {ovlp_pct:.1f}%)", fontsize=10)
    axes[0].legend(fontsize=9); axes[0].grid(alpha=0.3)

    smd_r = []; smd_w = []
    for fi in range(X.shape[1]):
        x1 = X[A == 1, fi]; x0 = X[A == 0, fi]
        sp = np.sqrt((x1.var() + x0.var()) / 2) + 1e-9
        smd_r.append(abs(x1.mean() - x0.mean()) / sp)
        smd_w.append(abs(np.average(x1, weights=sw[A == 1]) -
                          np.average(x0, weights=sw[A == 0])) / sp)
    fd = [{"V012": "Age", "V025": "Urban/Rural", "V106": "Education",
            "V136": "HH size", "parity": "Parity", "employed": "Employed",
            "watches_tv": "Media", "fp_autonomy": "FP autonomy"
            }.get(f, f) for f in FEATURES]
    yp = np.arange(len(fd))
    axes[1].scatter(smd_r, yp, color="#E74C3C", s=70, zorder=3, label="Unadjusted")
    axes[1].scatter(smd_w, yp, color="#2980B9", s=70, marker="D", zorder=3,
                     label="Trimmed IPTW (99th pct)")
    axes[1].axvline(0.1, color="gray", ls="--", lw=1, label="SMD=0.1")
    axes[1].set_yticks(yp); axes[1].set_yticklabels(fd, fontsize=9)
    axes[1].set_xlabel("SMD", fontsize=11)
    axes[1].set_title(f"Love plot (trimmed stabilised weights\n"
                       f"cap={TRIM_CAP:.2f}, n_tr={n_tr})", fontsize=10)
    axes[1].legend(fontsize=9); axes[1].grid(alpha=0.3)
    fig.suptitle("Figure 3. Propensity score overlap and covariate balance (BDHS 2022)",
                  fontsize=11, fontweight="bold", y=1.01)
    plt.tight_layout()
    plt.savefig("fig3_propensity_overlap.png", dpi=200, bbox_inches="tight", facecolor="white")
    plt.close(); print("  fig3_propensity_overlap.png saved.")

    # Figure 4: Calibration plots
    # [REVISION] Enhanced calibration interpretation (R2 Comment 5)
    fig, axes = plt.subplots(2, 3, figsize=(14, 9))
    fig.patch.set_facecolor("white")
    for idx, (name, cal) in enumerate(cal_res.items()):
        ax2 = axes[idx // 3][idx % 3]
        ax2.plot([0, 1], [0, 1], "k--", lw=1)
        ax2.plot(cal["mean_pred"], cal["frac_pos"], "o-",
                  color=COLORS.get(name, "#888"), lw=2, ms=5, label=name)
        ax2.set_xlabel("Mean predicted prob", fontsize=9)
        ax2.set_ylabel("Fraction positive" if idx % 3 == 0 else "", fontsize=9)
        ax2.set_title(f"{name}\nBrier={cal['brier']:.4f}  CV-AUC={all_CV_AUC[name]:.4f}",
                       fontsize=9, fontweight="bold")
        ax2.set_xlim(0, 1); ax2.set_ylim(0, 1)
        ax2.legend(fontsize=7); ax2.grid(alpha=0.3)
    fig.suptitle(
        "Figure 4. Calibration plots — all six outcome models (BDHS 2022)\n"
        "Individual ORs depend on predicted probability quality; "     # [REVISION]
        "miscalibration in high-probability range may inflate ORs for high-educated subgroups.",  # [REVISION]
        fontsize=10, fontweight="bold", y=1.01)
    plt.tight_layout()
    plt.savefig("fig4_calibration.png", dpi=200, bbox_inches="tight", facecolor="white")
    plt.close(); print("  fig4_calibration.png saved.")

    # Figure 5: Individual OR density distribution by age group and education
    EC = {"No education": "#7F8C8D", "Primary": "#2980B9",
           "Secondary":   "#27AE60", "Higher":   "#E74C3C"}
    EL = {"No education": "No educ", "Primary": "Primary",
           "Secondary":   "Secondary", "Higher": "Higher educ"}
    fig, axes5 = plt.subplots(1, 3, figsize=(15, 5.5))
    fig.patch.set_facecolor("white")
    xr = np.linspace(0, 14, 600)

    for ai, (ag, tt) in enumerate(zip(["15-24", "25-34", "35-49"],
                                        ["Age 15–24", "Age 25–34", "Age 35–49"])):
        ax5 = axes5[ai]; ax5.set_facecolor("#FAFAFA")
        ax5.axvspan(1, 14, alpha=0.04, color="#E74C3C")
        ax5.axvspan(0, 1,  alpha=0.04, color="#2980B9")
        ax5.axvline(1, color="#2C3E50", lw=2, alpha=0.8, zorder=5)
        try:
            kde_all = gaussian_kde(np.clip(ot, 0.01, 12), bw_method=0.25)
            ax5.fill_between(xr, kde_all(xr), alpha=0.06, color="black")
        except:
            pass
        pct_above = 0; n_age = 0
        for edu in ["No education", "Primary", "Secondary", "Higher"]:
            sub = df[(df.age_group == ag) & (df.edu_label == edu)]
            if len(sub) < 30:
                continue
            vals = np.clip(sub.ind_OR_trim.values, 0.01, 12)
            pct_above += (vals > 1).sum(); n_age += len(vals)
            try:
                kde = gaussian_kde(vals, bw_method=0.4)
                yy  = kde(xr)
                ls  = "--" if edu == "Higher" else "-"
                lw  = 2.8 if edu == "Higher" else 1.8
                mor = sub.ind_OR_trim.median()
                ax5.plot(xr, yy, color=EC[edu], lw=lw, ls=ls,
                          label=f"{EL[edu]} (n={len(sub):,}, OR={mor:.2f})")
                if edu == "Higher":
                    ax5.fill_between(xr, yy, alpha=0.10, color=EC[edu])
                ax5.axvline(mor, color=EC[edu], lw=0.9, alpha=0.45, ls=":")
            except:
                pass
        pct_ab = pct_above / n_age * 100 if n_age > 0 else 0
        ax5.set_xlim(0, 12); ax5.set_ylim(bottom=0)
        ax5.set_xlabel("Individual wealth OR", fontsize=10)
        if ai == 0: ax5.set_ylabel("Density", fontsize=10)
        ax5.set_title(tt, fontsize=11, fontweight="bold", pad=8)
        ax5.legend(fontsize=7.5, loc="upper right", framealpha=0.92)
        ax5.grid(alpha=0.2)
        ax5.annotate(f"~{pct_ab:.0f}% OR>1",
                      xy=(0.97, 0.62), xycoords="axes fraction", fontsize=8,
                      ha="right", va="top", color="#7D0000",
                      bbox=dict(boxstyle="round,pad=0.3", facecolor="#FADBD8", alpha=0.75))

    fig.text(0.25, -0.03, "← Wealth → LOWER contraceptive odds (equity-achieved)",
              ha="center", fontsize=8, color="#1A5276", style="italic")
    fig.text(0.75, -0.03, "Wealth → HIGHER contraceptive odds → (wealth-barrier)",
              ha="center", fontsize=8, color="#922B21", style="italic")
    fig.suptitle(
        f"Fig 5. Individual wealth-associated OR distribution by age group & education\n"
        f"{SUBGROUP_MODEL} T-learner; BDHS 2022; Gini(|log OR|)={gv:.4f} "
        f"(95% CI: {gini_ci_lo:.4f}–{gini_ci_hi:.4f})",   # [REVISION] CI added to title
        fontsize=10, fontweight="bold", y=1.02)
    plt.tight_layout()
    plt.savefig("fig5_or_distribution.png", dpi=200, bbox_inches="tight", facecolor="white")
    plt.close(); print("  fig5_or_distribution.png saved.")

    # Figure 6: E-value forest plot
    fig, ax4 = plt.subplots(figsize=(10, 7))
    fig.patch.set_facecolor("white")
    lev  = [f"{r['Age Group']} / {r['Education']}" for _, r in t6_eval.iterrows()]
    ev_  = [float(r["E-value"])   for _, r in t6_eval.iterrows()]
    mor_ = [float(r["Median OR"]) for _, r in t6_eval.iterrows()]
    cv_  = ["#E74C3C" if e < 2 else "#F39C12" if e < 3 else "#27AE60" for e in ev_]
    yp4  = np.arange(len(lev))
    ax4.barh(yp4, ev_, color=cv_, alpha=0.8, edgecolor="white", height=0.6)
    ax4.scatter(mor_, yp4, color="navy", s=50, zorder=5, label="Median OR")
    ax4.axvline(1, color="black", lw=1.5)
    ax4.axvline(2, color="#F39C12", lw=1, ls="--", alpha=0.7)
    ax4.axvline(3, color="#27AE60", lw=1, ls="--", alpha=0.7)
    ax4.set_yticks(yp4); ax4.set_yticklabels(lev, fontsize=9)
    ax4.set_xlabel("E-value", fontsize=10)
    ax4.set_title(f"Figure 6. E-value analysis\n"
                   f"({SUBGROUP_MODEL} model; BDHS 2022)", fontsize=10, fontweight="bold")
    ax4.legend(handles=[
        mpatches.Patch(color="#E74C3C", alpha=0.8, label="Weak (E<2)"),
        mpatches.Patch(color="#F39C12", alpha=0.8, label="Moderate (2≤E<3)"),
        mpatches.Patch(color="#27AE60", alpha=0.8, label="Strong (E≥3)"),
    ], fontsize=9, loc="lower right")
    ax4.grid(axis="x", alpha=0.3)
    plt.tight_layout()
    plt.savefig("fig6_evalue_plot.png", dpi=200, bbox_inches="tight", facecolor="white")
    plt.close(); print("  fig6_evalue_plot.png saved.")

except ImportError as e:
    print(f"  matplotlib/scipy not available: {e}")

# ─── FINAL SUMMARY ────────────────────────────────────────────────────────────
print("\n" + "=" * 70)
print("RESULTS SUMMARY")
print("=" * 70)
print(f"\nSAMPLE: n={len(df):,}  poor={n_poor:,}  rich={n_rich:,}")
print(f"FEATURES ({len(FEATURES)}): {FEATURES}")
print(f"BEST MODEL (CV-AUC): {BEST}  ({all_CV_AUC[BEST]:.4f})")
print(f"SUBGROUP MODEL: {SUBGROUP_MODEL}")
print(f"Gini(|log OR|)={gv:.4f}  95% CI: [{gini_ci_lo:.4f}, {gini_ci_hi:.4f}]")
print(f"OR>1: {pa:.1f}%  OR<1: {pb:.1f}%")
print(f"PS overlap: {ovlp_pct:.1f}%  Untrimmed max: {max_raw:.2f}  "
      f"Trimmed ESS: {ess:.0f} ({ess_pct:.1f}%)")
print(f"Bootstrap: {blabel}")
print("\n--- TABLE 1 ---"); print(t1.to_string(index=False))
print("\n--- TABLE 2 (model performance + AIPW) ---"); print(t2.to_string(index=False))
print("\n--- TABLE 3 (Age × Education subgroup ORs + ARD) ---"); print(t3.to_string(index=False))
print("\n--- TABLE 4 (Residence × Education subgroup ORs) ---"); print(t4.to_string(index=False))
if not t5_int.empty:
    print("\n--- TABLE 5 (Interaction tests) ---"); print(t5_int.to_string(index=False))
print("\n--- TABLE 6 (E-values) ---"); print(t6_eval.to_string(index=False))
print("\n--- TABLE 7 (Programmatic targeting) ---")
print(t8[["Age group", "Education", "Median OR", "Priority"]].to_string(index=False))
if shap_rows:
    print("\n--- TABLE 8 (SHAP feature importance) ---")
    print(pd.DataFrame(shap_rows)[["Rank", "Feature", "Mean |SHAP|"]].to_string(index=False))
print("\n--- TRIMMING SENSITIVITY (Supp Table S3) ---")
print(t3_sens[["Age Group","Education","Median OR (99th pct)","Median OR (95th pct)","Shift","Substantive change?"]].to_string(index=False))
print("\n" + "=" * 70)

# ─── FILE LIST ────────────────────────────────────────────────────────────────
outputs = [
    "results_table1_sample.csv",
    "results_table2_model_performance.csv",
    "results_table3_age_education.csv",           # now includes ARD column
    "results_table4_residence_education.csv",
    "results_table5_interaction.csv",
    "results_table6_evalue.csv",
    "results_table8_policy_targeting.csv",
    "results_table9_shap_importance.csv",
    "results_weight_diagnostics.csv",
    "results_gini_heterogeneity.csv",             # now includes bootstrap CI
    "results_table2_aipw_detail.csv",
    "results_suppTableS1_hyperparameters.csv",    # [REVISION] new
    "results_suppTableS2_aipw_rfps_robustness.csv",  # [REVISION] new
    "results_suppTableS3_trimming95_sensitivity.csv", # [REVISION] new
    "fig1_dag.png",
    "fig2_shap_summary.png",
    "fig3_propensity_overlap.png",
    "fig4_calibration.png",
    "fig5_or_distribution.png",
    "fig6_evalue_plot.png",
    "individual_ORs_lr.npy",
    "individual_ORs_rf.npy",
    "individual_ORs_gb.npy",
    "individual_ORs_xgb.npy",
    "individual_ORs_lgb.npy",
    "individual_ORs_stack.npy",
]
print("\nOutput files:")
for f in outputs:
    print(f"  {'OK     ' if os.path.exists(f) else 'MISSING'} {f}")
print("\nDone.")