"""Primary analysis: task-weighted completion by arm, prespecified contrasts (D−A primary; B−A, C−B, D−C secondary with Holm),
task-cluster bootstrap CIs and paired sign-flip permutation p-values. Also secondary endpoints by arm."""
import os, sys, json
sys.path.insert(0, os.path.dirname(__file__))
from common import *

def arm_table(df):
    rows = []
    for a in ARMS:
        d = df[df["treatment"] == a]
        rows.append({"arm": a, "n_runs": int(len(d)), "n_tasks": int(d["task_id"].nunique()), "n_evaluated": int((d["status"] == "evaluated").sum()),
                     "Y_task_weighted": task_weighted_rate(d, a, "Y"), "Y_run_mean": float(d["Y"].mean()) if len(d) else np.nan,
                     "correct_any_task_weighted": task_weighted_rate(d, a, "Y_any"), "violation_rate": float(d["viol"].mean()) if len(d) else np.nan,
                     "n_runs_with_violation": int(d["viol"].sum()), "reproducible_rate": float(d["reproducibility"].mean()) if len(d) else np.nan,
                     "quality_mean": float(d["quality_score"].mean()) if len(d) else np.nan, "q_correctness": float(d["q_correctness"].mean()) if len(d) else np.nan,
                     "q_reporting": float(d["q_reporting"].mean()) if len(d) else np.nan, "q_baseline": float(d["q_baseline"].mean()) if len(d) else np.nan,
                     "cost_mean_usd": float(d["spent_usd"].mean()) if len(d) else np.nan, "cost_median_usd": float(d["spent_usd"].median()) if len(d) else np.nan,
                     "cost_p90_usd": float(d["spent_usd"].quantile(0.9)) if len(d) else np.nan, "cost_total_usd": float(d["spent_usd"].sum()),
                     "reviewer_usd_total": float(d["reviewer_usd"].fillna(0).sum()), "wall_median_s": float(d["spent_wall_s"].median()) if len(d) else np.nan,
                     "wall_p90_s": float(d["spent_wall_s"].quantile(0.9)) if len(d) else np.nan, "n_timeout": int((d["terminal_reason"] == "wall_clock_limit").sum()),
                     "n_budget_exhausted": int(d["terminal_reason"].fillna("").str.contains("budget").sum()), "n_sessions_mean": float(d["n_sessions"].mean()) if len(d) else np.nan,
                     "cost_per_success_usd": (float(d["spent_usd"].sum() / d["Y"].sum()) if d["Y"].sum() > 0 else None)})
    return pd.DataFrame(rows)

def contrasts(df, col="Y", B=5000, n_perm=20000):
    out = []
    pv = {}
    for (a, b, role) in [("D", "A", "primary"), ("B", "A", "secondary"), ("C", "B", "secondary"), ("D", "C", "secondary"), ("C", "A", "exploratory"), ("D", "B", "exploratory")]:
        d = paired_task_diff(df, a, b, col); est, lo, hi = cluster_bootstrap(d, B); p = permutation_p(d, n_perm)
        out.append({"contrast": f"{a}-{b}", "role": role, "endpoint": col, "n_task_clusters": int(len(d)), "estimate": est, "ci_low": lo, "ci_high": hi, "perm_p": p,
                    "n_tasks_favoring_first": int((d > 0).sum()), "n_tasks_favoring_second": int((d < 0).sum()), "n_tasks_tied": int((d == 0).sum())})
        if role == "secondary": pv[f"{a}-{b}"] = p
    adj = holm(pv) if pv else {}
    for r in out:
        r["holm_p"] = adj.get(r["contrast"])
    return pd.DataFrame(out)

def main(split="confirmatory"):
    os.makedirs(EST, exist_ok=True)
    df = load_runs(split)
    if df.empty:
        print(f"no runs for split {split}"); return
    at = arm_table(df); at.to_csv(os.path.join(EST, f"{split}_arm_table.csv"), index=False)
    ct = pd.concat([contrasts(df, "Y"), contrasts(df, "Y_any"), contrasts(df, "viol"), contrasts(df, "quality_score")]); ct.to_csv(os.path.join(EST, f"{split}_contrasts.csv"), index=False)
    fam = df.groupby(["family", "treatment"]).apply(lambda d: pd.Series({"Y": d.groupby("task_id")["Y"].mean().mean(), "n_tasks": d["task_id"].nunique(), "n_runs": len(d), "viol": d["viol"].mean(), "cost_mean": d["spent_usd"].mean()}), include_groups=False).reset_index()
    fam.to_csv(os.path.join(EST, f"{split}_family_table.csv"), index=False)
    print(at.round(3).to_string()); print(); print(ct.round(4).to_string())

if __name__ == "__main__":
    main(sys.argv[1] if len(sys.argv) > 1 else "confirmatory")