#!/usr/bin/env python3
"""lab.lib.stats — reusable statistics for clustered, multi-seed experiments.

Extracted verbatim from sandbox/stage2/stat_hardening_g/aggregate.py (the
30-run hardening study) and generalized. Conventions:

- Class ids follow G_hf_dataset: 0=malignant tumor (MT), 1=non-malignant
  lesion (NML), 2=non-ROI. `minority` selects the class of interest.
- Cluster-aware everywhere: frames within one lesion/patient are NOT iid.
- Multi-seed by default: report mean±std over seeds; compare configs with
  per-seed McNemar + Holm correction across the family of tests.

Typical use:
    from lab.lib.stats import mcnemar_exact, holm, wilson, lesion_cluster_bootstrap

    res = mcnemar_exact(pred_a, pred_b, true)          # exact two-sided p
    adj = holm([r["p"] for r in seed_results])         # family-wise control
    ci  = wilson(k, n)                                  # proportion CI
    bs  = lesion_cluster_bootstrap(true, pred, lids)    # cluster bootstrap CI
"""
from __future__ import annotations

from collections import defaultdict

import numpy as np
from scipy import stats as st

__all__ = [
    "CLASS_NAMES", "mean_std", "wilson",
    "lesion_cluster_bootstrap", "mcnemar_exact", "pooled_mcnemar_exact",
    "holm", "benjamini_hochberg", "majority_vote",
]

# G_hf_dataset 3-class convention (verify against dataset metadata if a new
# variant reorders classes).
CLASS_NAMES = ["malignant tumor", "non-malignant lesion", "non-ROI"]


def valid_lid(lid) -> bool:
    """True when a frame belongs to an actual lesion cluster (not None/NaN)."""
    if lid is None:
        return False
    if isinstance(lid, float) and np.isnan(lid):
        return False
    return True


# ── Descriptives ──────────────────────────────────────────────────────────────

def mean_std(vals):
    """(mean, std[ddof=1], n); std=0.0 for single value; None for empty."""
    vals = [v for v in vals if v is not None]
    if not vals:
        return None, None, 0
    sd = float(np.std(vals, ddof=1)) if len(vals) > 1 else 0.0
    return float(np.mean(vals)), sd, len(vals)


def wilson(k: int, n: int, alpha: float = 0.05):
    """Wilson score interval for a binomial proportion. Returns (lo, hi)."""
    if n == 0:
        return (0.0, 0.0)
    z = st.norm.ppf(1 - alpha / 2)
    ph = k / n
    den = 1 + z**2 / n
    ctr = (ph + z**2 / (2 * n)) / den
    hw = z * np.sqrt(ph * (1 - ph) / n + z**2 / (4 * n**2)) / den
    return ctr - hw, ctr + hw


# ── Cluster bootstrap ─────────────────────────────────────────────────────────

def lesion_cluster_bootstrap(true, pred, lesion_ids,
                             metric: str = "nml_recall",
                             minority: int = 1,
                             n_boot: int = 10_000, seed: int = 123,
                             alpha: float = 5.0):
    """Percentile bootstrap CIs resampling LESIONS (not frames).

    All frames of a lesion move together — this respects the clustering that
    makes frame-iid intervals anti-conservative. Frames with invalid
    lesion_ids carry no cluster structure and are dropped.

    metric:
      - "nml_recall"  → recall of `minority` class over selected frames
      - "mcc"         → Matthews corrcoef over selected frames (needs ≥2 classes)
      - callable(t, p) → any custom statistic
    """
    rng = np.random.default_rng(seed)
    clusters: dict[str, list[int]] = defaultdict(list)
    for i, lid in enumerate(lesion_ids):
        if valid_lid(lid):
            clusters[str(lid)].append(i)
    idx_by_lid = [np.array(v) for v in clusters.values()]
    L = len(idx_by_lid)

    def stat(sel_idx):
        t = np.asarray(true)[sel_idx]
        p = np.asarray(pred)[sel_idx]
        if callable(metric):
            return float(metric(t, p))
        if metric == "nml_recall":
            denom = (t == minority).sum()
            return float((p[t == minority] == minority).sum() / denom) if denom else np.nan
        if metric == "mcc":
            from sklearn.metrics import matthews_corrcoef
            if len(set(t.tolist())) < 2:
                return np.nan
            return float(matthews_corrcoef(t, p))
        raise ValueError(f"unknown metric {metric!r}")

    point = stat(np.concatenate(idx_by_lid))
    boots = []
    for _ in range(n_boot):
        pick = rng.integers(0, L, L)
        v = stat(np.concatenate([idx_by_lid[j] for j in pick]))
        if not np.isnan(v):
            boots.append(v)
    lo, hi = np.percentile(boots, [alpha / 2, 100 - alpha / 2])
    return {"point": float(point), "ci_low": float(lo), "ci_high": float(hi),
            "n_lesions": L, "n_boot": len(boots)}


