alnn_rewrite/tasks/analysis.py

"""Step 7: analysis over saved evaluation netCDF files.

Analysis is deliberately decoupled from models and data: it reads ONLY the
evaluation ``.nc`` files in ``<work_dir>/evaluations/`` and writes artifacts to
``<work_dir>/analysis/``. That makes an analysis rerun cheap and safe -- the only
outputs it can clobber are the easily-regeneratable files under ``analysis/``.

To add an analysis, register a callable in :data:`ANALYSES`. Each receives the
loaded evaluation datasets (keyed by file stem, e.g. ``"normal"``,
``"noisy_bayesian"``), an output directory, and a logger. Because every analysis
is a pure function of the schema-locked datasets, new plots (coverage,
accuracy-vs-noise, uncertainty-vs-noise, calibration, ...) drop in here without
touching the rest of the pipeline.
"""

import datetime
import itertools
import json
import math
import pathlib as pl
from typing import Any, Callable, Dict, List, Tuple

import matplotlib

matplotlib.use("Agg")  # headless; analysis runs in a background worker thread

import matplotlib.pyplot as plt  # noqa: E402
import numpy as np  # noqa: E402
import typst  # noqa: E402
import xarray as xr  # noqa: E402
from matplotlib.ticker import MaxNLocator  # noqa: E402

from evaluation import (  # noqa: E402
    CLASS_NAMES,
    MODEL_KIND_BAYESIAN,
    MODEL_KIND_NORMAL,
    load_evaluation,
)
from util.progress import ProgressTracker  # noqa: E402
from util.ui_logger import PipelineLogger  # noqa: E402

# Categorical colors (dataviz palette slots 1 & 2, CVD-safe in fixed order).
_FAMILY_COLORS = {
    MODEL_KIND_NORMAL: "#003f5c",  # blue
    MODEL_KIND_BAYESIAN: "#ffa600",  # orange
}

# Number of random member-subset configurations sampled per ensemble size
# (fewer only when the total number of distinct subsets is smaller).
_TARGET_CONFIGS = 30

# Bootstrap CI settings for full-ensemble accuracy.
_BOOTSTRAP_ITERS = 2000
_CI_LEVEL = 0.95

EVAL_SUBDIR = "evaluations"
ANALYSIS_SUBDIR = "analysis"

# Analysis registry: name -> function(datasets, out_dir, log) -> None.
# Each analysis is a pure function of the schema-locked evaluation datasets.
AnalysisFn = Callable[[Dict[str, xr.Dataset], pl.Path, PipelineLogger], None]

_TEMPLATE_DIR = pl.Path(__file__).parent / "templates"
_MODEL_REPORT_TEMPLATE = _TEMPLATE_DIR / "model_report.typ"

# Display labels + ordering for known model families.
_FAMILY_LABELS = {
    MODEL_KIND_NORMAL: "Normal (deterministic)",
    MODEL_KIND_BAYESIAN: "Bayesian (MC)",
}
_FAMILY_ORDER = [MODEL_KIND_NORMAL, MODEL_KIND_BAYESIAN]


def _baseline_noise_index(ds: xr.Dataset) -> int:
    """Index of the noise level closest to the clean baseline (sigma = 0)."""
    levels = np.asarray(ds["noise_level"].values, dtype=float)
    return int(np.abs(levels).argmin())


def _family_metrics(kind: str, source: str, ds: xr.Dataset) -> Dict[str, Any]:
    """Per-model + ensemble metrics for one model family at the clean baseline."""
    baseline = _baseline_noise_index(ds)
    kinds = [str(k) for k in np.atleast_1d(ds["model_kind"].values)]
    positions = [i for i, k in enumerate(kinds) if k == kind]

    correct = ds["correct"].isel(noise_level=baseline).values  # (model, sample)
    entropy = ds["predictive_entropy"].isel(noise_level=baseline).values
    mut_info = ds["mutual_information"].isel(noise_level=baseline).values
    labels = np.atleast_1d(ds["model"].values)

    models: List[Dict[str, Any]] = []
    accuracies: List[float] = []
    for pos in positions:
        acc = float(np.mean(correct[pos]))
        accuracies.append(acc)
        models.append(
            {
                "index": int(labels[pos]) + 1,  # 1-based for display
                "accuracy": acc,
                "mean_entropy": float(np.mean(entropy[pos])),
                "mean_mi": float(np.mean(mut_info[pos])),
            }
        )

    acc_arr = np.asarray(accuracies, dtype=float)
    return {
        "name": _FAMILY_LABELS.get(kind, kind),
        "kind": kind,
        "source": source,
        "noise_sigma": float(np.asarray(ds["noise_level"].values, dtype=float)[baseline]),
        "n_models": len(models),
        "n_samples": int(ds.sizes["sample"]),
        "n_mc": int(ds.attrs.get("n_mc", 1)),
        "accuracy_mean": float(acc_arr.mean()) if acc_arr.size else 0.0,
        "accuracy_std": float(acc_arr.std()) if acc_arr.size else 0.0,
        "accuracy_best": float(acc_arr.max()) if acc_arr.size else 0.0,
        "models": models,
    }


