"""Phase 4 MAIN TABLE Part B — freshness-over-time curve on tau2-bench.

Packages the existing multi-session tau2-bench longitudinal result as an explicit
freshness curve: per-session MAT for each arm (no_memory / toolspec /
static_global / personal_memory), x = session index, one line per arm. This is
the plot that shows WHY persistence is necessary: the fixed store degrades in
relative terms as novel per-user calls accumulate while the live store holds.

CPU-only: replays the frozen tau2 decision points against their cached on-policy
trace targets (no server calls). Emits results/phase4_freshness_curve.json.

Run from the repo root:  python -m harness.phase4_partb
"""
from __future__ import annotations

import json
import os
from collections import defaultdict
from pathlib import Path

from . import metrics
from .data import Task
from .memory import (Embedder, NoMemory, PersonalMemory, StaticGlobal,
                     ToolSpecBaseline)
from .run_accept import _parse_target
from .simulate import build_users

ROOT = Path(__file__).resolve().parent.parent
RESULTS = ROOT / "results"
# Tokenizer for the token-LCP accept metric: HF hub id by default;
# override with a local snapshot path if running offline.
MODEL_PATH = os.environ.get("SPECMEM_TOKENIZER", "openai/gpt-oss-120b")
DOMAINS = ("airline", "retail", "telecom")
ARMS = ["no_memory", "toolspec", "static_global", "personal_memory"]


def main():
    metrics.get_tokenizer(MODEL_PATH)
    dp = [json.loads(l) for l in
          (RESULTS / "tau2_live_decision_points.jsonl").read_text().splitlines()]
    tools = {d: json.loads((ROOT / "data" / "tau2" /
                            f"tools_{d}.json").read_text()) for d in DOMAINS}
    tasks = [Task(id=r["id"], query=r["query"], functions=tools[r["domain"]],
                  origin_id=r["id"]) for r in dp]
    targets = {r["query"]: r["target"] for r in dp}
    emb = Embedder()

    per_seed = {a: defaultdict(list) for a in ARMS}   # arm -> session -> [seed MAT]
    for sd in (0, 1, 2):
        inst = build_users(tasks, n_users=40, tasks_per_user=15, n_sessions=12,
                           queries_per_session=6, seed=sd, perturb_prob=0.0)
        inst.sort(key=lambda x: (x.session, x.user_id))
        arms = [NoMemory(), ToolSpecBaseline(), StaticGlobal(),
                PersonalMemory(capacity=48, eviction="lru")]
        agg = {a.name: defaultdict(list) for a in arms}
        cur = -1
        for ins in inst:
            tgt = targets.get(ins.query)
            if tgt is None:
                continue
            if ins.session != cur:
                cur = ins.session
                if cur == 1:
                    for a in arms:
                        if hasattr(a, "freeze"):
                            a.freeze()
            for a in arms:
                agg[a.name][ins.session].append(metrics.score(
                    a.draft(ins.query, ins.functions, ins.user_id, emb), tgt))
            cn, ca = _parse_target(tgt)
            for a in arms[1:]:
                a.observe(ins.query, ins.functions, ins.user_id, cn, ca, emb)
                if isinstance(a, PersonalMemory) and ins.session == 0:
                    a.seed_shared(ins.query, cn, ca, emb)
        for a in arms:
            for s, xs in agg[a.name].items():
                per_seed[a.name][s].append(
                    sum(x["accept_length"] for x in xs) / len(xs))
        print(f"seed {sd} done", flush=True)

    curve = {a: {str(s): round(sum(v) / len(v), 3)
                 for s, v in sorted(per_seed[a].items())} for a in ARMS}
    # relative freshness: personal advantage over static per session (post-warmup)
    rel = {}
    for s in sorted(per_seed["personal_memory"]):
        if s == 0:
            continue
        pm = sum(per_seed["personal_memory"][s]) / len(per_seed["personal_memory"][s])
        sg = sum(per_seed["static_global"][s]) / len(per_seed["static_global"][s])
        rel[str(s)] = round(100 * (pm - sg) / sg, 1)
    out = {"arms": ARMS, "per_session_MAT": curve,
           "personal_over_static_pct_by_session": rel,
           "note": ("tau2-bench 3-domain frozen decision points, cached "
                    "on-policy trace targets, 3 seeds; session 0 = warmup. "
                    "Fixed stores (toolspec/static) plateau/degrade while "
                    "personal (live) holds -> the freshness curve.")}
    (RESULTS / "phase4_freshness_curve.json").write_text(json.dumps(out, indent=2))
    print(json.dumps({"per_session_MAT": curve,
                      "personal_over_static_pct_by_session": rel}, indent=1))


if __name__ == "__main__":
    main()