extract_data.py

#!/usr/bin/env python3
"""
Extract all data and model results from the mixed-effects notebook
and save them as self-contained files in ./data/ for the Streamlit dashboard.

Run once on the cluster, then the entire Stroke_Dashboard/ directory can be
moved to any machine.
"""

import warnings, json, pickle
from pathlib import Path
from itertools import combinations

import numpy as np
import pandas as pd
from scipy.special import expit
import statsmodels.formula.api as smf

OUT = Path(__file__).parent / "data"
OUT.mkdir(exist_ok=True)

# ── 1. Discover CSVs ─────────────────────────────────────────────────────────
RUNS_ROOT = Path("/home/rbielski/stroke_cleaned/ARC_ATLAS_Combined/ARC_ATLAS_Train_v3/runs")
RUN_DIRS = sorted(RUNS_ROOT.glob("*/test_eval"))
assert RUN_DIRS, f"No test_eval directories found under {RUNS_ROOT}"

def latest_csv_any(pattern, img_path_hint=None):
    matches = sorted(path for run in RUN_DIRS for path in run.glob(pattern))
    if not matches:
        return None
    if img_path_hint is None:
        return matches[-1]
    hinted = []
    for path in matches:
        try:
            probe = pd.read_csv(path, usecols=["img_path"], nrows=5)
        except Exception:
            continue
        if probe["img_path"].astype(str).str.contains(img_path_hint, regex=False).any():
            hinted.append(path)
    return hinted[-1] if hinted else matches[-1]

VARIANT_GLOBS = [
    {"variant": "hires", "cohort": "hires", "role": "natural",
     "glob": "test_hires_metrics_with_manifest_*.csv"},
    {"variant": "lower_resolution", "cohort": "lores", "role": "holdout",
     "glob": "test_lores_metrics_with_manifest_*.csv", "img_path_hint": "/test_lores/"},
]
for family, prefix in [
    ("crude", "crude"),
    ("thick_slice", "thickslices"),
    ("inplane_coarsening", "inplane"),
    ("reduced_snr", "reducedsnr"),
    ("rigid_jitter", "rigidjitter"),
]:
    for level in range(1, 6):
        VARIANT_GLOBS.append({
            "variant": f"{family}_v{level}",
            "cohort": "hires",
            "role": "degraded",
            "glob": f"test_{prefix}_v{level}_metrics_with_manifest_*.csv",
        })

FILE_SPECS = []
for entry in VARIANT_GLOBS:
    path = latest_csv_any(entry["glob"], entry.get("img_path_hint"))
    if path is None:
        print(f"[warn] no CSV for variant '{entry['variant']}'; skipping.")
        continue
    FILE_SPECS.append({
        "variant": entry["variant"], "cohort": entry["cohort"],
        "role": entry["role"], "path": str(path),
    })

# ── 2. Load & merge ──────────────────────────────────────────────────────────
REQUIRED = ["key", "soft_dice", "img_path", "mf_fwhm_mm", "mf_hi_freq_energy", "mf_lap_var"]

dfs = {}
for spec in FILE_SPECS:
    df = pd.read_csv(spec["path"])
    if "dataset_tag" not in df.columns:
        df["dataset_tag"] = spec["cohort"]
    df["cohort"] = spec["cohort"]
    df["variant"] = spec["variant"]
    df["role"] = spec["role"]
    dfs[spec["variant"]] = df

all_df = pd.concat(dfs.values(), ignore_index=True)

# Image-level quality metrics
QM_DIR = Path("/home/rbielski/stroke_cleaned/ARC_ATLAS_Combined/ARC_ATLAS_Train_v3/Image_Quality_Metrics")
qm_csvs = sorted(QM_DIR.glob("image_quality_metrics_all_variants_*.csv"))
if qm_csvs:
    qm = pd.read_csv(qm_csvs[-1])
    QM_MAP = {"test_hires": "hires", "test_lores": "lower_resolution"}
    for fi, fo in [("Crude","crude"),("ThickSlices","thick_slice"),
                   ("InPlaneCoarse","inplane_coarsening"),("ReducedSNR","reduced_snr"),
                   ("RigidJitter","rigid_jitter")]:
        for lv in range(1,6):
            QM_MAP[f"{fi}_v{lv}"] = f"{fo}_v{lv}"
    qm["variant"] = qm["variant"].map(QM_MAP).fillna(qm["variant"])
    qm["key"] = qm["key"].astype(str)
    qm = qm.rename(columns={"fwhm_mm":"mf_fwhm_mm","hi_freq_energy":"mf_hi_freq_energy","lap_var":"mf_lap_var"})
    qm = qm[["key","variant","mf_fwhm_mm","mf_hi_freq_energy","mf_lap_var"]]
    for c in ["mf_fwhm_mm","mf_hi_freq_energy","mf_lap_var"]:
        if c in all_df.columns:
            all_df = all_df.drop(columns=[c])
    all_df = all_df.merge(qm, on=["key","variant"], how="left")

