"""Paired, cluster-bootstrap statistics (cluster = document). Holm correction for the registered primary family."""
from __future__ import annotations
import numpy as np
import pandas as pd

def cluster_bootstrap_mean(df: pd.DataFrame, value: str, cluster: str = "document_id", n_boot: int = 2000, seed: int = 0, ci: float = 0.95) -> dict:
    """Mean of `value` with a cluster bootstrap CI (resample documents with replacement; equal weight per observation)."""
    d = df[[cluster, value]].dropna()
    if d.empty:
        return {"mean": np.nan, "lo": np.nan, "hi": np.nan, "n": 0, "n_clusters": 0}
    groups = [g[value].to_numpy() for _, g in d.groupby(cluster)]
    rng = np.random.default_rng(seed); G = len(groups)
    sums = np.array([g.sum() for g in groups]); cnts = np.array([len(g) for g in groups])
    boots = np.empty(n_boot)
    for b in range(n_boot):
        idx = rng.integers(0, G, G)
        boots[b] = sums[idx].sum() / cnts[idx].sum()
    a = (1 - ci) / 2
    return {"mean": float(d[value].mean()), "lo": float(np.quantile(boots, a)), "hi": float(np.quantile(boots, 1 - a)), "n": int(len(d)), "n_clusters": int(G)}

def paired_difference(df: pd.DataFrame, value: str, cond_col: str, a: str, b: str, keys=("document_id", "question_id", "pair"), cluster: str = "document_id",
                      n_boot: int = 2000, seed: int = 0) -> dict:
    """Paired difference a - b on identical (document, question, pair) instances; cluster bootstrap CI and two-sided
    bootstrap p-value (proportion of bootstrap means on the other side of zero, doubled)."""
    A = df[df[cond_col] == a].set_index(list(keys))[value]; B = df[df[cond_col] == b].set_index(list(keys))[value]
    j = pd.concat([A.rename("a"), B.rename("b")], axis=1, join="inner").dropna()
    if j.empty:
        return {"a": a, "b": b, "diff": np.nan, "lo": np.nan, "hi": np.nan, "p": np.nan, "n": 0, "n_clusters": 0}
    j["d"] = j["a"] - j["b"]; j = j.reset_index()
    groups = [g["d"].to_numpy() for _, g in j.groupby(cluster)]
    rng = np.random.default_rng(seed); G = len(groups)
    sums = np.array([g.sum() for g in groups]); cnts = np.array([len(g) for g in groups])
    boots = np.empty(n_boot)
    for i in range(n_boot):
        idx = rng.integers(0, G, G); boots[i] = sums[idx].sum() / cnts[idx].sum()
    d = float(j["d"].mean())
    p = 2 * min((boots <= 0).mean(), (boots >= 0).mean()) if d != 0 else 1.0
    return {"a": a, "b": b, "diff": d, "lo": float(np.quantile(boots, 0.025)), "hi": float(np.quantile(boots, 0.975)), "p": float(min(1.0, p)), "n": int(len(j)), "n_clusters": int(G),
            "mean_a": float(j["a"].mean()), "mean_b": float(j["b"].mean()), "rel_change": float(d / j["b"].mean()) if j["b"].mean() not in (0, np.nan) else np.nan}

def holm(pvals: list[float]) -> list[float]:
    m = len(pvals); order = np.argsort(pvals); adj = np.empty(m)
    running = 0.0
    for rank, i in enumerate(order):
        v = min(1.0, (m - rank) * pvals[i]); running = max(running, v); adj[i] = running
    return adj.tolist()

def summarize_latency(x: pd.Series) -> dict:
    x = x.dropna()
    if x.empty: return {"p50": np.nan, "p95": np.nan, "n": 0}
    return {"p50": float(x.quantile(0.5)), "p95": float(x.quantile(0.95)), "mean": float(x.mean()), "n": int(len(x))}