alnn_rewrite/analysis/evaluate_models.py

# pyright: basic

from __future__ import annotations

import argparse
from pathlib import Path
from typing import Any

import pandas as pd

from analysis.analysis_modules import (
    run_calibration,
    run_longitudinal,
    run_performance,
    run_physician,
)
from analysis.data_access import load_backend_evaluation, load_clinical_table
from analysis.defaults import (
    DEFAULT_BACKENDS,
    DEFAULT_BAYESIAN_MC_PASSES,
    DEFAULT_CALIBRATION_BINS,
    DEFAULT_DECISION_THRESHOLD,
    DEFAULT_POSITIVE_CLASS_INDEX,
    noise_factor_grid,
    threshold_grid,
)
from analysis.holdout_evaluation import ensure_backend_netcdf
from analysis.noise_analysis import run_noise_analysis
from analysis.runtime import backend_dir, init_runtime_paths, load_config, write_json


def _plot_description(filename: str) -> str:
    descriptions = {
        "performance_threshold_sweep.png": "Accuracy and F1 as the decision threshold varies.",
        "performance_uncertainty_cutoff.png": "Performance while progressively restricting to lower-uncertainty predictions.",
        "performance_uncertainty_percentile_cutoff.png": "Percentile-ranked low-uncertainty subset performance from least to most restricted.",
        "calibration_reliability.png": "Reliability diagram comparing predicted probability to empirical outcome frequency.",
        "physician_confidence_boxplot.png": "Model confidence grouped by physician confidence ratings.",
        "physician_std_boxplot.png": "Model secondary uncertainty grouped by physician confidence ratings.",
        "physician_predictive_entropy_boxplot.png": "Predictive entropy grouped by physician confidence ratings.",
        "longitudinal_cohort_confidence.png": "Longitudinal cohort comparison using model confidence.",
        "longitudinal_cohort_std.png": "Longitudinal cohort comparison using ensemble standard deviation uncertainty.",
        "longitudinal_cohort_predictive_entropy.png": "Longitudinal cohort comparison using predictive entropy uncertainty.",
        "noise_sensitivity.png": "Performance metrics across increasing Gaussian noise factors.",
        "noise_uncertainty.png": "Uncertainty metrics across increasing Gaussian noise factors.",
        "noise_confidence_certainty.png": "Confidence certainty trend across increasing Gaussian noise factors.",
        "ensemble_noise_examples.png": "Representative image slices with progressively larger Gaussian noise factors.",
        "bayesian_noise_examples.png": "Representative image slices with progressively larger Gaussian noise factors.",
    }
    return descriptions.get(filename, "Generated analysis plot.")


def _write_backend_plot_report(backend: str, out_dir: Path) -> Path:
    plots_dir = out_dir / "plots"
    images = sorted(plots_dir.rglob("*.png")) if plots_dir.exists() else []

    report_path = out_dir / "plots_report.md"
    lines = [
        f"# {backend.title()} Analysis Plot Report",
        "",
        "This document lists generated analysis plots with brief descriptions.",
        "",
    ]
    if not images:
        lines.append("No plot images were generated for this backend run.")
    else:
        for image_path in images:
            rel = image_path.relative_to(out_dir).as_posix()
            title = image_path.stem.replace("_", " ").title()
            lines.append(f"## {title}")
            lines.append(_plot_description(image_path.name))
            lines.append("")
            lines.append(f"![{title}]({rel})")
            lines.append("")

    report_path.write_text("\n".join(lines), encoding="utf-8")
    return report_path


def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "Run modular evaluation analyses for ensemble and bayesian models. "
            "All outputs are written to alnn_rewrite/analysis_output."
        )
    )
    parser.add_argument(
        "--backend",
        nargs="+",
        choices=["ensemble", "bayesian"],
        default=DEFAULT_BACKENDS,
        help="Backends to evaluate.",
    )
    parser.add_argument(
        "--run-name",
        default=None,
        help="Optional run directory name under analysis_output.",
    )
    parser.add_argument(
        "--skip-noise",
        action="store_true",
        help="Skip Gaussian noise sensitivity analysis.",
    )
    return parser.parse_args()


