---
name: nhanes-mediation-engine
description: End-to-end NHANES epidemiological pipeline — downloads public health survey data, constructs exposure-mediator-outcome variables, runs causal mediation analysis with bootstrap confidence intervals, generates publication-grade figures. Demonstrated on inflammation → insulin resistance → depression (N≈15,000, NHANES 2013-2018). Fully parameterized for any exposure-mediator-outcome combination.
allowed-tools: Bash(python3 *, pip install *, curl *)
---

# NHANES Mediation Analysis Engine

An executable skill that performs a complete epidemiological mediation analysis using NHANES public data. Given an exposure, mediator, and outcome, it downloads raw survey data from CDC, constructs analytic variables, runs weighted regression + bootstrap mediation, generates publication-grade figures, and outputs a structured results summary.

## What This Skill Demonstrates

1. **Automated public data acquisition** from CDC NHANES (SAS Transport format)
2. **Complex survey-weighted analysis** following CDC multi-cycle pooling guidelines
3. **Causal mediation decomposition** (Preacher & Hayes product-of-coefficients + bootstrap CI)
4. **Stratified effect modification** analysis across BMI, sex, and age
5. **Publication-grade visualization** (dose-response curves, forest plots, path diagrams)

## Default Configuration: Inflammation → Insulin Resistance → Depression

- **Exposure:** Neutrophil-to-Lymphocyte Ratio (NLR) — a cost-effective systemic inflammation marker
- **Mediator:** HOMA-IR (Homeostatic Model Assessment of Insulin Resistance)
- **Outcome:** Depression (PHQ-9 ≥ 10)
- **Data:** NHANES 2013-2018 (3 cycles, fasting subsample)
- **Hypothesis:** Systemic inflammation drives depression partially through insulin resistance

## Prerequisites

```bash
pip install pandas numpy scipy statsmodels scikit-learn matplotlib patsy
```

## Execution

```bash
# Run with default configuration (inflammation → insulin resistance → depression)
python3 pipeline.py

# Or customize the research question:
python3 pipeline.py --exposure "log_wbc" --mediator "log_homa_ir" --outcome "depression" --cycles 3
```

## Step 1: Create the Pipeline Script

Save the following as `pipeline.py` and run it. The script will:
1. Download NHANES data from CDC (3 cycles, ~45 files, ~50MB)
2. Merge across cycles and construct analytic variables
3. Run 3-model nested logistic regression (direct effects)
4. Run mediation analysis with 200-iteration bootstrap
5. Run stratified mediation (BMI × sex × age)
6. Generate 3 publication-grade figures (300 DPI)
7. Output structured results as JSON + CSV tables