# ── Paired comparisons ────────────────────────────────────────────────────────

def mcnemar_exact(pred_a, pred_b, true):
    """Exact two-sided McNemar on discordant error pairs (same eval set).

    b = A wrong & B right, c = A right & B wrong.
    p = P(X ≤ min(b,c)) · 2 over Binomial(n=b+c, 0.5), capped at 1.
    """
    ea = np.asarray(pred_a) != np.asarray(true)
    eb = np.asarray(pred_b) != np.asarray(true)
    b = int((ea & ~eb).sum())
    c = int((~ea & eb).sum())
    n = b + c
    if n == 0:
        return {"b": b, "c": c, "p": 1.0}
    k = min(b, c)
    p = sum(st.binom.pmf(i, n, 0.5) for i in range(0, k + 1)) * 2
    return {"b": b, "c": c, "p": float(min(p, 1.0))}


def pooled_mcnemar_exact(pair_results):
    """Pool (b, c) counts from multiple seeds into one exact test.

    Use when the same comparison is repeated across seeds on the same eval
    frames and you want the aggregate direction + significance. Report BOTH
    pooled and per-seed-Holm results — pooling can mask sign flips.
    """
    b = sum(r["b"] for r in pair_results)
    c = sum(r["c"] for r in pair_results)
    n = b + c
    if n == 0:
        return {"b": b, "c": c, "p": 1.0}
    k = min(b, c)
    p = sum(st.binom.pmf(i, n, 0.5) for i in range(0, k + 1)) * 2
    return {"b": b, "c": c, "p": float(min(p, 1.0))}


# ── Multiple-comparison correction ────────────────────────────────────────────

def holm(pvals):
    """Holm–Bonferroni step-down. Returns adjusted p-values, original order."""
    pvals = np.asarray(pvals, dtype=float)
    m = len(pvals)
    order = np.argsort(pvals)
    adj = np.empty(m)
    running = 0.0
    for rank, idx in enumerate(order):
        val = (m - rank) * pvals[idx]
        running = max(running, val)
        adj[idx] = min(1.0, running)
    return adj


def benjamini_hochberg(pvals):
    """BH FDR correction. Returns adjusted q-values, original order."""
    pvals = np.asarray(pvals, dtype=float)
    m = len(pvals)
    order = np.argsort(pvals)[::-1]          # descending
    adj = np.empty(m)
    running = 0.0
    for rank, idx in enumerate(order):       # rank = m-1 … 0
        val = pvals[idx] * m / (rank + 1)
        running = max(running, val)
        adj[idx] = min(1.0, running)
    return adj


# ── Aggregation rules (frames → lesions) ──────────────────────────────────────

def majority_vote(pred, lesion_ids, n_classes: int = 3):
    """Per-lesson hard majority vote broadcast to frames.

    WARNING (hardening result, 2026-08-21): majority voting suppresses the
    minority class under domain shift (NML recall 0.26→0.02 at unchanged
    MCC). Do not use when minority recall matters; kept for completeness and
    as the negative-control aggregation rule.
    """
    l2i = defaultdict(list)
    for i, lid in enumerate(lesion_ids):
        if valid_lid(lid):
            l2i[lid].append(i)
    result = np.array(pred, copy=True)
    for _lid, idx in l2i.items():
        result[idx] = np.bincount(
            result[idx], minlength=max(n_classes, len(np.unique(result)))
        ).argmax()
    return result


def mean_confidence_vote(probabilities, lesion_ids, n_classes: int = 3):
    """Per-lesion average of frame probabilities, argmax → label, broadcast.

    Minority-friendlier alternative to hard majority vote: soft evidence is
    averaged before the decision instead of after, so a confident minority
    block inside a lesion is not erased by a slim non-lesion plurality.
    Not yet validated in the hardening study — treat as experimental.
    """
    l2i = defaultdict(list)
    for i, lid in enumerate(lesion_ids):
        if valid_lid(lid):
            l2i[lid].append(i)
    probabilities = np.asarray(probabilities, dtype=float)
    result = probabilities.argmax(1).copy()
    for _lid, idx in l2i.items():
        result[idx] = probabilities[idx].mean(0).argmax()
    return result