#!/usr/bin/env python3
"""Paired A/B analysis for focuschain-ab.

Inputs: trials/focuschain-ab/summary/focuschain-ab/results.csv (from
summarize_results.py; one row per (variant=block, task)) and the shared pack's
pack_detail.csv (stratum per task).

Primary: per-task mean score per arm (runs averaged), paired over tasks.
- McNemar exact test on binarized per-task means (success := mean score >= 0.5),
  overall and per stratum; also per run-pairing (control-1 vs treatment-1,
  control-2 vs treatment-2) as sensitivity.
- Bootstrap (10k, resample tasks) CI for mean paired score delta.
- Cost: $/task using OpenRouter upstream cost when present, else est_cost_usd.
- Cache hit rate per arm (prompt-token-weighted); flags treatment < control.

Writes trials/focuschain-ab/analysis.json and prints a readable summary.
"""
import csv
import json
import math
import random
from collections import defaultdict
from pathlib import Path

ROOT = Path("/srv/controller-data")
SUMMARY = ROOT / "trials/focuschain-ab/summary/focuschain-ab/results.csv"
PACK = ROOT / "trials/shared/osworld-40-pack/pack_detail.csv"
OUT = ROOT / "trials/focuschain-ab/analysis.json"

ARM_OF = {"control-1": "control", "control-2": "control",
          "treatment-1": "treatment", "treatment-2": "treatment"}
SUCCESS_THRESHOLD = 0.5
N_BOOT = 10_000


def binom_two_sided_p(k: int, n: int) -> float:
    """Exact two-sided binomial(n, 0.5) p-value (McNemar exact)."""
    if n == 0:
        return 1.0
    def pmf(i):
        return math.comb(n, i) * 0.5 ** n
    p_k = pmf(k)
    return min(1.0, sum(pmf(i) for i in range(n + 1) if pmf(i) <= p_k + 1e-12))


def mcnemar(pairs):
    """pairs: list of (ctrl_success, treat_success) booleans."""
    b = sum(1 for c, t in pairs if c and not t)   # control-only wins
    c = sum(1 for c, t in pairs if t and not c)   # treatment-only wins
    return {"control_only": b, "treatment_only": c,
            "p_exact": round(binom_two_sided_p(min(b, c), b + c), 5)}


def bootstrap_ci(deltas, n_boot=N_BOOT, seed=20260703):
    rng = random.Random(seed)
    n = len(deltas)
    if n == 0:
        return (float("nan"),) * 3
    means = []
    for _ in range(n_boot):
        means.append(sum(deltas[rng.randrange(n)] for _ in range(n)) / n)
    means.sort()
    return (sum(deltas) / n,
            means[int(0.025 * n_boot)],
            means[int(0.975 * n_boot)])


