# =============================================================================
#  Project   : modelmap
#  File      : src/modelmap/stats/replication.py
#  Purpose   : Bootstrap CIs, BH-FDR control, selectivity, replication metrics
#  Author    : Simon-Pierre Boucher
#  Contact   : contact@spboucher.ai
#  Website   : https://modelmap.io
#  Created   : 2026-08-12
#  Modified  : 2026-08-12
#  Platform  : macOS / Apple Silicon (arm64)
#  License   : All rights reserved (research code)
# =============================================================================
"""Statistical core mandated by charter §8.2 and the Dead-Salmons doctrine
(notes §4.8): bootstrap confidence intervals, Benjamini-Hochberg FDR control
for unit scans, probe selectivity, and map-replication metrics.
"""

from __future__ import annotations

import numpy as np


def bootstrap_ci(
    values: np.ndarray,
    stat=np.mean,
    n_boot: int = 10_000,
    alpha: float = 0.05,
    rng: np.random.Generator | None = None,
) -> tuple[float, float, float]:
    """Percentile bootstrap CI. Returns (point, lo, hi)."""
    values = np.asarray(values, dtype=np.float64)
    if values.size == 0:
        raise ValueError("bootstrap_ci: empty sample")
    rng = rng or np.random.default_rng(0)
    idx = rng.integers(0, values.size, size=(n_boot, values.size))
    boots = stat(values[idx], axis=1)
    return (
        float(stat(values)),
        float(np.quantile(boots, alpha / 2)),
        float(np.quantile(boots, 1 - alpha / 2)),
    )


def bh_fdr(p_values: np.ndarray, q: float = 0.05) -> np.ndarray:
    """Benjamini-Hochberg: boolean mask of discoveries at FDR level q."""
    p = np.asarray(p_values, dtype=np.float64)
    m = p.size
    if m == 0:
        return np.zeros(0, dtype=bool)
    order = np.argsort(p)
    ranked = p[order]
    thresh = q * (np.arange(1, m + 1) / m)
    passing = np.nonzero(ranked <= thresh)[0]
    out = np.zeros(m, dtype=bool)
    if passing.size:
        out[order[: passing.max() + 1]] = True
    return out


def selectivity(task_accuracy: float, control_accuracy: float) -> float:
    """Hewitt & Liang selectivity: task acc minus shuffled-label control acc."""
    return float(task_accuracy - control_accuracy)


def jaccard(a: set, b: set) -> float:
    """Overlap of 'important unit' sets — the replication metric of expD."""
    if not a and not b:
        return 1.0
    return len(a & b) / len(a | b)


def replication_rate(unit_sets: list[set]) -> tuple[float, float, float]:
    """Mean pairwise Jaccard across seeds/datasets, with bootstrap CI.

    This is THE first-class atlas metric (charter §8.3): does the map
    reproduce under resampling?
    """
    n = len(unit_sets)
    if n < 2:
        raise ValueError("replication_rate needs >=2 replicates")
    pairs = [jaccard(unit_sets[i], unit_sets[j]) for i in range(n) for j in range(i + 1, n)]
    return bootstrap_ci(np.array(pairs))