"""Reconstruct per-task Opus 4.6 cost for the ARC-AGI pilots from
Claude Code session JSONL logs, and compute a paired Wilcoxon
signed-rank test on the baseline vs g6_full cost arrays.

Why this exists: the pilot runner in benchmarking/arc_pilot/ routed
LLM calls through claude_code_headless but the poetiq_bridge cost
ledger (_CC_USAGE) was process-local and never persisted into
results.jsonl. The *session* logs under
~/.claude/projects/<workspace>/ do retain full per-message usage, so
we can price each assistant message against Anthropic's published
Opus 4.6 rates, then allocate session cost across (task, arm) windows
by time-overlap weighting.

Usage:
    python benchmarks/test_sets/arc_agi_pilot_reconstruct.py \\
        --run benchmarking/arc_pilot/runs/20260414T232841Z \\
        --session-dir ~/.claude/projects/<slug>

Caveats: overlap allocation is only reliable when the run window has
no concurrent unrelated Claude Code activity. The overnight v1 run
(2026-04-14T23:28 - 2026-04-15T02:47) is clean; the daytime v2 run
(2026-04-15T06:05 - 13:29) is not, and we publish no p-value from it.
"""
from __future__ import annotations

import argparse
import json
import os
from collections import defaultdict
from datetime import datetime, timedelta
from pathlib import Path

# Opus 4.6 published pricing (USD per MTok)
PRICE = {
    "input": 15.0,
    "output": 75.0,
    "cache_write_5m": 18.75,
    "cache_write_1h": 30.0,
    "cache_read": 1.5,
}


def parse_ts(s: str) -> datetime:
    return datetime.fromisoformat(s.replace("Z", "+00:00"))


def session_cost(msgs: list[dict]) -> float:
    c = 0.0
    for m in msgs:
        if m.get("type") != "assistant":
            continue
        msg = m.get("message") or {}
        u = msg.get("usage") if isinstance(msg, dict) else None
        if not u:
            continue
        cc = u.get("cache_creation") or {}
        c += (
            u.get("input_tokens", 0) * PRICE["input"]
            + u.get("output_tokens", 0) * PRICE["output"]
            + u.get("cache_read_input_tokens", 0) * PRICE["cache_read"]
            + cc.get("ephemeral_5m_input_tokens", 0) * PRICE["cache_write_5m"]
            + cc.get("ephemeral_1h_input_tokens", 0) * PRICE["cache_write_1h"]
        ) / 1_000_000
    return c


def reconstruct(run_dir: Path, session_dir: Path, buffer_min: int = 5):
    manifest = json.loads((run_dir / "manifest.json").read_text())
    t0 = parse_ts(manifest["started_at"])
    rows = [json.loads(l) for l in (run_dir / "results.jsonl").read_text().splitlines() if l.strip()]

    windows = []
    acc = 0.0
    for r in rows:
        start = t0 + timedelta(seconds=acc)
        acc += r["wall_sec"]
        end = t0 + timedelta(seconds=acc)
        windows.append({"task": r["task_id"], "arm": r["arm"], "start": start, "end": end})
    t_end = windows[-1]["end"] if windows else t0

    sessions = []
    for f in session_dir.glob("*.jsonl"):
        msgs = []
        for ln in f.read_text(encoding="utf-8", errors="ignore").splitlines():
            try:
                msgs.append(json.loads(ln))
            except Exception:
                continue
        times = [parse_ts(m["timestamp"]) for m in msgs if m.get("timestamp")]
        if not times:
            continue
        ts, te = min(times), max(times)
        if te < t0 - timedelta(minutes=1) or ts > t_end + timedelta(minutes=buffer_min):
            continue
        c = session_cost(msgs)
        if c > 0:
            sessions.append({"start": ts, "end": te, "cost": c})

    per = defaultdict(float)
    alloc = 0.0
    for s in sessions:
        dur = (s["end"] - s["start"]).total_seconds()
        if dur <= 0:
            for w in windows:
                if w["start"] <= s["start"] <= w["end"]:
                    per[(w["task"], w["arm"])] += s["cost"]
                    alloc += s["cost"]
                    break
            continue
        ov = []
        for w in windows:
            o = (min(w["end"], s["end"]) - max(w["start"], s["start"])).total_seconds()
            if o > 0:
                ov.append((w, o))
        to = sum(o for _, o in ov)
        if to <= 0:
            continue
        for w, o in ov:
            share = s["cost"] * (o / to)
            per[(w["task"], w["arm"])] += share
            alloc += share

    return rows, dict(per), sum(s["cost"] for s in sessions), alloc


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--run", required=True, help="arc_pilot run dir")
    ap.add_argument("--session-dir", required=True, help="~/.claude/projects/<slug> dir")
    ap.add_argument("--arm-a", default="baseline")
    ap.add_argument("--arm-b", default="g6_full")
    args = ap.parse_args()

    run_dir = Path(args.run)
    session_dir = Path(os.path.expanduser(args.session_dir))
    rows, per, total, alloc = reconstruct(run_dir, session_dir)

    by_arm = defaultdict(list)
    for r in rows:
        by_arm[r["arm"]].append(r["task_id"])
    paired = sorted(set(by_arm[args.arm_a]) & set(by_arm[args.arm_b]))

    a_total = sum(per.get((t, args.arm_a), 0.0) for t in paired)
    b_total = sum(per.get((t, args.arm_b), 0.0) for t in paired)
    correct = {(r["task_id"], r["arm"]): r["correct"] for r in rows}
    a_c = sum(1 for t in paired if correct.get((t, args.arm_a)))
    b_c = sum(1 for t in paired if correct.get((t, args.arm_b)))

    print(f"run={run_dir.name}  paired_n={len(paired)}")
    print(f"sessions_total=${total:.2f}  allocated=${alloc:.2f}  ({100*alloc/total:.1f}%)")
    print(f"{args.arm_a}: total=${a_total:.2f} mean=${a_total/len(paired):.3f} correct={a_c}/{len(paired)}")
    print(f"{args.arm_b}: total=${b_total:.2f} mean=${b_total/len(paired):.3f} correct={b_c}/{len(paired)}")

    diffs = [per.get((t, args.arm_b), 0.0) - per.get((t, args.arm_a), 0.0) for t in paired]
    try:
        from scipy.stats import wilcoxon

        W, p = wilcoxon(diffs, zero_method="wilcox", alternative="two-sided")
        print(f"Wilcoxon two-sided: W={W:.1f}  p={p:.4f}")
        W1, p1 = wilcoxon(diffs, alternative="less")
        print(f"Wilcoxon one-sided ({args.arm_b} < {args.arm_a}): W={W1:.1f}  p={p1:.4f}")
    except Exception as exc:
        print(f"wilcoxon unavailable: {exc}")


if __name__ == "__main__":
    main()