"""Shared helpers for the WP29 analysis pipeline (frozen interfaces)."""
import os, json, numpy as np, pandas as pd
REPO = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
DATA = os.path.join(REPO, "results", "analysis_data"); EST = os.path.join(REPO, "results", "estimates")
ARMS = ["A", "B", "C", "D"]

def load_runs(split=None, path=None, include_postfix=False):
    df = pd.read_csv(path or os.path.join(DATA, "runs.csv"))
    if split: df = df[df["split"] == split]
    df = df.copy()
    if "notes" in df.columns:
        df["postfix_verification"] = df["notes"].fillna("").str.contains("seed 102|seed 402|seed 403|seed 802")
        if not include_postfix: df = df[~df["postfix_verification"]]
    no_intervention = df["human_interventions"].fillna("").eq("") if "human_interventions" in df.columns else True
    df["Y"] = ((df["scientific_success"] == 1) & (df["status"].isin(["evaluated"])) & (df["terminal_reason"].fillna("").str.startswith("executor_finished")) & no_intervention).astype(int)   # max_turns_reached and result_* are outside the envelope
    # Y requires completion within the envelope: timeout/budget_exhausted → 0 (status is 'evaluated' for all evaluated runs; terminal_reason keeps the cause)
    df.loc[df["terminal_reason"].fillna("").isin(["wall_clock_limit"]) | df["terminal_reason"].fillna("").str.contains("budget", na=False), "Y"] = 0
    df["Y_any"] = (df["scientific_success"] == 1).astype(int)          # correct regardless of envelope
    df["viol"] = (df["n_violations"].fillna(0) > 0).astype(int)
    df["family"] = df["task_id"].str.split("-").str[0]
    return df

def task_weighted_rate(df, arm, col="Y"):
    d = df[df["treatment"] == arm]
    if d.empty: return np.nan
    return d.groupby("task_id")[col].mean().mean()

def paired_task_diff(df, a, b, col="Y"):
    """Task-level paired differences (mean over repeats within task) a − b, as a Series indexed by task."""
    ma = df[df["treatment"] == a].groupby("task_id")[col].mean(); mb = df[df["treatment"] == b].groupby("task_id")[col].mean()
    common = ma.index.intersection(mb.index)
    return (ma.loc[common] - mb.loc[common])

def cluster_bootstrap(diffs, B=5000, seed=20261030):
    rng = np.random.default_rng(seed); x = diffs.values; n = len(x)
    if n == 0: return (np.nan, np.nan, np.nan)
    boots = np.array([x[rng.integers(0, n, n)].mean() for _ in range(B)])
    return float(x.mean()), float(np.percentile(boots, 2.5)), float(np.percentile(boots, 97.5))

def permutation_p(diffs, n_perm=20000, seed=20261030):
    rng = np.random.default_rng(seed); x = diffs.values; n = len(x)
    if n == 0: return np.nan
    obs = abs(x.mean()); signs = rng.choice([-1, 1], size=(n_perm, n))
    perm = np.abs((signs * x).mean(axis=1))
    return float((np.sum(perm >= obs - 1e-12) + 1) / (n_perm + 1))

def holm(pvals: dict):
    items = sorted(pvals.items(), key=lambda kv: kv[1]); m = len(items); out = {}; prev = 0.0
    for i, (k, p) in enumerate(items):
        adj = min(1.0, max(prev, (m - i) * p)); out[k] = adj; prev = adj
    return out