"""q_confirm.py: the frozen confirmatory estimators (H1, H2, H3 and their robustness variants).

Every function takes an analysis population (a DataFrame of paths with the columns y, family, run_id, N and
the features) and the registered settings, and returns a flat dict of results: the estimate, its percentile
interval, the one-sided p-value with its null, and the share of undefined bootstrap draws. The functions read
no values file and write nothing; the Phase 4 script calls them with the settings named in plan.json and
registers the returned numbers under the names listed there. Phase 3 runs them only on synthetic data
(q04_check.py), never on the dataset. This file is frozen with the plan.

Design choices fixed here (see plan.json, methods M3 to M9):
  - one run-bootstrap weight matrix per population, stratified by tool family (the runs of each family are
    resampled within the family), shared by both evaluation schemes so that the AUC difference is paired;
  - the run-held-out statistic is the mean over repeats of each repeat's pooled AUC (auc_over_repeats);
  - predictions are held fixed across resamples;
  - p-values test the stated null (AUC difference at or below zero; kappa at or below its null, the null
    recomputed in every resample) with the plus-one rule; undefined draws are counted, never silently kept.
"""
import numpy as np
import pandas as pd

import q_common as QC

AUC_NULL = 0.0     # null value of an AUC difference
KAPPA_NULL = 0.0   # null value of a kappa excess over its independence expectation

# Every robustness variant as the single change it makes (the check asserts these equal the registered judgments).
VARIANTS = {
    "v1": {"population": "no_twins"},
    "v2": {"population": "no_early"},
    "v3": {"kind": "rf"},
    "v5": {"winsor": True},
    "v6a": {"comb_half_bins": 1},
    "v6b": {"comb_half_bins": 3},
    "v6c": {"rel_speed_tol": 0.005},
    "v7": {"band": (200.0, 5000.0)},
    "v9": {"strata": "family_speed"},
}
VARIANTS_FOR = {"H1": ["v1", "v2", "v3", "v5"], "H2": ["v1", "v2", "v3", "v5"],
                "H3": ["v1", "v2", "v6a", "v6b", "v6c", "v7", "v9"]}


def population(paths, name, early_runs):
    """The analysis population of a variant: all paths, without defective-tool twins, or without early-stopped
    runs (early_runs is the set of run numbers whose summary remark records an early stop)."""
    if name == "no_twins":
        return paths[~paths.Tool.str.endswith("*")].reset_index(drop=True)
    if name == "no_early":
        return paths[~paths.run.isin(early_runs)].reset_index(drop=True)
    return paths.reset_index(drop=True)


def h3_strata(df, edge, kind="speed"):
    """H3 strata: spindle-speed class, or tool family crossed with speed class (variant V9)."""
    cls = speed_class(df.N.values, edge)
    if kind == "family_speed":
        return np.array([f"{f}|{c}" for f, c in zip(df.family.values, cls)])
    return cls


def boot_weights(df, B, seed):
    """Run-bootstrap path weights, stratified by tool family (a path's weight is its run's multiplicity)."""
    return QC.run_weights_strat(df.run_id.values, df.family.values, B, seed)


def _summ(name, est, draws, null, level):
    lo, hi = QC.percentile_ci(draws, level)
    out = {name: float(est), f"{name}_lo": lo, f"{name}_hi": hi, f"{name}_nonfinite": QC.nonfinite_share(draws)}
    if null is not None:
        out[f"{name}_p"] = QC.one_sided_p(draws, null)
    return out


def auc_drop(df, fset, B, seed, level=0.95, repeats=20, k=5, kind="lr", C=1.0, winsor=False, W=None):
    """H1 and H2 (and their variants V1, V2, V3, V5): pooled out-of-fold AUC under run-held-out evaluation
    (mean over repeats of each repeat's AUC) minus the pooled AUC under tool-family-held-out evaluation."""
    df = df.reset_index(drop=True)
    y = df.y.values
    S = QC.oof_run_matrix(df, fset, repeats, k, kind, C, winsor, True, seed)
    s_fam = QC.oof_family(df, fset, kind, C, winsor, True)
    a_run, a_fam = QC.auc_over_repeats(y, S), QC.auc_w(y, s_fam)
    W = boot_weights(df, B, seed) if W is None else W
    d_run = np.array([QC.auc_over_repeats(y, S, w) for w in W])
    d_fam = np.array([QC.auc_w(y, s_fam, w) for w in W])
    out = {}
    out.update(_summ("auc_run", a_run, d_run, None, level))
    out.update(_summ("auc_family", a_fam, d_fam, None, level))
    out.update(_summ("drop", a_run - a_fam, d_run - d_fam, AUC_NULL, level))
    out["nonfinite"] = max(out["drop_nonfinite"], out["auc_run_nonfinite"], out["auc_family_nonfinite"])
    out["n_paths"], out["n_runs"], out["n_families"] = len(df), int(df.run_id.nunique()), int(df.family.nunique())
    return out, {"scores_run": S, "scores_family": s_fam}


def speed_class(N, edge):
    """H3 strata: 0 below the registered edge, 1 at or above it."""
    return (np.asarray(N, float) >= edge).astype(int)


def cooccurrence(ia, ib, df, strata, B, seed, level=0.95, frac=0.5, W=None):
    """H3 (and its variants V1, V2, V6a, V6b, V6c, V7, V9): Cohen's kappa of two binary indicators on the same
    paths, its null (independence within strata), the largest kappa the per-stratum shares allow, the
    threshold rule, and the one-sided p-value of the excess over the null (null recomputed in each resample)."""
    ia = np.asarray(ia, float); ib = np.asarray(ib, float); strata = np.asarray(strata)
    keep = np.isfinite(ia) & np.isfinite(ib)
    df = df.loc[keep].reset_index(drop=True)
    ia, ib, strata = ia[keep].astype(int), ib[keep].astype(int), strata[keep]
    k = QC.kappa_w(ia, ib)
    nul = QC.kappa_null_w(ia, ib, strata)
    kmax = QC.kappa_max_w(ia, ib, strata)
    W = boot_weights(df, B, seed) if W is None else W
    dk = np.array([QC.kappa_w(ia, ib, w) for w in W])
    dn = np.array([QC.kappa_null_w(ia, ib, strata, w) for w in W])
    out = {}
    out.update(_summ("kappa", k, dk, None, level))
    out.update(_summ("excess", k - nul, dk - dn, KAPPA_NULL, level))
    out["kappa_null"], out["kappa_max"] = float(nul), float(kmax)
    out["thr"] = float(QC.kappa_threshold(nul, kmax, frac))
    out["nonfinite"] = max(out["kappa_nonfinite"], out["excess_nonfinite"])
    out["n_paths"], out["n_dropped"] = int(keep.sum()), int((~keep).sum())
    out["n_runs"] = int(df.run_id.nunique())
    return out


def holm_adjust(pvals, tested):
    """Holm over the tested predictions only; pvals maps a prediction id to its p-value, tested is the set of
    ids that passed both testability conditions (q_rules.testable). Untested predictions get NaN."""
    ids = [k for k, p in pvals.items() if k in tested and np.isfinite(p)]
    adj = QC.holm([pvals[k] for k in ids]) if ids else []
    out = {k: float("nan") for k in pvals}
    out.update({k: float(a) for k, a in zip(ids, adj)})
    return out