"""Confidence intervals for full-ensemble accuracy.

Patient-clustered percentile bootstrap: whole PATIENTS are resampled with
replacement, not individual scans, because scans from one patient are correlated
and a per-scan bootstrap would understate the interval. Holds the trained models
fixed, so it captures test-set sampling noise only -- an honest lower bound on
total uncertainty (it excludes training/initialisation variability).
"""

from typing import Any, Dict

import numpy as np

from analysis.plotting import panel_grid, save_figure, save_json, style_axes
from analysis.sources import FAMILY_LABELS, ordered_families
from analysis.context import AnalysisContext, FAMILY_COLORS
from analysis.registry import register

BOOTSTRAP_ITERS = 2000
CI_LEVEL = 0.95


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.

    Reduces to an ordinary bootstrap when every patient has exactly 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)
        boot[b] = correct[np.concatenate([groups[c] for c in chosen])].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,
    }


@register(
    "bootstrap_ensemble_ci",
    title="Bootstrap confidence intervals for ensemble accuracy",
)
def bootstrap_ensemble_ci(ctx: AnalysisContext) -> None:
    base_seed = ctx.base_seed()

    results: Dict[str, Dict[str, Any]] = {}
    for index, family in enumerate(ctx.families):
        selected = ctx.clean_source(family)
        if selected is None:
            continue
        stem, ds = selected
        probs = ctx.probs(family, ds)
        if probs.shape[0] < 1:
            continue
        # Full-ensemble prediction: argmax of the mean over ALL members.
        correct = (
            probs.mean(axis=0).argmax(axis=1) == ctx.true_index(ds)
        ).astype(float)
        rng = np.random.default_rng([base_seed, 1000 + index])
        ci = _bootstrap_accuracy_ci(
            correct, ctx.ptids(ds), BOOTSTRAP_ITERS, CI_LEVEL, rng
        )
        ci["source"] = stem
        results[family] = ci
        ctx.log.info(
            f"bootstrap_ensemble_ci: {family} full-ensemble acc = "
            f"{ci['point'] * 100:.1f}% (95% CI {ci['ci_low'] * 100:.1f}-"
            f"{ci['ci_high'] * 100:.1f}%, n={ci['n_samples']} scans / "
            f"{ci['n_patients']} patients)."
        )

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

    ordered = ordered_families(results.keys())
    fig, axes = panel_grid(
        1, n_cols=1, panel_size=(7.0, 1.4 + 0.8 * len(ordered))
    )
    ax = axes[0]
    for row, family in enumerate(ordered):
        ci = results[family]
        color = FAMILY_COLORS.get(family, "#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 accuracy (95% patient-clustered bootstrap CI)")
    ax.set_title("Ensemble accuracy with bootstrap confidence intervals")
    ax.margins(x=0.18, y=0.35)
    style_axes(ax, grid_axis="x")
    ax.spines["left"].set_visible(False)
    ax.tick_params(left=False)

    png = save_figure(fig, ctx.out_dir, "bootstrap_ci")
    save_json(
        {"iters": BOOTSTRAP_ITERS, "ci_level": CI_LEVEL, "families": results},
        ctx.out_dir,
        "bootstrap_ci",
    )
    ctx.log.info(
        f"bootstrap_ensemble_ci: wrote {png.name} ({len(results)} family/families)."
    )