"""Exploratory recovery experiment (A-011): A3 vs A4 with the first session interrupted at 120 s.
ARI = P(final self-check pass | interrupted, A4); acceptance after recovery vs the interrupted A3 control and vs uninterrupted A3/A4 (pilot)."""
import os, sys, json
import numpy as np, pandas as pd
REPO = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")); sys.path.insert(0, os.path.join(REPO, "analysis", "econometrics"))
from estimate import load, cluster_bootstrap, permutation_p
def main():
    rec = load("recovery", 70); pilot = load("pilot", 70)
    rows = []
    for arm, g in rec.groupby("arm"):
        acc = g["S"].sum()
        rows.append({"condition": f"{arm} interrupted@120s", "n": len(g), "acceptance": g.groupby("task_id")["S"].mean().mean(), "det": g["D"].mean(), "quality": g["quality"].mean(),
                     "critical": g["E"].mean(), "cost_mean": g["C"].mean(), "cost_per_accepted": g["C"].sum() / acc if acc else np.nan, "wall_total_median": g["wall_total_s"].median(),
                     "recovery_attempted": g["recovery_attempted"].astype(bool).mean(), "selfcheck_first": g["selfcheck_pass_first"].mean(), "selfcheck_final": g["selfcheck_pass_final"].mean()})
    for arm in ("A3", "A4"):
        g = pilot[pilot.arm == arm]; acc = g["S"].sum()
        rows.append({"condition": f"{arm} uninterrupted (pilot)", "n": len(g), "acceptance": g.groupby("task_id")["S"].mean().mean(), "det": g["D"].mean(), "quality": g["quality"].mean(), "critical": g["E"].mean(),
                     "cost_mean": g["C"].mean(), "cost_per_accepted": g["C"].sum() / acc if acc else np.nan, "wall_total_median": g["wall_total_s"].median(), "recovery_attempted": g["recovery_attempted"].astype(bool).mean(),
                     "selfcheck_first": g["selfcheck_pass_first"].mean(), "selfcheck_final": g["selfcheck_pass_final"].mean()})
    t = pd.DataFrame(rows); t.to_csv(os.path.join(REPO, "analysis", "tables", "table10_recovery.csv"), index=False)
    a4 = rec[rec.arm == "A4"]; interrupted = a4[a4["recovery_attempted"].astype(bool)]
    ari = {"n_A4_interrupted": int(len(a4)), "n_recovery_attempted": int(len(interrupted)), "n_selfcheck_pass_final": int((interrupted["selfcheck_pass_final"] == 1).sum()),
           "ARI": float((interrupted["selfcheck_pass_final"] == 1).mean()) if len(interrupted) else None, "acceptance_after_recovery": float(interrupted["S"].mean()) if len(interrupted) else None,
           "total_cost_mean_A4_interrupted_cumulative": float(a4["C"].mean()), "note": "A4 interrupted cost = cumulative client total over both sessions (no imputation); A3 interrupted cost imputed from wall time", "recovery_wall_mean": float(a4.get("wall_total_s", pd.Series()).mean() if len(a4) else 0)}
    rng = np.random.default_rng(7); est, lo, hi, J = cluster_bootstrap(rec, "S", "A4", "A3", 5000, rng); p = permutation_p(rec, "S", "A4", "A3", 20000, rng)
    out = {"ari": ari, "contrast_A4int_minus_A3int_S": {"estimate": est, "ci_lo": lo, "ci_hi": hi, "perm_p": p, "J": J}}
    json.dump(out, open(os.path.join(REPO, "analysis", "derived", "recovery_experiment.json"), "w"), indent=1)
    print(t.round(3).to_string()); print(json.dumps(out, indent=1))
if __name__ == "__main__": main()