"""Session-reset baseline (reviewer ask, round 7).

Separates CROSS-SESSION persistence from mere WITHIN-SESSION liveness: a
"reset" arm ingests live during a session but is restored to its end-of-warmup
state at every session boundary. Then, post-warmup:
    reset - static      = value of within-session liveness alone
    personal - reset    = value of persisting across sessions
Replays the phase-2 seed-0 stream against cached targets (no GPU/model calls).

Usage: python -m harness.reset_arm  (from code/)
Writes results/phase2_reset_arm.json.
"""
from __future__ import annotations

import copy
import json
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 = StaticGlobal()
    personal = PersonalMemory(capacity=48, eviction="lru")
    reset = PersonalMemory(capacity=48, eviction="lru")
    warm_snapshot = None

    agg = {n: defaultdict(list) for n in ("static", "personal", "reset")}
    cur_session = -1
    for ins in instances:
        tgt = targets.get(ins.query)
        if tgt is None:
            continue
        if ins.session != cur_session:
            cur_session = ins.session
            if cur_session == 1:
                static.freeze()
                warm_snapshot = copy.deepcopy(reset)   # end-of-warmup state
            elif cur_session > 1:
                reset = copy.deepcopy(warm_snapshot)   # wipe session memory
        for name, a in (("static", static), ("personal", personal),
                        ("reset", reset)):
            draft = a.draft(ins.query, ins.functions, ins.user_id, embedder)
            agg[name][ins.session].append(metrics.score(draft, tgt))
        cname, cargs = _parse_target(tgt)
        for a in (static, personal, reset):
            a.observe(ins.query, ins.functions, ins.user_id, cname, cargs,
                      embedder)
            if isinstance(a, PersonalMemory) and ins.session == 0:
                a.seed_shared(ins.query, cname, cargs, embedder)

    out = {}
    for name in agg:
        scores = [x for s, xs in agg[name].items() if s > 0 for x in xs]
        n = len(scores)
        out[name] = {
            "n": n,
            "MAT": round(sum(x["accept_length"] for x in scores) / n, 3),
            "exact_rate": round(sum(1 for x in scores if x["exact"]) / n, 4),
            "by_session": {str(s): round(sum(x["accept_length"] for x in xs)
                                         / len(xs), 2)
                           for s, xs in sorted(agg[name].items())},
        }
    result = {"config": {"users": 40, "tasks_per_user": 15, "sessions": 12,
                         "queries_per_session": 6, "seed": 0, "capacity": 48,
                         "targets": "phase2_targets_seed0.json (cached)"},
              "arms": out,
              "decomposition": {
                  "within_session_liveness (reset - static)":
                      round(out["reset"]["MAT"] - out["static"]["MAT"], 3),
                  "cross_session_persistence (personal - reset)":
                      round(out["personal"]["MAT"] - out["reset"]["MAT"], 3),
              }}
    (RESULTS / "phase2_reset_arm.json").write_text(json.dumps(result,
                                                              indent=2))
    print(json.dumps(result["arms"]["reset"], indent=1))
    print(json.dumps(result["decomposition"], indent=1))


if __name__ == "__main__":
    main()