def _family_sources(
    datasets: Dict[str, xr.Dataset]
) -> Dict[str, Tuple[str, xr.Dataset]]:
    """Pick one evaluation dataset per model family, preferring clean files.

    ``normal.nc`` and ``noisy_normal.nc`` share the same members at sigma=0, so
    only the first source seen per family is kept (non-noisy stems sort first) to
    avoid analyzing a family twice.
    """
    chosen: Dict[str, Tuple[str, xr.Dataset]] = {}
    for stem in sorted(datasets, key=lambda s: ("noisy" in s, s)):
        ds = datasets[stem]
        if "model_kind" not in ds.coords:
            continue
        for kind in {str(k) for k in np.atleast_1d(ds["model_kind"].values)}:
            chosen.setdefault(kind, (stem, ds))
    return chosen


def _ordered_families(kinds: Any) -> List[str]:
    """Sort family kinds into the canonical Normal-then-Bayesian order."""
    order = {kind: i for i, kind in enumerate(_FAMILY_ORDER)}
    return sorted(kinds, key=lambda k: order.get(k, len(order)))


def _collect_families(datasets: Dict[str, xr.Dataset]) -> List[Dict[str, Any]]:
    """One metrics block per model family, in canonical order."""
    sources = _family_sources(datasets)
    families = [
        _family_metrics(kind, source, ds) for kind, (source, ds) in sources.items()
    ]
    families.sort(key=lambda f: _ordered_families(sources.keys()).index(f["kind"]))
    return families


def model_report(
    datasets: Dict[str, xr.Dataset], out_dir: pl.Path, log: PipelineLogger
) -> None:
    """Compile a short per-model PDF report (accuracy + uncertainty) via Typst."""
    families = _collect_families(datasets)
    if not families:
        log.error("model_report: no model families found in evaluations; skipping.")
        return

    meta = next(iter(datasets.values())).attrs
    seed = meta.get("seed")
    payload = {
        "title": "Model Evaluation Report",
        "generated": datetime.datetime.now(datetime.timezone.utc).strftime(
            "%Y-%m-%d %H:%M UTC"
        ),
        "work_dir": str(out_dir.parent),
        "schema_version": str(meta.get("schema_version", "?")),
        "seed": str(seed) if seed is not None else "n/a",
        "git_commit": str(meta.get("git_commit", "unknown"))[:10],
        "families": families,
    }
    payload_json = json.dumps(payload)

    # The template parses this JSON via `json(bytes(sys.inputs.data))`.
    pdf_bytes = typst.compile(
        str(_MODEL_REPORT_TEMPLATE),
        sys_inputs={"data": payload_json},
        format="pdf",
    )
    (out_dir / "model_report.pdf").write_bytes(pdf_bytes)
    # Keep the JSON alongside for debugging / reuse (regeneratable).
    (out_dir / "model_report.json").write_text(payload_json)

    total_models = sum(f["n_models"] for f in families)
    log.info(
        f"model_report: wrote model_report.pdf "
        f"({len(families)} family/families, {total_models} model(s))."
    )


def _baseline_probs(kind: str, ds: xr.Dataset) -> Tuple[np.ndarray, np.ndarray]:
    """Return ``(probs (M, S, C), true_idx (S,))`` for ``kind``'s members at the
    clean-baseline noise level."""
    baseline = _baseline_noise_index(ds)
    kinds = [str(k) for k in np.atleast_1d(ds["model_kind"].values)]
    positions = [i for i, k in enumerate(kinds) if k == kind]
    probs = np.asarray(ds["prob"].isel(noise_level=baseline).values, dtype=float)
    probs = probs[positions]  # (M, S, C)
    true_idx = np.array(
        [CLASS_NAMES.index(str(t)) for t in np.atleast_1d(ds["true_class"].values)],
        dtype=int,
    )
    return probs, true_idx