# Lesion volume
MANIFEST = Path("/home/rbielski/stroke_cleaned/ARC_ATLAS_Combined/ARC_ATLAS_Train_v4/data/splits/50_25_25/meta/_resolution_manifest_v2.csv")
if MANIFEST.exists():
    mf = pd.read_csv(MANIFEST)[["key","mask_ml_clean"]]
    mf["lesion_mm3"] = mf["mask_ml_clean"] * 1000
    mf = mf[["key","lesion_mm3"]]
    if "lesion_mm3" in all_df.columns:
        all_df = all_df.drop(columns=["lesion_mm3"])
    all_df = all_df.merge(mf, on="key", how="left")

# ── 3. Build design matrix (same logic as notebook) ──────────────────────────
DEGRADED_VARIANTS = [s["variant"] for s in FILE_SPECS if s["role"] in ("natural","degraded")]
HOLDOUT_VARIANTS = [s["variant"] for s in FILE_SPECS if s["role"] == "holdout"]
OUTCOME = "soft_dice"
GROUP_COL = "mf_key" if ("mf_key" in all_df.columns and all_df["mf_key"].notna().all()) else "key"
MODEL_QUALITY_COLS = ["mf_fwhm_mm","mf_lap_var"]
ALL_QUALITY_COLS = ["mf_fwhm_mm","mf_hi_freq_energy","mf_lap_var"]

needed = [OUTCOME,"cohort","variant","role","dataset_tag","key",GROUP_COL,"img_path","lesion_mm3"] + ALL_QUALITY_COLS
if "mf_bin" in all_df.columns:
    needed.append("mf_bin")
work_df = all_df[[c for c in needed if c in all_df.columns]].copy()
for col in [OUTCOME] + ALL_QUALITY_COLS:
    work_df[col] = pd.to_numeric(work_df[col], errors="coerce")
work_df[GROUP_COL] = work_df[GROUP_COL].astype(str)
work_df["key"] = work_df["key"].astype(str)

if "lesion_mm3" in work_df.columns:
    work_df["log_lesion_mm3"] = np.log1p(work_df["lesion_mm3"].clip(lower=0))
else:
    work_df["log_lesion_mm3"] = np.nan

# Misalign flag
variant_hint = work_df["variant"].astype(str).str.lower().str.contains("jitter|rigid|misalign|shift", regex=True)
if "img_path" in work_df.columns:
    img_hint = work_df["img_path"].astype(str).str.lower().str.contains("jitter|rigid|misalign|shift", regex=True)
else:
    img_hint = pd.Series(False, index=work_df.index)
work_df["misalign"] = (variant_hint | img_hint).astype(int)

train_df = work_df[work_df["variant"].isin(DEGRADED_VARIANTS)].copy()
holdout_df = work_df[work_df["variant"].isin(HOLDOUT_VARIANTS)].copy()
train_df = train_df.dropna(subset=[OUTCOME, GROUP_COL] + MODEL_QUALITY_COLS).copy()
holdout_df = holdout_df.dropna(subset=[OUTCOME] + MODEL_QUALITY_COLS).copy()

# Standardise
scaler = {}
for col in MODEL_QUALITY_COLS:
    mu = float(train_df[col].mean())
    sd = float(train_df[col].std(ddof=0))
    if not np.isfinite(sd) or sd < 1e-12:
        sd = 1.0
    scaler[col] = {"mean": mu, "std": sd}
    zcol = f"z_{col}"
    train_df[zcol] = (train_df[col] - mu) / sd
    holdout_df[zcol] = (holdout_df[col] - mu) / sd

