"""Analyze harbor arm results: pooled pass rate, Fisher exact vs A0, cluster bootstrap.

Reads harbor job dirs (results/jobs/<arm>/) and extracts per-(task,trial) pass/fail
from each trial's result. Reports per-arm pooled pass rate, per-task pass counts,
Fisher's exact test vs the baseline arm, and a task-cluster bootstrap 95% CI on ΔP
(resample the 17 tasks with replacement, keeping all k trials per task — the honest
test given within-task correlation).

Win criterion (pre-registered): bootstrap CI excludes 0, OR Fisher p<0.05.

Usage: .venv/bin/python scripts/31_analysis/stats.py --baseline A0 --arms A0 A1 A_comb
       (arm = subdir name under results/jobs/)
"""
import argparse, glob, json, sys
from collections import defaultdict
from pathlib import Path
import numpy as np
from scipy.stats import fisher_exact

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from common import RESULTS


def load_arm(arm):
    """Return dict task -> list of 0/1 outcomes across trials."""
    base = RESULTS / "jobs" / arm
    task_res = defaultdict(list)
    # harbor writes each trial under <job>/<task>__<hash>/ with a result/reward file
    for trial_dir in glob.glob(str(base / "*" / "*__*")):
        td = Path(trial_dir)
        task = td.name.split("__")[0]
        passed = None
        # try trial result.json
        for cand in ["result.json", "results.json"]:
            p = td / cand
            if p.exists():
                try:
                    d = json.loads(p.read_text())
                    passed = _extract_pass(d)
                except Exception:
                    pass
        if passed is None:
            rw = td / "reward.txt"
            if rw.exists():
                try:
                    passed = float(rw.read_text().strip()) >= 1.0
                except Exception:
                    pass
        if passed is not None:
            task_res[task].append(1 if passed else 0)
    return task_res


def _extract_pass(d):
    for k in ("resolved", "is_resolved", "passed", "success"):
        if k in d:
            return bool(d[k])
    if "reward" in d:
        return float(d["reward"]) >= 1.0
    if isinstance(d.get("results"), dict):
        s = d["results"].get("summary", {})
        return s.get("failed", 1) == 0 and s.get("passed", 0) > 0
    return None


def pooled(task_res):
    flat = [x for v in task_res.values() for x in v]
    return sum(flat), len(flat)


def cluster_bootstrap(a_res, b_res, n_boot=10000, seed=0):
    """ΔP = P(a) - P(b), resampling tasks with replacement."""
    tasks = sorted(set(a_res) | set(b_res))
    rng = np.random.default_rng(seed)
    deltas = []
    for _ in range(n_boot):
        pick = rng.choice(len(tasks), len(tasks), replace=True)
        a_hit = a_tot = b_hit = b_tot = 0
        for j in pick:
            t = tasks[j]
            a_hit += sum(a_res.get(t, [])); a_tot += len(a_res.get(t, []))
            b_hit += sum(b_res.get(t, [])); b_tot += len(b_res.get(t, []))
        if a_tot and b_tot:
            deltas.append(a_hit / a_tot - b_hit / b_tot)
    lo, hi = np.percentile(deltas, [2.5, 97.5])
    return float(lo), float(hi)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--baseline", default="A0")
    ap.add_argument("--arms", nargs="+", required=True)
    args = ap.parse_args()

    base = load_arm(args.baseline)
    b_hit, b_tot = pooled(base)
    all_tasks = sorted(base)
    print(f"baseline {args.baseline}: {b_hit}/{b_tot} = {b_hit/max(b_tot,1):.1%}\n")

    for arm in args.arms:
        res = load_arm(arm)
        a_hit, a_tot = pooled(res)
        all_tasks = sorted(set(all_tasks) | set(res))
        line = f"{arm}: {a_hit}/{a_tot} = {a_hit/max(a_tot,1):.1%}"
        if arm != args.baseline and b_tot:
            # Fisher on pooled 2x2
            _, p = fisher_exact([[a_hit, a_tot - a_hit], [b_hit, b_tot - b_hit]])
            lo, hi = cluster_bootstrap(res, base)
            win = "WIN" if (lo > 0 or p < 0.05) else ""
            line += f"  | Fisher p={p:.3f}  bootstrap ΔP 95% CI=[{lo:+.1%},{hi:+.1%}]  {win}"
        print(line)

    # per-task table
    print("\nper-task pass counts:")
    header = "task".ljust(30) + "".join(a[:10].rjust(11) for a in args.arms)
    print(header)
    arm_res = {a: load_arm(a) for a in args.arms}
    for t in all_tasks:
        row = t.ljust(30)
        for a in args.arms:
            v = arm_res[a].get(t, [])
            row += f"{sum(v)}/{len(v)}".rjust(11)
        print(row)


if __name__ == "__main__":
    main()