def _sample_subsets(
    n_models: int, size: int, target: int, rng: np.random.Generator
) -> List[List[int]]:
    """Distinct member-index subsets of ``size``.

    Returns every subset when the number of distinct combinations is <= ``target``
    (e.g. size == 1, size == M, or size near M); otherwise ``target`` random
    distinct subsets.
    """
    total = math.comb(n_models, size)
    if total <= target:
        return [list(combo) for combo in itertools.combinations(range(n_models), size)]

    seen: set[Tuple[int, ...]] = set()
    subsets: List[List[int]] = []
    while len(subsets) < target:
        pick = tuple(
            sorted(int(i) for i in rng.choice(n_models, size=size, replace=False))
        )
        if pick not in seen:
            seen.add(pick)
            subsets.append(list(pick))
    return subsets


def _ensemble_size_curve(
    probs: np.ndarray, true_idx: np.ndarray, target: int, rng: np.random.Generator
) -> Dict[str, Any]:
    """Sampled-configuration accuracy for every ensemble size ``1..M``.

    An ensemble prediction is the argmax of the mean of its members' output
    probabilities. For each size we average accuracy over sampled member subsets
    and report the spread (std) as the uncertainty estimate.
    """
    n_models = probs.shape[0]
    sizes: List[int] = []
    means: List[float] = []
    stds: List[float] = []
    n_configs: List[int] = []
    for size in range(1, n_models + 1):
        subsets = _sample_subsets(n_models, size, target, rng)
        accuracies = []
        for subset in subsets:
            ensemble_prob = probs[subset].mean(axis=0)  # (S, C)
            predictions = ensemble_prob.argmax(axis=1)
            accuracies.append(float(np.mean(predictions == true_idx)))
        acc = np.asarray(accuracies, dtype=float)
        sizes.append(size)
        means.append(float(acc.mean()))
        stds.append(float(acc.std()))
        n_configs.append(len(subsets))
    return {"sizes": sizes, "mean": means, "std": stds, "n_configs": n_configs}


def ensemble_size_accuracy(
    datasets: Dict[str, xr.Dataset], out_dir: pl.Path, log: PipelineLogger
) -> None:
    """Plot ensemble accuracy vs. number of member models, with an uncertainty
    band from sampling different member configurations, for each family."""
    sources = _family_sources(datasets)
    if not sources:
        log.error("ensemble_size_accuracy: no model families found; skipping.")
        return

    base_seed = int(next(iter(datasets.values())).attrs.get("seed", 0) or 0)

    curves: Dict[str, Dict[str, Any]] = {}
    for fam_index, kind in enumerate(_ordered_families(sources.keys())):
        source, ds = sources[kind]
        probs, true_idx = _baseline_probs(kind, ds)
        if probs.shape[0] < 1:
            continue
        # Distinct, reproducible RNG stream per family for subset sampling.
        rng = np.random.default_rng([base_seed, fam_index])
        curve = _ensemble_size_curve(probs, true_idx, _TARGET_CONFIGS, rng)
        curve["source"] = source
        curves[kind] = curve
        log.info(
            f"ensemble_size_accuracy: {kind} - sizes 1..{probs.shape[0]}, "
            f"up to {_TARGET_CONFIGS} configs each."
        )

    if not curves:
        log.error("ensemble_size_accuracy: no usable probabilities; skipping.")
        return

    fig, ax = plt.subplots(figsize=(7.0, 4.3))
    for kind in _ordered_families(curves.keys()):
        curve = curves[kind]
        color = _FAMILY_COLORS.get(kind, "#666666")
        x = np.asarray(curve["sizes"])
        mean = np.asarray(curve["mean"])
        std = np.asarray(curve["std"])
        ax.fill_between(x, mean - std, mean + std, color=color, alpha=0.18, linewidth=0)
        ax.plot(
            x, mean, color=color, marker="o", markersize=4, linewidth=2,
            label=_FAMILY_LABELS.get(kind, kind),
        )

    ax.set_xlabel("Ensemble size (number of models)")
    ax.set_ylabel("Test accuracy")
    ax.set_title("Ensemble accuracy vs. number of models")
    ax.xaxis.set_major_locator(MaxNLocator(integer=True))
    ax.margins(x=0.02)
    ax.grid(True, color="#000000", alpha=0.08, linewidth=0.7)
    for spine in ("top", "right"):
        ax.spines[spine].set_visible(False)
    ax.legend(frameon=False, loc="lower right")
    fig.tight_layout()

    png_path = out_dir / "ensemble_size_accuracy.png"
    fig.savefig(png_path, dpi=150)
    fig.savefig(out_dir / "ensemble_size_accuracy.pdf")
    plt.close(fig)

    # Persist the underlying curve data for reuse / regeneration.
    (out_dir / "ensemble_size_accuracy.json").write_text(
        json.dumps({"target_configs": _TARGET_CONFIGS, "families": curves})
    )
    log.info(
        f"ensemble_size_accuracy: wrote {png_path.name} "
        f"({len(curves)} family/families)."
    )