les_mu = float(train_df["log_lesion_mm3"].mean())
les_sd = float(train_df["log_lesion_mm3"].std(ddof=0)) or 1.0
scaler["log_lesion_mm3"] = {"mean": les_mu, "std": les_sd}
train_df["z_log_lesion_mm3"] = (train_df["log_lesion_mm3"] - les_mu) / les_sd
holdout_df["z_log_lesion_mm3"] = (holdout_df["log_lesion_mm3"] - les_mu) / les_sd

# Logit-transform
eps = 1e-6
for df in (train_df, holdout_df):
    clipped = np.clip(df[OUTCOME].astype(float).values, eps, 1 - eps)
    df["logit_dice"] = np.log(clipped / (1 - clipped))

# Family column
def _variant_family(variant):
    text = str(variant)
    if text in {"hires","lower_resolution"}: return text
    for prefix, family in [("crude_v","crude"),("thick_slice_v","thick_slice"),
                           ("inplane_coarsening_v","inplane_coarsening"),
                           ("reduced_snr_v","reduced_snr"),("rigid_jitter_v","rigid_jitter")]:
        if text.startswith(prefix): return family
    return text

train_df["family"] = train_df["variant"].map(_variant_family)
holdout_df["family"] = holdout_df["variant"].map(_variant_family)

# ── 4. Fit all 7 models ─────────────────────────────────────────────────────
compare_df = train_df.copy()

def fit_mixedlm(formula, data):
    model = smf.mixedlm(formula=formula, data=data, groups=data[GROUP_COL], re_formula="1")
    for method in ["bfgs","cg","powell","nm","lbfgs"]:
        try:
            with warnings.catch_warnings(record=True):
                warnings.simplefilter("always")
                result = model.fit(reml=False, method=method, maxiter=2000, disp=False)
            if result.converged:
                return result, method
        except Exception:
            pass
    raise RuntimeError(f"No optimizer converged for: {formula}")

FAMILY_LABELS = {
    "hires": "Natural High-Quality",
    "crude": "Crude Downsample",
    "thick_slice": "Thick Slices",
    "inplane_coarsening": "In-Plane Coarse",
    "reduced_snr": "Reduced SNR",
    "rigid_jitter": "Rigid Jitter",
    "lower_resolution": "Natural Low-Quality",
}

model_specs = [
    {"model":"FWHM only",
     "formula":"logit_dice ~ z_mf_fwhm_mm",
     "role":"Unadjusted continuous blur model",
     "group":"Building up","has_holdout":True,
     "has_lapvar":False,"has_lesion":False,"has_misalign":False},
    {"model":"FWHM + lesion",
     "formula":"logit_dice ~ z_mf_fwhm_mm + z_log_lesion_mm3",
     "role":"Blur adjusted for lesion volume",
     "group":"Building up","has_holdout":True,
     "has_lapvar":False,"has_lesion":True,"has_misalign":False},
    {"model":"FWHM + lesion + misalign",
     "formula":"logit_dice ~ z_mf_fwhm_mm + z_log_lesion_mm3 + misalign",
     "role":"Blur plus lesion volume, with misalignment flag",
     "group":"Building up","has_holdout":True,
     "has_lapvar":False,"has_lesion":True,"has_misalign":True},
    {"model":"FWHM + LapVar",
     "formula":"logit_dice ~ z_mf_fwhm_mm + z_mf_lap_var",
     "role":"Add LapVar to FWHM without lesion volume",
     "group":"LapVar models","has_holdout":True,
     "has_lapvar":True,"has_lesion":False,"has_misalign":False},
    {"model":"FWHM + LapVar + lesion",
     "formula":"logit_dice ~ z_mf_fwhm_mm + z_mf_lap_var + z_log_lesion_mm3",
     "role":"Continuous quality model with lesion volume",
     "group":"LapVar models","has_holdout":True,
     "has_lapvar":True,"has_lesion":True,"has_misalign":False},
    {"model":"FWHM + LapVar + lesion + misalign",
     "formula":"logit_dice ~ z_mf_fwhm_mm + z_mf_lap_var + z_log_lesion_mm3 + misalign",
     "role":"Continuous quality model with lesion volume and misalignment",
     "group":"LapVar models","has_holdout":True,
     "has_lapvar":True,"has_lesion":True,"has_misalign":True},
    {"model":"FWHM x LapVar + lesion + misalign",
     "formula":"logit_dice ~ z_mf_fwhm_mm * z_mf_lap_var + z_log_lesion_mm3 + misalign",
     "role":"Full interaction model: blur by edge-energy plus lesion and misalignment",
     "group":"Interaction","has_holdout":True,
     "has_lapvar":True,"has_lesion":True,"has_misalign":True},
]