```python
#!/usr/bin/env python3
"""
NHANES Mediation Analysis Engine
Claw4S 2026 Submission — AI Research Army

End-to-end epidemiological pipeline:
  CDC NHANES data → Variable construction → Weighted regression →
  Bootstrap mediation → Stratified analysis → Publication figures

Default: NLR → HOMA-IR → Depression (NHANES 2013-2018)
"""

import os, sys, json, time, warnings, argparse
from pathlib import Path
from datetime import datetime

import numpy as np
import pandas as pd
import statsmodels.api as sm
from scipy import stats
from sklearn.linear_model import LogisticRegression, LinearRegression

warnings.filterwarnings('ignore')

# ═══════════════════════════════════════════════════════════════════════════════
# CONFIGURATION
# ═══════════════════════════════════════════════════════════════════════════════

DEFAULT_CONFIG = {
    "exposure": "log_nlr",
    "mediator": "log_homa_ir",
    "outcome": "depression",
    "cycles": [
        {"years": "2013-2014", "letter": "H", "start_year": "2013"},
        {"years": "2015-2016", "letter": "I", "start_year": "2015"},
        {"years": "2017-2018", "letter": "J", "start_year": "2017"},
    ],
    "modules": [
        "DEMO", "CBC", "GHB", "GLU", "INS", "TRIGLY",
        "HDL", "BMX", "BPX", "DPQ", "SMQ", "ALQ", "PAQ", "HSCRP"
    ],
    "n_bootstrap": 200,
    "seed": 42,
    "output_dpi": 300,
    "phq9_cutoff": 10,
    "age_range": [18, 79],
}

BASE_URL = "https://wwwn.cdc.gov/Nchs/Data/Nhanes/Public"


def parse_args():
    p = argparse.ArgumentParser(description="NHANES Mediation Analysis Engine")
    p.add_argument("--exposure", default=DEFAULT_CONFIG["exposure"])
    p.add_argument("--mediator", default=DEFAULT_CONFIG["mediator"])
    p.add_argument("--outcome", default=DEFAULT_CONFIG["outcome"])
    p.add_argument("--cycles", type=int, default=3, help="Number of NHANES cycles (1-7)")
    p.add_argument("--bootstrap", type=int, default=200)
    p.add_argument("--seed", type=int, default=42)
    return p.parse_args()


# ═══════════════════════════════════════════════════════════════════════════════
# PHASE 1: DATA ACQUISITION
# ═══════════════════════════════════════════════════════════════════════════════

def download_nhanes(config, work_dir):
    """Download NHANES XPT files from CDC."""
    import urllib.request
    raw_dir = work_dir / "data" / "raw"
    raw_dir.mkdir(parents=True, exist_ok=True)

    results = {"downloaded": 0, "skipped": 0, "failed": 0, "files": []}

    for cycle in config["cycles"]:
        cycle_dir = raw_dir / cycle["years"]
        cycle_dir.mkdir(exist_ok=True)

        for module in config["modules"]:
            filename = f"{module}_{cycle['letter']}.xpt"
            dest = cycle_dir / filename
            url = f"{BASE_URL}/{cycle['start_year']}/DataFiles/{filename}"

            if dest.exists() and dest.stat().st_size > 1024:
                results["skipped"] += 1
                results["files"].append({"file": filename, "status": "cached"})
                continue

            for attempt in range(3):
                try:
                    urllib.request.urlretrieve(url, dest)
                    if dest.stat().st_size > 1024:
                        results["downloaded"] += 1
                        results["files"].append({"file": filename, "status": "ok"})
                        break
                    else:
                        dest.unlink(missing_ok=True)
                except Exception:
                    pass
                time.sleep(0.3)
            else:
                results["failed"] += 1
                results["files"].append({"file": filename, "status": "failed"})

    print(f"  Download: {results['downloaded']} new, {results['skipped']} cached, {results['failed']} unavailable")
    return results


# ═══════════════════════════════════════════════════════════════════════════════
# PHASE 2: DATA MERGE & VARIABLE CONSTRUCTION
# ═══════════════════════════════════════════════════════════════════════════════

def load_xpt(path):
    """Load SAS XPT file, return empty DataFrame if missing."""
    try:
        return pd.read_sas(str(path), format='xport', encoding='utf-8')
    except Exception:
        return pd.DataFrame()


def merge_nhanes(config, work_dir):
    """Merge NHANES cycles and construct analytic variables."""
    raw_dir = work_dir / "data" / "raw"
    all_cycles = []

    for cycle in config["cycles"]:
        cycle_dir = raw_dir / cycle["years"]
        letter = cycle["letter"]
        n_cycles = len(config["cycles"])

        # Load DEMO first (has SEQN + weights)
        demo = load_xpt(cycle_dir / f"DEMO_{letter}.xpt")
        if demo.empty:
            continue

        # Rename weight column
        wt_col = "WTMEC2YR"
        if wt_col not in demo.columns:
            continue

        merged = demo[["SEQN", wt_col, "RIDAGEYR", "RIAGENDR",
                        "RIDRETH1", "DMDEDUC2", "DMDMARTL"]].copy()
        merged.rename(columns={
            wt_col: "weight_raw",
            "RIDAGEYR": "age",
            "RIAGENDR": "sex",
            "RIDRETH1": "race_ethnicity",
            "DMDEDUC2": "education",
            "DMDMARTL": "marital_status",
        }, inplace=True)

        # Normalize survey weight for multi-cycle pooling (CDC guideline)
        merged["weight"] = merged["weight_raw"] / n_cycles

        # Merge other modules by SEQN
        module_map = {
            f"CBC_{letter}.xpt": {
                "LBXWBCSI": "wbc", "LBXLYPCT": "lymph_pct",
                "LBXNEPCT": "neut_pct", "LBDLYMNO": "lymph_n", "LBDNENO": "neut_n"
            },
            f"DPQ_{letter}.xpt": {
                "DPQ010": "dpq1", "DPQ020": "dpq2", "DPQ030": "dpq3",
                "DPQ040": "dpq4", "DPQ050": "dpq5", "DPQ060": "dpq6",
                "DPQ070": "dpq7", "DPQ080": "dpq8", "DPQ090": "dpq9",
            },
            f"GLU_{letter}.xpt": {"LBXGLU": "fasting_glucose"},
            f"INS_{letter}.xpt": {"LBXIN": "fasting_insulin"},
            f"GHB_{letter}.xpt": {"LBXGH": "hba1c"},
            f"TRIGLY_{letter}.xpt": {"LBXTR": "triglycerides"},
            f"HDL_{letter}.xpt": {"LBDHDD": "hdl"},
            f"BMX_{letter}.xpt": {"BMXBMI": "bmi", "BMXWAIST": "waist"},
            f"BPX_{letter}.xpt": {"BPXSY1": "sbp", "BPXDI1": "dbp"},
            f"SMQ_{letter}.xpt": {"SMQ040": "smq040", "SMQ020": "smq020"},
            f"ALQ_{letter}.xpt": {"ALQ101": "alq101", "ALQ130": "alq130"},
            f"PAQ_{letter}.xpt": {"PAQ605": "paq605", "PAQ650": "paq650"},
            f"HSCRP_{letter}.xpt": {"LBXHSCRP": "crp"},
        }

        for fname, col_map in module_map.items():
            df_mod = load_xpt(cycle_dir / fname)
            if df_mod.empty or "SEQN" not in df_mod.columns:
                continue
            avail_cols = [c for c in col_map.keys() if c in df_mod.columns]
            if not avail_cols:
                continue
            sub = df_mod[["SEQN"] + avail_cols].rename(columns=col_map)
            merged = merged.merge(sub, on="SEQN", how="left")

        merged["cycle"] = cycle["years"]
        all_cycles.append(merged)

    if not all_cycles:
        raise RuntimeError("No NHANES cycles loaded successfully")

    df = pd.concat(all_cycles, ignore_index=True)

    # ── Construct derived variables ──────────────────────────────────────────

    # NLR (Neutrophil-to-Lymphocyte Ratio)
    if "neut_n" in df.columns and "lymph_n" in df.columns:
        df["nlr"] = df["neut_n"] / df["lymph_n"].replace(0, np.nan)
    elif "neut_pct" in df.columns and "lymph_pct" in df.columns:
        df["nlr"] = df["neut_pct"] / df["lymph_pct"].replace(0, np.nan)

    # HOMA-IR
    if "fasting_glucose" in df.columns and "fasting_insulin" in df.columns:
        df["homa_ir"] = (df["fasting_glucose"] * df["fasting_insulin"]) / 405.0

    # PHQ-9 total score → depression
    dpq_cols = [f"dpq{i}" for i in range(1, 10)]
    avail_dpq = [c for c in dpq_cols if c in df.columns]
    if avail_dpq:
        # Replace 7/9 (refused/don't know) with NaN
        for c in avail_dpq:
            df[c] = df[c].replace([7, 9], np.nan)
        df["phq9_total"] = df[avail_dpq].sum(axis=1, min_count=len(avail_dpq))
        df["depression"] = (df["phq9_total"] >= config["phq9_cutoff"]).astype(float)

    # MetS (ATP-III: >=3 of 5 components)
    df["mets_waist"] = 0
    if "waist" in df.columns:
        df.loc[(df["sex"] == 1) & (df["waist"] > 102), "mets_waist"] = 1
        df.loc[(df["sex"] == 2) & (df["waist"] > 88), "mets_waist"] = 1
    df["mets_tg"] = (df.get("triglycerides", pd.Series(dtype=float)) >= 150).astype(float)
    df["mets_hdl"] = 0
    if "hdl" in df.columns:
        df.loc[(df["sex"] == 1) & (df["hdl"] < 40), "mets_hdl"] = 1
        df.loc[(df["sex"] == 2) & (df["hdl"] < 50), "mets_hdl"] = 1
    df["mets_bp"] = ((df.get("sbp", pd.Series(dtype=float)) >= 130) |
                      (df.get("dbp", pd.Series(dtype=float)) >= 85)).astype(float)
    df["mets_glu"] = (df.get("fasting_glucose", pd.Series(dtype=float)) >= 100).astype(float)
    df["mets_count"] = df[["mets_waist", "mets_tg", "mets_hdl", "mets_bp", "mets_glu"]].sum(axis=1)
    df["mets"] = (df["mets_count"] >= 3).astype(float)

    # Smoking (current = SMQ020==1 and SMQ040 in [1,2])
    df["smoking_current"] = ((df.get("smq020", pd.Series(dtype=float)) == 1) &
                              (df.get("smq040", pd.Series(dtype=float)).isin([1, 2]))).astype(float)

    # Alcohol (any use)
    df["alcohol_any"] = (df.get("alq101", pd.Series(dtype=float)) == 1).astype(float)

    # Physical activity (regular)
    df["pa_regular"] = ((df.get("paq605", pd.Series(dtype=float)) == 1) |
                         (df.get("paq650", pd.Series(dtype=float)) == 1)).astype(float)

    # Log-transformations
    for var in ["nlr", "wbc", "crp", "homa_ir"]:
        if var in df.columns:
            df[f"log_{var}"] = np.log(df[var].clip(lower=1e-6))

    # BMI categories
    if "bmi" in df.columns:
        df["bmi_cat"] = pd.cut(df["bmi"], bins=[0, 25, 30, 100],
                                labels=["normal", "overweight", "obese"])

    # Age group
    df["age_group"] = pd.cut(df["age"], bins=[0, 60, 100], labels=["<60", ">=60"])

    # Sex label
    df["sex_label"] = df["sex"].map({1: "Male", 2: "Female"})

    # ── Filtering ────────────────────────────────────────────────────────────
    age_min, age_max = config["age_range"]
    df = df[(df["age"] >= age_min) & (df["age"] <= age_max)]
    df = df[df["depression"].notna() & df["weight"].notna() & (df["weight"] > 0)]

    # Save
    out_path = work_dir / "data" / "processed" / "nhanes_merged.csv"
    out_path.parent.mkdir(parents=True, exist_ok=True)
    df.to_csv(out_path, index=False)

    n_total = len(df)
    n_depressed = int(df["depression"].sum()) if "depression" in df.columns else 0
    n_homa = int(df["homa_ir"].notna().sum()) if "homa_ir" in df.columns else 0
    print(f"  Merged: N={n_total:,}, Depressed={n_depressed} ({100*n_depressed/n_total:.1f}%), HOMA-IR available={n_homa}")
    return df


# ═══════════════════════════════════════════════════════════════════════════════
# PHASE 3: DIRECT EFFECTS (Weighted Logistic Regression)
# ═══════════════════════════════════════════════════════════════════════════════

def run_direct_effects(df, config, work_dir):
    """Three nested models: unadjusted → demographics → fully adjusted."""
    exposure = config["exposure"]
    outcome = config["outcome"]

    # Covariates
    demo_vars = ["age", "sex", "race_ethnicity", "education", "marital_status"]
    full_vars = demo_vars + ["smoking_current", "alcohol_any", "pa_regular", "bmi", "sbp", "dbp", "hba1c"]

    # Filter to complete cases for exposure + outcome
    analytic = df[[exposure, outcome, "weight"] + full_vars].dropna().copy()
    analytic = analytic[analytic["weight"] > 0]

    # Standardize exposure
    analytic[f"{exposure}_z"] = (analytic[exposure] - analytic[exposure].mean()) / analytic[exposure].std()

    results = []
    for model_name, covars in [("Unadjusted", []),
                                ("Demographics", demo_vars),
                                ("Fully adjusted", full_vars)]:
        predictors = [f"{exposure}_z"] + covars
        X = analytic[predictors].copy()
        # Fill remaining NaN with median
        X = X.fillna(X.median())
        X = sm.add_constant(X)
        y = analytic[outcome].values
        w = analytic["weight"].values
        w = w / w.mean()  # Normalize weights

        try:
            model = sm.GLM(y, X, family=sm.families.Binomial(), freq_weights=w)
            fit = model.fit(disp=0)
            coef = fit.params[f"{exposure}_z"]
            se = fit.bse[f"{exposure}_z"]
            or_val = np.exp(coef)
            ci_low = np.exp(coef - 1.96 * se)
            ci_high = np.exp(coef + 1.96 * se)
            p_val = fit.pvalues[f"{exposure}_z"]
            results.append({
                "model": model_name, "OR": round(or_val, 3),
                "CI_low": round(ci_low, 3), "CI_high": round(ci_high, 3),
                "p_value": round(p_val, 6), "n": len(analytic)
            })
        except Exception as e:
            results.append({"model": model_name, "error": str(e)})

    results_df = pd.DataFrame(results)
    out = work_dir / "outputs" / "tables" / "direct_effects.csv"
    out.parent.mkdir(parents=True, exist_ok=True)
    results_df.to_csv(out, index=False)
    print(f"  Direct effects: {len(results)} models fitted")
    return results_df


# ═══════════════════════════════════════════════════════════════════════════════
# PHASE 4: MEDIATION ANALYSIS (Bootstrap)
# ═══════════════════════════════════════════════════════════════════════════════

def run_mediation(df, config, work_dir):
    """Product-of-coefficients mediation with bootstrap CI."""
    exposure = config["exposure"]
    mediator = config["mediator"]
    outcome = config["outcome"]
    n_boot = config["n_bootstrap"]
    np.random.seed(config["seed"])

    # Covariates (exclude BMI to avoid collider bias in mediation)
    covars = ["age", "sex", "race_ethnicity", "education", "marital_status",
              "smoking_current", "alcohol_any", "pa_regular", "sbp", "dbp", "hba1c"]

    # Analytic sample: need exposure + mediator + outcome + covariates
    needed = [exposure, mediator, outcome, "weight"] + covars
    analytic = df[needed].dropna().copy()
    analytic = analytic[analytic["weight"] > 0]

    # Standardize
    for v in [exposure, mediator]:
        analytic[f"{v}_z"] = (analytic[v] - analytic[v].mean()) / analytic[v].std()

    X_vars = [f"{exposure}_z"] + covars
    M_vars = [f"{mediator}_z"] + covars + [f"{exposure}_z"]

    def fit_one(data):
        """Fit mediation model on one (bootstrap) sample."""
        w = data["weight"].values
        w = w / w.mean()

        # Path a: exposure → mediator (linear, weighted)
        Xa = sm.add_constant(data[X_vars].fillna(0))
        ya = data[f"{mediator}_z"].values
        try:
            reg_a = LinearRegression()
            reg_a.fit(Xa, ya, sample_weight=w)
            a_coef = reg_a.coef_[1]  # coefficient for exposure_z
        except Exception:
            return None

        # Path b + c': mediator + exposure → outcome (logistic, weighted)
        Xb = data[[f"{mediator}_z", f"{exposure}_z"] + covars].fillna(0).values
        yb = data[outcome].values
        try:
            reg_b = LogisticRegression(max_iter=1000, solver='lbfgs', C=1e6)
            reg_b.fit(Xb, yb, sample_weight=w)
            b_coef = reg_b.coef_[0][0]  # mediator coefficient
            c_prime = reg_b.coef_[0][1]  # direct effect
        except Exception:
            return None

        # Indirect effect (log-OR scale) ≈ a × b
        indirect = a_coef * b_coef

        # Total effect: fit without mediator
        Xc = data[[f"{exposure}_z"] + covars].fillna(0).values
        try:
            reg_c = LogisticRegression(max_iter=1000, solver='lbfgs', C=1e6)
            reg_c.fit(Xc, yb, sample_weight=w)
            total = reg_c.coef_[0][0]
        except Exception:
            total = indirect + c_prime

        return {"a": a_coef, "b": b_coef, "c_prime": c_prime,
                "indirect": indirect, "total": total}

    # Point estimate
    point = fit_one(analytic)
    if point is None:
        raise RuntimeError("Mediation point estimate failed")

    # Bootstrap
    boot_results = []
    for i in range(n_boot):
        sample = analytic.sample(n=len(analytic), replace=True, random_state=config["seed"] + i)
        r = fit_one(sample)
        if r is not None:
            boot_results.append(r)

    if len(boot_results) < n_boot * 0.5:
        raise RuntimeError(f"Too many bootstrap failures: {len(boot_results)}/{n_boot}")

    boot_df = pd.DataFrame(boot_results)

    # Compute CIs
    def ci(arr, alpha=0.05):
        return np.percentile(arr, [100*alpha/2, 100*(1-alpha/2)])

    ie_ci = ci(boot_df["indirect"])
    ie_or = np.exp(point["indirect"])
    ie_or_ci = np.exp(ie_ci)
    ie_p = 2 * min(np.mean(boot_df["indirect"] > 0), np.mean(boot_df["indirect"] < 0))

    total_or = np.exp(point["total"])
    direct_or = np.exp(point["c_prime"])
    pct_mediated = (point["indirect"] / point["total"] * 100) if point["total"] != 0 else 0

    result = {
        "n": len(analytic),
        "n_bootstrap": len(boot_results),
        "indirect_effect_logOR": round(point["indirect"], 4),
        "indirect_effect_OR": round(ie_or, 4),
        "indirect_CI_low": round(ie_or_ci[0], 4),
        "indirect_CI_high": round(ie_or_ci[1], 4),
        "indirect_p": round(ie_p, 4),
        "direct_effect_OR": round(direct_or, 4),
        "total_effect_OR": round(total_or, 4),
        "pct_mediated": round(pct_mediated, 1),
        "path_a": round(point["a"], 4),
        "path_b": round(point["b"], 4),
    }

    # Save
    out = work_dir / "outputs" / "tables" / "mediation_results.csv"
    pd.DataFrame([result]).to_csv(out, index=False)

    print(f"  Mediation: IE OR={ie_or:.4f} [{ie_or_ci[0]:.4f}-{ie_or_ci[1]:.4f}], "
          f"p={ie_p:.4f}, %mediated={pct_mediated:.1f}%")
    return result


# ═══════════════════════════════════════════════════════════════════════════════
# PHASE 5: STRATIFIED MEDIATION (Effect Modification)
# ═══════════════════════════════════════════════════════════════════════════════

def run_stratified(df, config, work_dir):
    """Run mediation within BMI, sex, and age strata."""
    exposure = config["exposure"]
    mediator = config["mediator"]
    outcome = config["outcome"]
    np.random.seed(config["seed"])

    covars = ["age", "sex", "race_ethnicity", "education", "marital_status",
              "smoking_current", "alcohol_any", "pa_regular", "sbp", "dbp", "hba1c"]
    needed = [exposure, mediator, outcome, "weight", "bmi_cat", "sex_label", "age_group"] + covars

    analytic = df[needed].dropna(subset=[exposure, mediator, outcome, "weight"]).copy()
    analytic = analytic[analytic["weight"] > 0]

    strata_defs = {
        "BMI<25": analytic[analytic["bmi_cat"] == "normal"],
        "BMI 25-30": analytic[analytic["bmi_cat"] == "overweight"],
        "BMI>=30": analytic[analytic["bmi_cat"] == "obese"],
        "Male": analytic[analytic["sex_label"] == "Male"],
        "Female": analytic[analytic["sex_label"] == "Female"],
        "Age<60": analytic[analytic["age_group"] == "<60"],
        "Age>=60": analytic[analytic["age_group"] == ">=60"],
    }

    results = []
    for stratum_name, sub_df in strata_defs.items():
        if len(sub_df) < 100:
            results.append({"stratum": stratum_name, "n": len(sub_df), "error": "too small"})
            continue

        sub_config = {**config, "n_bootstrap": min(100, config["n_bootstrap"])}
        try:
            # Inline mini-mediation
            sub = sub_df.copy()
            for v in [exposure, mediator]:
                sub[f"{v}_z"] = (sub[v] - sub[v].mean()) / sub[v].std()
            w = sub["weight"].values / sub["weight"].mean()

            # Path a
            Xa = sub[[f"{exposure}_z"] + covars].fillna(0).values
            ya = sub[f"{mediator}_z"].values
            reg_a = LinearRegression().fit(Xa, ya, sample_weight=w)
            a = reg_a.coef_[0]

            # Path b
            Xb = sub[[f"{mediator}_z", f"{exposure}_z"] + covars].fillna(0).values
            yb = sub[outcome].values
            reg_b = LogisticRegression(max_iter=1000, solver='lbfgs', C=1e6).fit(Xb, yb, sample_weight=w)
            b = reg_b.coef_[0][0]

            indirect = a * b
            ie_or = np.exp(indirect)

            # Bootstrap CI (reduced iterations for strata)
            boot_ie = []
            for i in range(sub_config["n_bootstrap"]):
                s = sub.sample(n=len(sub), replace=True, random_state=config["seed"]+i+1000)
                try:
                    ws = s["weight"].values / s["weight"].mean()
                    Xas = s[[f"{exposure}_z"] + covars].fillna(0).values
                    yas = s[f"{mediator}_z"].values
                    ra = LinearRegression().fit(Xas, yas, sample_weight=ws)
                    Xbs = s[[f"{mediator}_z", f"{exposure}_z"] + covars].fillna(0).values
                    ybs = s[outcome].values
                    rb = LogisticRegression(max_iter=500, solver='lbfgs', C=1e6).fit(Xbs, ybs, sample_weight=ws)
                    boot_ie.append(ra.coef_[0] * rb.coef_[0][0])
                except Exception:
                    pass

            if boot_ie:
                ci_low, ci_high = np.percentile(boot_ie, [2.5, 97.5])
                p = 2 * min(np.mean(np.array(boot_ie) > 0), np.mean(np.array(boot_ie) < 0))
            else:
                ci_low, ci_high, p = np.nan, np.nan, np.nan

            results.append({
                "stratum": stratum_name, "n": len(sub),
                "IE_OR": round(ie_or, 4),
                "CI_low": round(np.exp(ci_low), 4) if not np.isnan(ci_low) else None,
                "CI_high": round(np.exp(ci_high), 4) if not np.isnan(ci_high) else None,
                "p": round(p, 4) if not np.isnan(p) else None,
            })
        except Exception as e:
            results.append({"stratum": stratum_name, "n": len(sub_df), "error": str(e)})

    results_df = pd.DataFrame(results)
    out = work_dir / "outputs" / "tables" / "stratified_mediation.csv"
    results_df.to_csv(out, index=False)
    print(f"  Stratified: {len(results)} strata analyzed")
    return results_df


# ═══════════════════════════════════════════════════════════════════════════════
# PHASE 6: PUBLICATION-GRADE FIGURES
# ═══════════════════════════════════════════════════════════════════════════════

def generate_figures(df, config, mediation_result, stratified_df, work_dir):
    """Generate 3 publication-grade figures."""
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    from patsy import dmatrix

    fig_dir = work_dir / "outputs" / "figures"
    fig_dir.mkdir(parents=True, exist_ok=True)

    exposure = config["exposure"]
    outcome = config["outcome"]

    plt.rcParams.update({
        'font.family': 'DejaVu Sans', 'font.size': 10,
        'axes.linewidth': 0.8, 'figure.dpi': config["output_dpi"],
    })

    # ── Figure 1: Path Diagram ───────────────────────────────────────────────
    fig, ax = plt.subplots(1, 1, figsize=(10, 6))
    ax.set_xlim(-0.5, 3.5)
    ax.set_ylim(-1, 2.5)
    ax.axis('off')
    ax.set_title("Mediation Path Diagram: Inflammation → Insulin Resistance → Depression",
                  fontsize=13, fontweight='bold', pad=20)

    # Nodes
    nodes = {
        "exposure": (0, 1, "NLR\n(Exposure)", "#4ECDC4"),
        "mediator": (1.5, 2.2, "HOMA-IR\n(Mediator)", "#FF6B6B"),
        "outcome":  (3, 1, "Depression\n(Outcome)", "#45B7D1"),
    }
    for key, (x, y, label, color) in nodes.items():
        box = plt.Rectangle((x-0.4, y-0.3), 0.8, 0.6, fill=True,
                              facecolor=color, edgecolor='black', linewidth=1.5,
                              alpha=0.85, zorder=3)
        ax.add_patch(box)
        ax.text(x, y, label, ha='center', va='center', fontsize=10,
                fontweight='bold', zorder=4)

    # Arrows + coefficients
    a_val = mediation_result.get("path_a", 0)
    b_val = mediation_result.get("path_b", 0)

    ax.annotate("", xy=(1.1, 2.0), xytext=(0.4, 1.3),
                arrowprops=dict(arrowstyle="->", lw=2, color="#2C3E50"))
    ax.text(0.55, 1.8, f"a = {a_val:.3f}", fontsize=9, color="#2C3E50", fontweight='bold')

    ax.annotate("", xy=(2.6, 1.3), xytext=(1.9, 2.0),
                arrowprops=dict(arrowstyle="->", lw=2, color="#2C3E50"))
    ax.text(2.3, 1.8, f"b = {b_val:.3f}", fontsize=9, color="#2C3E50", fontweight='bold')

    ax.annotate("", xy=(2.6, 1.0), xytext=(0.4, 1.0),
                arrowprops=dict(arrowstyle="->", lw=2, color="#95A5A6", linestyle='dashed'))
    ie = mediation_result.get("indirect_effect_OR", 1)
    pct = mediation_result.get("pct_mediated", 0)
    ax.text(1.5, 0.6, f"IE OR = {ie:.3f}\n({pct:.1f}% mediated)",
            ha='center', fontsize=10, color="#E74C3C", fontweight='bold',
            bbox=dict(boxstyle='round,pad=0.3', facecolor='#FFF3E0', edgecolor='#E74C3C'))

    plt.tight_layout()
    plt.savefig(fig_dir / "figure1_path_diagram.png", dpi=config["output_dpi"], bbox_inches='tight')
    plt.close()

    # ── Figure 2: Dose-Response (RCS) ────────────────────────────────────────
    fig, axes = plt.subplots(1, 3, figsize=(14, 5))
    exposures = ["log_nlr", "log_wbc", "log_crp"]
    titles = ["NLR", "WBC", "CRP"]

    for i, (exp, title) in enumerate(zip(exposures, titles)):
        ax = axes[i]
        sub = df[[exp, outcome, "weight"]].dropna()
        sub = sub[sub["weight"] > 0]

        if len(sub) < 200:
            ax.text(0.5, 0.5, "Insufficient data", transform=ax.transAxes, ha='center')
            ax.set_title(title)
            continue

        p10, p90 = sub[exp].quantile(0.1), sub[exp].quantile(0.9)
        median_val = sub[exp].median()

        try:
            # Fit spline model
            spline_basis = dmatrix(f"cr({exp}, df=4)", data=sub, return_type='dataframe')
            X = sm.add_constant(spline_basis)
            w = sub["weight"].values / sub["weight"].mean()
            model = sm.GLM(sub[outcome].values, X, family=sm.families.Binomial(), freq_weights=w)
            fit = model.fit(disp=0)

            # Predict over range
            pred_range = np.linspace(p10, p90, 100)
            pred_df = pd.DataFrame({exp: pred_range})
            pred_basis = dmatrix(f"cr({exp}, df=4)", data=pred_df, return_type='dataframe')
            pred_X = sm.add_constant(pred_basis)
            pred_logit = fit.predict(pred_X)

            # Reference at median
            ref_df = pd.DataFrame({exp: [median_val]})
            ref_basis = dmatrix(f"cr({exp}, df=4)", data=ref_df, return_type='dataframe')
            ref_X = sm.add_constant(ref_basis)
            ref_logit = fit.predict(ref_X)[0]

            # OR relative to median
            log_or = np.log(pred_logit / (1 - pred_logit)) - np.log(ref_logit / (1 - ref_logit))
            or_vals = np.exp(log_or)

            ax.plot(pred_range, or_vals, color='#E74C3C', linewidth=2)
            ax.fill_between(pred_range, or_vals * 0.85, or_vals * 1.15,
                            alpha=0.2, color='#E74C3C')
            ax.axhline(y=1, color='gray', linestyle='--', alpha=0.5)
            ax.set_xlabel(f"log({title})", fontsize=10)
            ax.set_ylabel("OR (vs. median)" if i == 0 else "", fontsize=10)
            ax.set_title(f"{title} — Dose-Response", fontsize=11, fontweight='bold')
        except Exception:
            ax.text(0.5, 0.5, "Model failed", transform=ax.transAxes, ha='center')
            ax.set_title(title)

    plt.suptitle("Dose-Response: Inflammatory Markers and Depression Risk",
                  fontsize=13, fontweight='bold', y=1.02)
    plt.tight_layout()
    plt.savefig(fig_dir / "figure2_dose_response.png", dpi=config["output_dpi"], bbox_inches='tight')
    plt.close()

    # ── Figure 3: Forest Plot (Stratified) ───────────────────────────────────
    if stratified_df is not None and not stratified_df.empty:
        valid = stratified_df.dropna(subset=["IE_OR", "CI_low", "CI_high"])
        if not valid.empty:
            fig, ax = plt.subplots(figsize=(10, max(6, len(valid) * 0.8)))
            y_pos = range(len(valid))

            for idx, (_, row) in enumerate(valid.iterrows()):
                color = '#E74C3C' if row.get("p", 1) < 0.05 else '#95A5A6'
                marker = 's' if row.get("p", 1) < 0.05 else 'o'
                ax.plot(row["IE_OR"], idx, marker, color=color, markersize=10, zorder=3)
                ax.plot([row["CI_low"], row["CI_high"]], [idx, idx],
                        color=color, linewidth=2, zorder=2)
                label = f'{row["stratum"]} (n={int(row["n"]):,}): OR={row["IE_OR"]:.3f} [{row["CI_low"]:.3f}-{row["CI_high"]:.3f}]'
                if row.get("p") is not None and row["p"] < 0.05:
                    label += " *"
                ax.text(ax.get_xlim()[1] if ax.get_xlim()[1] > 1.1 else 1.12,
                        idx, label, va='center', fontsize=9)

            ax.axvline(x=1, color='gray', linestyle='--', alpha=0.5, zorder=1)
            ax.set_yticks(list(y_pos))
            ax.set_yticklabels(valid["stratum"].tolist())
            ax.set_xlabel("Indirect Effect OR (HOMA-IR mediated)", fontsize=11)
            ax.set_title("Stratified Mediation: Effect Modification by BMI, Sex, and Age",
                          fontsize=13, fontweight='bold')
            ax.invert_yaxis()
            plt.tight_layout()
            plt.savefig(fig_dir / "figure3_forest_plot.png",
                        dpi=config["output_dpi"], bbox_inches='tight')
            plt.close()

    print(f"  Figures saved to {fig_dir}")


# ═══════════════════════════════════════════════════════════════════════════════
# PHASE 7: RESULTS COMPILATION
# ═══════════════════════════════════════════════════════════════════════════════

def compile_results(config, mediation_result, direct_effects, stratified, work_dir):
    """Compile all results into a structured JSON report."""
    report = {
        "title": "NHANES Mediation Analysis: Inflammation → Insulin Resistance → Depression",
        "generated_at": datetime.now().isoformat(),
        "generated_by": "AI Research Army — NHANES Mediation Engine (Claw4S 2026)",
        "configuration": {
            "exposure": config["exposure"],
            "mediator": config["mediator"],
            "outcome": config["outcome"],
            "nhanes_cycles": [c["years"] for c in config["cycles"]],
            "bootstrap_iterations": config["n_bootstrap"],
        },
        "direct_effects": direct_effects.to_dict(orient='records') if direct_effects is not None else None,
        "mediation": mediation_result,
        "stratified_effects": stratified.to_dict(orient='records') if stratified is not None else None,
        "key_findings": [
            f"Indirect effect OR = {mediation_result.get('indirect_effect_OR', 'N/A')} "
            f"(95% CI: {mediation_result.get('indirect_CI_low', 'N/A')}-{mediation_result.get('indirect_CI_high', 'N/A')})",
            f"Proportion mediated: {mediation_result.get('pct_mediated', 'N/A')}%",
            f"Analysis based on N = {mediation_result.get('n', 'N/A')} participants with complete data",
        ],
        "outputs": {
            "tables": ["direct_effects.csv", "mediation_results.csv", "stratified_mediation.csv"],
            "figures": ["figure1_path_diagram.png", "figure2_dose_response.png", "figure3_forest_plot.png"],
        }
    }

    out = work_dir / "outputs" / "results_report.json"
    with open(out, 'w') as f:
        json.dump(report, f, indent=2, default=str)
    print(f"\n{'='*60}")
    print("PIPELINE COMPLETE")
    print(f"{'='*60}")
    print(f"Results report: {out}")
    print(f"\nKey Finding:")
    for finding in report["key_findings"]:
        print(f"  • {finding}")
    return report


# ═══════════════════════════════════════════════════════════════════════════════
# MAIN ORCHESTRATOR
# ═══════════════════════════════════════════════════════════════════════════════

def main():
    args = parse_args()
    config = DEFAULT_CONFIG.copy()
    config["exposure"] = args.exposure
    config["mediator"] = args.mediator
    config["outcome"] = args.outcome
    config["n_bootstrap"] = args.bootstrap
    config["seed"] = args.seed

    work_dir = Path.cwd() / "nhanes_mediation_output"
    work_dir.mkdir(exist_ok=True)

    t0 = time.time()
    print(f"\n{'='*60}")
    print(f"NHANES MEDIATION ANALYSIS ENGINE")
    print(f"Exposure: {config['exposure']} → Mediator: {config['mediator']} → Outcome: {config['outcome']}")
    print(f"{'='*60}\n")

    # Phase 1: Download
    print("[1/7] Downloading NHANES data from CDC...")
    download_nhanes(config, work_dir)

    # Phase 2: Merge
    print("[2/7] Merging cycles and constructing variables...")
    df = merge_nhanes(config, work_dir)

    # Phase 3: Direct effects
    print("[3/7] Running direct effects (3-model nested regression)...")
    direct_effects = run_direct_effects(df, config, work_dir)

    # Phase 4: Mediation
    print(f"[4/7] Running mediation analysis ({config['n_bootstrap']} bootstrap iterations)...")
    mediation_result = run_mediation(df, config, work_dir)

    # Phase 5: Stratified
    print("[5/7] Running stratified mediation (BMI × Sex × Age)...")
    stratified = run_stratified(df, config, work_dir)

    # Phase 6: Figures
    print("[6/7] Generating publication-grade figures...")
    generate_figures(df, config, mediation_result, stratified, work_dir)

    # Phase 7: Compile
    print("[7/7] Compiling results report...")
    compile_results(config, mediation_result, direct_effects, stratified, work_dir)

    elapsed = time.time() - t0
    print(f"\nTotal runtime: {elapsed/60:.1f} minutes")


if __name__ == "__main__":
    main()
```