def _full_ensemble_correct(kind: str, ds: xr.Dataset) -> Tuple[np.ndarray, np.ndarray]:
    """Per-sample correctness (0/1) of the FULL ensemble at baseline, plus ptids.

    The ensemble prediction is the argmax of the mean of ALL members' output
    probabilities -- i.e. the rightmost point of the ensemble-size curve.
    """
    probs, true_idx = _baseline_probs(kind, ds)  # (M, S, C), (S,)
    ensemble_pred = probs.mean(axis=0).argmax(axis=1)  # (S,)
    correct = (ensemble_pred == true_idx).astype(float)
    ptids = np.asarray(ds["ptid"].values).astype(str)
    return correct, ptids


def _bootstrap_accuracy_ci(
    correct: np.ndarray,
    ptids: np.ndarray,
    iters: int,
    ci_level: float,
    rng: np.random.Generator,
) -> Dict[str, Any]:
    """Patient-clustered percentile bootstrap CI for a mean-accuracy statistic.

    Resamples whole PATIENTS with replacement (not individual scans), because
    scans from one patient are correlated; a per-scan bootstrap would understate
    the interval. Reduces to an ordinary bootstrap when every patient has one scan.
    """
    point = float(np.mean(correct)) if correct.size else 0.0
    _, inverse = np.unique(ptids, return_inverse=True)
    n_patients = int(inverse.max()) + 1 if inverse.size else 0
    groups = [np.where(inverse == p)[0] for p in range(n_patients)]

    boot = np.empty(iters, dtype=float)
    for b in range(iters):
        chosen = rng.integers(0, n_patients, size=n_patients)
        idx = np.concatenate([groups[c] for c in chosen])
        boot[b] = correct[idx].mean()

    tail = (1.0 - ci_level) / 2.0
    return {
        "point": point,
        "ci_low": float(np.quantile(boot, tail)),
        "ci_high": float(np.quantile(boot, 1.0 - tail)),
        "n_samples": int(correct.size),
        "n_patients": n_patients,
    }