compare_results = {}
model_summaries = []

for spec in model_specs:
    print(f"Fitting: {spec['model']}...")
    result, method = fit_mixedlm(spec["formula"], compare_df)
    compare_results[spec["model"]] = result

    fe = result.fe_params
    ci = result.conf_int()
    ci.columns = ["ci_low","ci_high"]
    pvals = result.pvalues

    group_var = float(result.cov_re.iloc[0,0])
    resid_var = float(result.scale)
    icc = group_var / (group_var + resid_var)

    # R-squared (Nakagawa)
    fe_vals = np.array(result.model.exog @ result.fe_params, dtype=float)
    var_fixed = float(np.var(fe_vals))
    var_total = var_fixed + group_var + resid_var
    r2_marginal = var_fixed / var_total
    r2_conditional = (var_fixed + group_var) / var_total

    # Coefficient table
    coef_rows = []
    for term in fe.index:
        coef_rows.append({
            "term": term,
            "coef": float(fe[term]),
            "ci_low": float(ci.loc[term, "ci_low"]),
            "ci_high": float(ci.loc[term, "ci_high"]),
            "p_value": float(pvals.get(term, np.nan)),
            "is_intercept": term == "Intercept",
            "is_interaction": ":" in term,
        })

    # Holdout predictions
    holdout_metrics = {}
    if spec["has_holdout"]:
        ho = holdout_df.copy()
        try:
            ho["pred_logit"] = result.predict(ho)
            ho["pred_dice"] = expit(ho["pred_logit"])
            obs = ho[OUTCOME].to_numpy(float)
            pred = ho["pred_dice"].to_numpy(float)
            res = obs - pred
            holdout_metrics = {
                "MAE": float(np.mean(np.abs(res))),
                "RMSE": float(np.sqrt(np.mean(res**2))),
                "r": float(np.corrcoef(obs, pred)[0,1]),
                "Bias": float(np.mean(res)),
            }
        except Exception as e:
            print(f"  Holdout prediction failed for {spec['model']}: {e}")

    # Fitted values for diagnostics
    fitted_logit = result.fittedvalues
    residuals = result.resid

    # Random effects
    re_dict = result.random_effects  # {group_label: Series}
    re_vals = {str(k): float(v.iloc[0]) for k, v in re_dict.items()}

    model_summaries.append({
        "model": spec["model"],
        "formula": spec["formula"],
        "role": spec["role"],
        "group": spec["group"],
        "has_holdout": spec["has_holdout"],
        "has_lapvar": spec.get("has_lapvar", False),
        "has_lesion": spec.get("has_lesion", False),
        "has_misalign": spec.get("has_misalign", False),
        "optimizer": method,
        "AIC": float(result.aic),
        "BIC": float(result.bic),
        "logLik": float(result.llf),
        "ICC": icc,
        "group_var": group_var,
        "resid_var": resid_var,
        "R2_marginal": r2_marginal,
        "R2_conditional": r2_conditional,
        "n_fixed": len(fe) - 1,
        "coefficients": coef_rows,
        "holdout_metrics": holdout_metrics,
        "random_effects": re_vals,
        "fitted_logit": fitted_logit.tolist(),
        "residuals": residuals.tolist(),
    })
    print(f"  AIC={result.aic:.1f}  ICC={icc:.3f}  R2m={r2_marginal:.3f}  R2c={r2_conditional:.3f}")

# ── 5. Holdout predictions per model ─────────────────────────────────────────
holdout_preds = {}
for spec in model_specs:
    if not spec["has_holdout"]:
        continue
    result = compare_results[spec["model"]]
    ho = holdout_df.copy()
    try:
        ho["pred_logit"] = result.predict(ho)
        ho["pred_dice"] = expit(ho["pred_logit"])
        holdout_preds[spec["model"]] = ho[["key","variant","soft_dice","pred_dice","pred_logit"]].copy()
    except Exception:
        pass

