"""Bootstrap-CI scoring utilities for single-choice MCQ evaluations.

Use `bootstrap_ci(scores, strata)` for stratified resampling, `fmt(stat)` to
format a result as `mean [ci_lo, ci_hi]`, and `compute_simple_table(df, ...)`
for one-row-per-condition summaries with bootstrap CIs.
"""

import numpy as np
import pandas as pd


def bootstrap_ci(scores, strata_labels, n_bootstrap=10_000, alpha=0.05, seed=610):
    """Stratified bootstrap: resample within each stratum, compute overall mean.

    Vectorized — each stratum's resampling is one 2D `rng.choice` call across
    all `n_bootstrap` iterations at once, ~50-200x faster than the loop form
    for typical inputs (10K iterations x 60 strata) with identical statistical
    behavior.
    """
    scores = np.asarray(scores, dtype=float)
    rng = np.random.default_rng(seed)
    strata_labels = np.asarray([str(s) for s in strata_labels])
    _, inverse = np.unique(strata_labels, return_inverse=True)
    n_strata = int(inverse.max()) + 1 if len(inverse) else 0
    strata_indices = [np.where(inverse == i)[0] for i in range(n_strata)]

    n_total = len(scores)
    boot_sums = np.zeros(n_bootstrap, dtype=float)
    for idx in strata_indices:
        if len(idx) == 0:
            continue
        sampled = rng.choice(idx, size=(n_bootstrap, len(idx)), replace=True)
        boot_sums += scores[sampled].sum(axis=1)
    boot_means = boot_sums / max(n_total, 1)

    return {
        'mean':  float(np.mean(scores)) if n_total else float('nan'),
        'ci_lo': float(np.percentile(boot_means, 100 * alpha / 2)) if n_total else float('nan'),
        'ci_hi': float(np.percentile(boot_means, 100 * (1 - alpha / 2))) if n_total else float('nan'),
        'n':     n_total,
    }


def fmt(stat):
    """Format a bootstrap_ci result as 'mean [ci_lo, ci_hi]'."""
    return f"{stat['mean']:.3f} [{stat['ci_lo']:.3f}, {stat['ci_hi']:.3f}]"


def compute_simple_table(df, group_cols, strata_col, *,
                         metric_col: str = 'accuracy',
                         n_bootstrap: int = 10_000, alpha: float = 0.05):
    """One row per group: <group_cols>, n, accuracy + bootstrap CI columns.

    `dropna=False` so rows missing some condition column still form their own
    NaN-keyed group rather than being silently dropped.
    """
    rows = []
    if isinstance(group_cols, str):
        group_cols = [group_cols]
    for keys, grp in df.groupby(group_cols, dropna=False):
        row = dict(zip(group_cols, keys)) if isinstance(keys, tuple) else {group_cols[0]: keys}
        row['n'] = len(grp)
        stat = bootstrap_ci(grp[metric_col].values, grp[strata_col].values, n_bootstrap, alpha)
        row[metric_col]              = fmt(stat)
        row[f'{metric_col}_mean']    = stat['mean']
        row[f'{metric_col}_ci_lo']   = stat['ci_lo']
        row[f'{metric_col}_ci_hi']   = stat['ci_hi']
        rows.append(row)
    return pd.DataFrame(rows)