def bootstrap_ensemble_ci(
    datasets: Dict[str, xr.Dataset], out_dir: pl.Path, log: PipelineLogger
) -> None:
    """Patient-clustered bootstrap CI for each family's full-ensemble test
    accuracy; writes a JSON summary and a compact point-range plot."""
    sources = _family_sources(datasets)
    if not sources:
        log.error("bootstrap_ensemble_ci: no model families found; skipping.")
        return

    base_seed = int(next(iter(datasets.values())).attrs.get("seed", 0) or 0)

    results: Dict[str, Dict[str, Any]] = {}
    for fam_index, kind in enumerate(_ordered_families(sources.keys())):
        source, ds = sources[kind]
        correct, ptids = _full_ensemble_correct(kind, ds)
        if correct.size == 0:
            continue
        rng = np.random.default_rng([base_seed, 1000 + fam_index])
        ci = _bootstrap_accuracy_ci(correct, ptids, _BOOTSTRAP_ITERS, _CI_LEVEL, rng)
        ci["source"] = source
        results[kind] = ci
        log.info(
            f"bootstrap_ensemble_ci: {kind} full-ensemble acc = "
            f"{ci['point'] * 100:.1f}% (95% CI "
            f"{ci['ci_low'] * 100:.1f}-{ci['ci_high'] * 100:.1f}%, "
            f"n={ci['n_samples']} scans / {ci['n_patients']} patients)."
        )

    if not results:
        log.error("bootstrap_ensemble_ci: nothing to compute; skipping.")
        return

    ordered = _ordered_families(results.keys())
    fig, ax = plt.subplots(figsize=(7.0, 1.4 + 0.8 * len(ordered)))
    for row, kind in enumerate(ordered):
        ci = results[kind]
        color = _FAMILY_COLORS.get(kind, "#666666")
        ax.plot(
            [ci["ci_low"], ci["ci_high"]], [row, row],
            color=color, linewidth=2.5, solid_capstyle="round",
        )
        ax.plot([ci["point"]], [row], "o", color=color, markersize=8)
        ax.annotate(
            f"{ci['point'] * 100:.1f}%  "
            f"[{ci['ci_low'] * 100:.1f}, {ci['ci_high'] * 100:.1f}]",
            (ci["point"], row), xytext=(0, 10), textcoords="offset points",
            ha="center", va="bottom", fontsize=9, color="#222222",
        )
    ax.set_yticks(range(len(ordered)))
    ax.set_yticklabels([_FAMILY_LABELS.get(k, k) for k in ordered])
    ax.invert_yaxis()  # first family on top
    ax.set_xlabel("Full-ensemble test accuracy (95% patient-clustered bootstrap CI)")
    ax.set_title("Ensemble accuracy with bootstrap confidence intervals")
    ax.margins(x=0.18, y=0.35)
    ax.grid(True, axis="x", color="#000000", alpha=0.08, linewidth=0.7)
    for spine in ("top", "right", "left"):
        ax.spines[spine].set_visible(False)
    ax.tick_params(left=False)
    fig.tight_layout()
    fig.savefig(out_dir / "bootstrap_ci.png", dpi=150)
    fig.savefig(out_dir / "bootstrap_ci.pdf")
    plt.close(fig)

    (out_dir / "bootstrap_ci.json").write_text(
        json.dumps(
            {"iters": _BOOTSTRAP_ITERS, "ci_level": _CI_LEVEL, "families": results}
        )
    )
    log.info(
        f"bootstrap_ensemble_ci: wrote bootstrap_ci.png ({len(results)} family/families)."
    )


ANALYSES: Dict[str, AnalysisFn] = {
    "model_report": model_report,
    "ensemble_size_accuracy": ensemble_size_accuracy,
    "bootstrap_ensemble_ci": bootstrap_ensemble_ci,
}


def _load_evaluations(
    eval_dir: pl.Path, log: PipelineLogger
) -> Dict[str, xr.Dataset]:
    """Load every ``*.nc`` in ``eval_dir``, keyed by file stem."""
    datasets: Dict[str, xr.Dataset] = {}
    for nc in sorted(eval_dir.glob("*.nc")):
        try:
            ds = load_evaluation(str(nc))
        except Exception as exc:  # noqa: BLE001 - report and skip a bad file
            log.error(f"Could not load evaluation '{nc.name}': {exc}")
            continue
        datasets[nc.stem] = ds
        log.info(f"Loaded evaluation '{nc.stem}' with dims {dict(ds.sizes)}.")
    return datasets


def run_analysis_task(
    track: ProgressTracker,
    log: PipelineLogger,
    config: Dict[str, Any],
    state: Dict[str, Any],
) -> None:
    work_dir = pl.Path(config["work_dir"])
    eval_dir = work_dir / EVAL_SUBDIR
    out_dir = work_dir / ANALYSIS_SUBDIR

    if not eval_dir.is_dir():
        log.error(
            f"No evaluations directory at {eval_dir}. Run an evaluation scenario "
            "first (or point the working directory at one that has evaluations)."
        )
        return

    datasets = _load_evaluations(eval_dir, log)
    if not datasets:
        log.error(f"No evaluation .nc files found in {eval_dir}; nothing to analyze.")
        return

    out_dir.mkdir(parents=True, exist_ok=True)

    try:
        if not ANALYSES:
            log.info(
                f"[NOT IMPLEMENTED] No analyses registered yet (step 7). "
                f"Loaded {len(datasets)} evaluation file(s): "
                f"{', '.join(sorted(datasets))}. Register functions in "
                f"tasks.analysis.ANALYSES to produce artifacts under {out_dir}/."
            )
            return

        track.update(total=len(ANALYSES), advance=0)
        for name, analysis_fn in ANALYSES.items():
            log.info(f"Running analysis: {name}")
            analysis_fn(datasets, out_dir, log)
            track.update(advance=1)
        log.info(f"Analysis artifacts written to {out_dir}/.")
    finally:
        for ds in datasets.values():
            ds.close()