# ── 6. Misalignment contrasts ──────────────────────────────────────────────
misalign_pairs = [
    ("FWHM + lesion", "FWHM + lesion + misalign"),
    ("FWHM + LapVar + lesion", "FWHM + LapVar + lesion + misalign"),
]
misalignment_contrasts = []
ms_lookup = {m["model"]: m for m in model_summaries}
for without_name, with_name in misalign_pairs:
    if without_name not in ms_lookup or with_name not in ms_lookup:
        continue
    wo = ms_lookup[without_name]
    wi = ms_lookup[with_name]
    misalignment_contrasts.append({
        "comparison": f"{with_name} vs {without_name}",
        "without_model": without_name,
        "with_model": with_name,
        "delta_AIC": wi["AIC"] - wo["AIC"],
        "delta_BIC": wi["BIC"] - wo["BIC"],
        "delta_holdout_MAE": (wi["holdout_metrics"].get("MAE", float("nan"))
                              - wo["holdout_metrics"].get("MAE", float("nan"))),
        "delta_R2_marginal": wi["R2_marginal"] - wo["R2_marginal"],
        "delta_R2_conditional": wi["R2_conditional"] - wo["R2_conditional"],
        "delta_ICC": wi["ICC"] - wo["ICC"],
    })

# ── 7. Save everything ──────────────────────────────────────────────────────
# DataFrames
train_df.to_csv(OUT / "train_df.csv", index=False)
holdout_df.to_csv(OUT / "holdout_df.csv", index=False)
all_df_export = all_df.copy()
# Add family column to all_df for dashboard use
all_df_export["family"] = all_df_export["variant"].map(_variant_family)
all_df_export.to_csv(OUT / "all_df.csv", index=False)

# Per-variant summary
variant_order = [v for v in DEGRADED_VARIANTS + HOLDOUT_VARIANTS if v in set(all_df["variant"].unique())]
summary_rows = []
for v in variant_order:
    source = holdout_df if v in HOLDOUT_VARIANTS else train_df
    arr = source[source["variant"] == v]["soft_dice"].dropna().values
    if len(arr) == 0:
        continue
    summary_rows.append({
        "variant": v, "family": _variant_family(v),
        "label": FAMILY_LABELS.get(_variant_family(v), v),
        "role": "Holdout" if v in HOLDOUT_VARIANTS else "Training",
        "n": len(arr),
        "median_dice": float(np.median(arr)),
        "mean_dice": float(np.mean(arr)),
        "std_dice": float(np.std(arr)),
        "q25": float(np.percentile(arr, 25)),
        "q75": float(np.percentile(arr, 75)),
        "pct_zero": float(100.0 * np.mean(arr == 0)),
    })
pd.DataFrame(summary_rows).to_csv(OUT / "variant_summary.csv", index=False)

# Model results (JSON-serializable)
with open(OUT / "model_summaries.json", "w") as f:
    json.dump(model_summaries, f, indent=2)

# Holdout predictions
for name, df in holdout_preds.items():
    safe = name.replace(" ", "_").replace("×", "x")
    df.to_csv(OUT / f"holdout_pred_{safe}.csv", index=False)

# Misalignment contrasts
with open(OUT / "misalignment_contrasts.json", "w") as f:
    json.dump(misalignment_contrasts, f, indent=2)

# Scaler info
with open(OUT / "scaler.json", "w") as f:
    json.dump(scaler, f, indent=2)

# Metadata
meta = {
    "GROUP_COL": GROUP_COL,
    "OUTCOME": OUTCOME,
    "DEGRADED_VARIANTS": DEGRADED_VARIANTS,
    "HOLDOUT_VARIANTS": HOLDOUT_VARIANTS,
    "FAMILY_LABELS": FAMILY_LABELS,
    "MODEL_QUALITY_COLS": MODEL_QUALITY_COLS,
    "ALL_QUALITY_COLS": ALL_QUALITY_COLS,
    "model_names": [s["model"] for s in model_specs],
}
with open(OUT / "meta.json", "w") as f:
    json.dump(meta, f, indent=2)

print(f"\nAll data saved to {OUT.resolve()}")
print("Files:", sorted(p.name for p in OUT.iterdir()))