def main():
    stratum = {}
    with open(PACK) as f:
        for r in csv.DictReader(f):
            stratum[r["task_id"]] = r["stratum"]

    rows = list(csv.DictReader(open(SUMMARY)))
    rows = [r for r in rows if r["variant"] in ARM_OF]

    # (arm, task) -> list of per-run dicts
    runs = defaultdict(list)
    for r in rows:
        runs[(ARM_OF[r["variant"]], r["task_id"])].append(r)

    tasks = sorted({t for (_, t) in runs})
    per_task = {}
    for t in tasks:
        entry = {"stratum": stratum.get(t, "?")}
        for arm in ("control", "treatment"):
            rs = runs.get((arm, t), [])
            if not rs:
                continue
            score = sum(float(x["score"]) for x in rs) / len(rs)
            up = sum(float(x["upstream_cost"]) for x in rs) / len(rs)
            est = sum(float(x["est_cost_usd"]) for x in rs) / len(rs)
            entry[arm] = {
                "n_runs": len(rs),
                "score": score,
                "success": score >= SUCCESS_THRESHOLD,
                "cost": up if up > 0 else est,
                "steps": sum(int(x["steps"]) for x in rs) / len(rs),
                "hit_rate": sum(float(x["cache_hit_rate"]) for x in rs) / len(rs),
                "invalidations": sum(int(x["invalidation_events"]) for x in rs),
                "edits": sum(int(x["compactions"]) + int(x["image_trims"]) for x in rs),
                "scores_by_run": {x["variant"]: float(x["score"]) for x in rs},
            }
        per_task[t] = entry

    paired = {t: e for t, e in per_task.items() if "control" in e and "treatment" in e}

    def analyze(subset_name, task_ids):
        sub = [(t, paired[t]) for t in task_ids if t in paired]
        if not sub:
            return None
        deltas = [e["treatment"]["score"] - e["control"]["score"] for _, e in sub]
        mean_d, lo, hi = bootstrap_ci(deltas)
        mc = mcnemar([(e["control"]["success"], e["treatment"]["success"]) for _, e in sub])
        return {
            "n_tasks": len(sub),
            "control_score": round(sum(e["control"]["score"] for _, e in sub), 3),
            "treatment_score": round(sum(e["treatment"]["score"] for _, e in sub), 3),
            "control_success": sum(e["control"]["success"] for _, e in sub),
            "treatment_success": sum(e["treatment"]["success"] for _, e in sub),
            "mean_paired_delta": round(mean_d, 4),
            "bootstrap_95ci": [round(lo, 4), round(hi, 4)],
            "mcnemar": mc,
            "control_cost_per_task": round(sum(e["control"]["cost"] for _, e in sub) / len(sub), 4),
            "treatment_cost_per_task": round(sum(e["treatment"]["cost"] for _, e in sub) / len(sub), 4),
        }

    result = {
        "overall": analyze("overall", list(paired)),
        "medium": analyze("medium", [t for t in paired if paired[t]["stratum"] == "medium"]),
        "long": analyze("long", [t for t in paired if paired[t]["stratum"] == "long"]),
    }

    # Per run-pairing sensitivity McNemar (c1-t1, c2-t2)
    for ci, ti in (("control-1", "treatment-1"), ("control-2", "treatment-2")):
        pairs = []
        for t, e in paired.items():
            cs = e["control"]["scores_by_run"].get(ci)
            ts = e["treatment"]["scores_by_run"].get(ti)
            if cs is not None and ts is not None:
                pairs.append((cs >= SUCCESS_THRESHOLD, ts >= SUCCESS_THRESHOLD))
        result[f"mcnemar_{ci}_vs_{ti}"] = mcnemar(pairs)

    # Arm-level cache + cost + edit integrity (token-weighted over all runs)
    arm_stats = {}
    for arm in ("control", "treatment"):
        sub = [r for r in rows if ARM_OF[r["variant"]] == arm]
        tok = sum(float(r["raw_prompt_tokens"]) for r in sub)
        hit = (sum(float(r["cache_hit_rate"]) * float(r["raw_prompt_tokens"]) for r in sub) / tok) if tok else 0
        arm_stats[arm] = {
            "task_runs": len(sub),
            "hit_rate_weighted": round(hit, 4),
            "invalidations": sum(int(r["invalidation_events"]) for r in sub),
            "edits": sum(int(r["compactions"]) + int(r["image_trims"]) for r in sub),
            "total_upstream_cost": round(sum(float(r["upstream_cost"]) for r in sub), 4),
            "total_est_cost": round(sum(float(r["est_cost_usd"]) for r in sub), 4),
            "total_steps": sum(int(r["steps"]) for r in sub),
            "output_tokens": sum(int(r["output_tokens"]) for r in sub),
        }
    result["arm_stats"] = arm_stats
    if arm_stats["treatment"]["hit_rate_weighted"] < arm_stats["control"]["hit_rate_weighted"] - 0.02:
        result["WARNING"] = ("treatment cache hit rate is materially below control — "
                             "something is editing history; debug before trusting results")

    result["per_task"] = {
        t: {"stratum": e["stratum"],
            "control_score": round(e["control"]["score"], 3),
            "treatment_score": round(e["treatment"]["score"], 3)}
        for t, e in paired.items()
    }

    OUT.write_text(json.dumps(result, indent=2))
    print(json.dumps({k: v for k, v in result.items() if k != "per_task"}, indent=2))
    print(f"\nwrote {OUT}")


if __name__ == "__main__":
    main()
