"""Compute 1,000-sample bootstrap confidence intervals for the test-set metrics.

Reads outputs/predictions_test.csv and outputs/metrics_summary.csv.
For every system column in the predictions file, resamples the test set
1,000 times with replacement, recomputes all five metrics on each resample,
and reports the 2.5th and 97.5th percentiles as a 95% CI.

Output:
    outputs/metrics_with_ci.csv -- one row per system, point estimate + CI
"""
from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np
import pandas as pd

import sys
sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent))
from src.meta.train_eval import all_metrics


def bootstrap_metrics(
    y_true: np.ndarray,
    y_prob: np.ndarray,
    n_resamples: int = 1000,
    seed: int = 42,
) -> dict[str, tuple[float, float]]:
    """Return {metric: (lo, hi)} for the 95% bootstrap CI of each metric."""
    rng = np.random.default_rng(seed)
    n = len(y_true)
    samples: dict[str, list[float]] = {
        "accuracy": [], "f1_macro": [], "auroc": [],
        "fpr_at_tpr_95": [], "fpr_at_tpr_99": [], "ece": [],
    }
    for _ in range(n_resamples):
        idx = rng.integers(0, n, size=n)
        try:
            m = all_metrics(y_true[idx], y_prob[idx])
        except Exception:
            # Resample may produce a single-class subset; skip.
            continue
        for k, v in m.items():
            samples[k].append(v)

    cis = {}
    for k, vals in samples.items():
        if len(vals) < 10:
            cis[k] = (float("nan"), float("nan"))
        else:
            arr = np.array(vals)
            cis[k] = (float(np.percentile(arr, 2.5)), float(np.percentile(arr, 97.5)))
    return cis


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--outdir", type=str, default="outputs")
    parser.add_argument("--n-resamples", type=int, default=1000)
    parser.add_argument("--seed", type=int, default=42)
    args = parser.parse_args()

    outdir = Path(args.outdir)
    summary = pd.read_csv(outdir / "metrics_summary.csv")
    predictions = pd.read_csv(outdir / "predictions_test.csv")
    y_true = predictions["label"].to_numpy()

    # Discover which systems have predictions stored.
    system_cols = [c for c in predictions.columns
                   if c not in {"id", "domain", "generator", "label"}]

    rows = []
    for system in system_cols:
        if system not in summary["system"].values:
            print(f"  skip {system}: not in summary")
            continue
        print(f"  {system}: bootstrapping {args.n_resamples} resamples...")
        y_prob = predictions[system].to_numpy()
        cis = bootstrap_metrics(y_true, y_prob, args.n_resamples, args.seed)
        point = summary[summary["system"] == system].iloc[0]
        row = {"system": system}
        for metric in ["accuracy", "f1_macro", "auroc",
                       "fpr_at_tpr_95", "fpr_at_tpr_99", "ece"]:
            row[metric] = float(point[metric])
            row[f"{metric}_ci_lo"], row[f"{metric}_ci_hi"] = cis[metric]
        rows.append(row)

    out_df = pd.DataFrame(rows)
    out_df.to_csv(outdir / "metrics_with_ci.csv", index=False)
    print(f"\nSaved {outdir / 'metrics_with_ci.csv'}")

    # Pretty-print headline metrics.
    print("\n=== METRICS WITH 95% CI ===")
    for _, r in out_df.iterrows():
        print(f"\n{r['system']}")
        for metric in ["auroc", "fpr_at_tpr_95", "fpr_at_tpr_99", "ece"]:
            point = r[metric]
            lo = r[f"{metric}_ci_lo"]
            hi = r[f"{metric}_ci_hi"]
            print(f"  {metric:>15}: {point:.4f}  [{lo:.4f}, {hi:.4f}]")


if __name__ == "__main__":
    main()