"""Bootstrap/paired CIs + variance decomposition (external-review item 5).

Deterministically replays the phase-2 seed-0 stream against CACHED targets
(no model calls) to obtain per-instance paired (static, personal) accepts,
then reports:
  - paired bootstrap 95% CI for the post-warmup MAT gap and % gap
  - per-user, per-session, per-signature-task gap distributions
Writes results/phase2_bootstrap_ci.json.
"""
from __future__ import annotations

import json
import random
from collections import defaultdict
from pathlib import Path

from . import metrics
from .data import load_bfcl
from .memory import Embedder, PersonalMemory, StaticGlobal
from .run_accept import MODEL_PATH, _parse_target
from .simulate import build_users

ROOT = Path(__file__).resolve().parent.parent
RESULTS = ROOT / "results"


def main():
    metrics.get_tokenizer(MODEL_PATH)
    tasks = load_bfcl()
    embedder = Embedder()
    instances = build_users(tasks, n_users=40, tasks_per_user=15,
                            n_sessions=12, queries_per_session=6, seed=0)
    instances.sort(key=lambda x: (x.session, x.user_id))
    targets = json.loads((RESULTS / "phase2_targets_seed0.json").read_text())

    static, personal = StaticGlobal(), PersonalMemory(capacity=48,
                                                      eviction="lru")
    rows = []          # (user, session, sig, a_static, a_personal)
    cur = -1
    for ins in instances:
        tgt = targets.get(ins.query)
        if tgt is None:
            continue
        if ins.session != cur:
            cur = ins.session
            if cur == 1:
                static.freeze()
        a_s = metrics.score(static.draft(ins.query, ins.functions,
                                         ins.user_id, embedder),
                            tgt)["accept_length"]
        a_p = metrics.score(personal.draft(ins.query, ins.functions,
                                           ins.user_id, embedder),
                            tgt)["accept_length"]
        rows.append((ins.user_id, ins.session, ins.signature_id, a_s, a_p))
        cname, cargs = _parse_target(tgt)
        for a in (static, personal):
            a.observe(ins.query, ins.functions, ins.user_id, cname, cargs,
                      embedder)
        if ins.session == 0:
            personal.seed_shared(ins.query, cname, cargs, embedder)

    post = [r for r in rows if r[1] > 0]
    n = len(post)
    mat_s = sum(r[3] for r in post) / n
    mat_p = sum(r[4] for r in post) / n

    rng = random.Random(0)
    B = 10_000
    gaps, pct = [], []
    for _ in range(B):
        idx = [rng.randrange(n) for _ in range(n)]
        s = sum(post[i][3] for i in idx) / n
        p = sum(post[i][4] for i in idx) / n
        gaps.append(p - s)
        pct.append(100 * (p - s) / s)
    gaps.sort(); pct.sort()
    ci = lambda xs: (round(xs[int(0.025 * B)], 3), round(xs[int(0.975 * B)], 3))

    def group_gaps(key):
        g = defaultdict(lambda: [0.0, 0.0, 0])
        for r in post:
            k = key(r); g[k][0] += r[3]; g[k][1] += r[4]; g[k][2] += 1
        vals = sorted((v[1] - v[0]) / v[2] for v in g.values())
        m = len(vals)
        return {"n_groups": m,
                "mean_gap": round(sum(vals) / m, 3),
                "min": round(vals[0], 3), "p25": round(vals[m // 4], 3),
                "median": round(vals[m // 2], 3),
                "p75": round(vals[3 * m // 4], 3), "max": round(vals[-1], 3),
                "groups_with_negative_gap": sum(1 for v in vals if v < 0)}

    out = {
        "config": {"seed": 0, "targets": "phase2_targets_seed0.json (cached)",
                   "n_post_warmup_paired": n, "bootstrap_resamples": B},
        "MAT": {"static": round(mat_s, 3), "personal": round(mat_p, 3),
                "gap": round(mat_p - mat_s, 3),
                "gap_pct": round(100 * (mat_p - mat_s) / mat_s, 2)},
        "paired_bootstrap_95CI": {"gap_MAT": ci(gaps), "gap_pct": ci(pct)},
        "per_user_gap": group_gaps(lambda r: r[0]),
        "per_session_gap": group_gaps(lambda r: r[1]),
        "per_task_gap": group_gaps(lambda r: r[2]),
    }
    (RESULTS / "phase2_bootstrap_ci.json").write_text(json.dumps(out,
                                                                 indent=2))
    print(json.dumps(out, indent=2))


if __name__ == "__main__":
    main()