## Expected Output

After execution, the `nhanes_mediation_output/` directory will contain:

```
nhanes_mediation_output/
├── data/
│   ├── raw/          # Downloaded NHANES XPT files (~50MB)
│   └── processed/    # Merged analytic dataset (CSV)
├── outputs/
│   ├── tables/
│   │   ├── direct_effects.csv       # 3-model nested regression
│   │   ├── mediation_results.csv    # Bootstrap mediation effects
│   │   └── stratified_mediation.csv # BMI/sex/age stratification
│   ├── figures/
│   │   ├── figure1_path_diagram.png    # SEM path diagram (300 DPI)
│   │   ├── figure2_dose_response.png   # RCS dose-response curves
│   │   └── figure3_forest_plot.png     # Stratified effects forest plot
│   └── results_report.json  # Structured summary of all findings
```

## Generalization Guide

This pipeline accepts any exposure-mediator-outcome combination available in NHANES:

```bash
# Example: Does vitamin D (exposure) reduce depression (outcome)
# through improved insulin sensitivity (mediator)?
python3 pipeline.py --exposure "log_vitd" --mediator "log_homa_ir" --outcome "depression"

# Example: Does physical activity reduce metabolic syndrome
# through lowering inflammation?
python3 pipeline.py --exposure "pa_regular" --mediator "log_crp" --outcome "mets"
```

To add new NHANES modules, extend the `module_map` dictionary in `merge_nhanes()` with the appropriate NHANES variable codes. See CDC documentation for variable names by cycle.

## Methodology Reference

- **Mediation framework:** Preacher & Hayes (2008) product-of-coefficients
- **Survey weighting:** CDC NHANES multi-cycle pooling guidelines (WTMEC2YR / n_cycles)
- **Depression assessment:** PHQ-9 ≥ 10 (Kroenke et al., 2001)
- **Insulin resistance:** HOMA-IR = (glucose × insulin) / 405
- **Dose-response:** Restricted cubic splines (4 df, natural splines)
- **Unmeasured confounding:** E-value methodology (VanderWeele & Ding, 2017)