def _run_backend(
    config: dict[str, Any],
    root_dir: Path,
    backend: str,
    clinical_df: pd.DataFrame,
    skip_noise: bool,
    out_dir: Path,
) -> dict[str, Any]:
    netcdf_path = ensure_backend_netcdf(
        config=config,
        root_dir=root_dir,
        backend=backend,
        bayesian_mc_passes=DEFAULT_BAYESIAN_MC_PASSES,
    )

    evaluation = load_backend_evaluation(
        config=config,
        backend=backend,
        class_index=DEFAULT_POSITIVE_CLASS_INDEX,
    )

    thresholds = threshold_grid()
    noise_factors = noise_factor_grid()

    summary: dict[str, Any] = {
        "backend": backend,
        "netcdf": str(netcdf_path),
        "source_file": str(evaluation.source_file),
        "uncertainty_metric": evaluation.uncertainty_metric,
    }

    summary["performance"] = run_performance(
        evaluation=evaluation,
        output_dir=out_dir,
        thresholds=thresholds,
    )
    summary["calibration"] = run_calibration(
        evaluation=evaluation,
        output_dir=out_dir,
        bins=DEFAULT_CALIBRATION_BINS,
    )
    summary["physician"] = run_physician(
        evaluation=evaluation,
        clinical_df=clinical_df,
        output_dir=out_dir,
    )
    summary["longitudinal"] = run_longitudinal(
        evaluation=evaluation,
        clinical_df=clinical_df,
        output_dir=out_dir,
    )

    if skip_noise:
        summary["noise"] = {"skipped": True, "reason": "--skip-noise supplied"}
    else:
        try:
            summary["noise"] = run_noise_analysis(
                config=config,
                root_dir=root_dir,
                backend=backend,
                output_dir=out_dir,
                class_index=DEFAULT_POSITIVE_CLASS_INDEX,
                noise_sigmas=noise_factors,
                threshold=DEFAULT_DECISION_THRESHOLD,
                calibration_bins=DEFAULT_CALIBRATION_BINS,
                bayesian_mc_passes=DEFAULT_BAYESIAN_MC_PASSES,
            )
        except Exception as exc:
            summary["noise"] = {
                "skipped": True,
                "reason": f"Noise analysis failed: {exc}",
            }

    report_path = _write_backend_plot_report(backend=backend, out_dir=out_dir)
    summary["plots_report"] = str(report_path)
    write_json(out_dir / "backend_summary.json", summary)
    return summary


def main() -> None:
    args = _parse_args()

    analysis_dir = Path(__file__).resolve().parent
    paths = init_runtime_paths(analysis_dir=analysis_dir, run_name=args.run_name)
    config = load_config(paths.root_dir)
    clinical_df = load_clinical_table(config=config, root_dir=paths.root_dir)

    manifest: dict[str, Any] = {
        "run_dir": str(paths.run_dir),
        "output_root": str(paths.output_root),
        "positive_class_index": DEFAULT_POSITIVE_CLASS_INDEX,
        "threshold_sweep": {
            "values": [float(v) for v in threshold_grid().tolist()],
        },
        "calibration_bins": DEFAULT_CALIBRATION_BINS,
        "noise_factors": noise_factor_grid(),
        "bayesian_mc_passes": DEFAULT_BAYESIAN_MC_PASSES,
        "decision_threshold": DEFAULT_DECISION_THRESHOLD,
        "backends": {},
    }

    for backend in args.backend:
        out_dir = backend_dir(paths, backend)
        manifest["backends"][backend] = _run_backend(
            config=config,
            root_dir=paths.root_dir,
            backend=backend,
            clinical_df=clinical_df,
            skip_noise=bool(args.skip_noise),
            out_dir=out_dir,
        )

    write_json(paths.run_dir / "run_manifest.json", manifest)
    print(f"Analysis complete. Results saved to {paths.run_dir}")


if __name__ == "__main__":
    main()