#!/usr/bin/env python3
"""Persistent-graft check: load a SAVED adapter graft in a fresh process and evaluate the recipient with it
on freshly generated held-out splits (new seed). This is the only script that may claim 'saved modification';
the adapter was fitted in scripts/run_transplant.py and written to experiments/summaries/<run>/persistent/.
Usage: python scripts/eval_persistent.py --run <graft run id> [--seed-base 70000]
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import numpy as np

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))
sys.path.insert(0, str(ROOT / "scripts"))

from train_models import ROLES, model_dir  # noqa: E402
from wp_ai_0001.alignment.maps import AffineMap  # noqa: E402
from wp_ai_0001.datasets import opseq  # noqa: E402
from wp_ai_0001.datasets.common import Split  # noqa: E402
from wp_ai_0001.instrumentation.provenance import utc_now  # noqa: E402
from wp_ai_0001.models.toy_transformer import resolve_device  # noqa: E402
from wp_ai_0001.statistics.bootstrap import cluster_bootstrap_mean  # noqa: E402
from wp_ai_0001.training.train_toy import load as load_model  # noqa: E402
from wp_ai_0001.transplantation.experiment import build_eval_splits, evaluate_condition, omega_step_acc  # noqa: E402
from wp_ai_0001.transplantation.grafts import AdapterGraft, GraftSpec  # noqa: E402


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--run", required=True)
    ap.add_argument("--seed-base", type=int, default=70_000)
    args = ap.parse_args()
    pdir = ROOT / "experiments" / "summaries" / args.run / "persistent"
    z = np.load(pdir / "adapter_graft.npz")
    hook, rkey = str(z["hook"]), str(z["recipient"])
    role, seed = rkey.rsplit("_s", 1)
    device = resolve_device("auto")
    model, vocab = load_model(model_dir(role, int(seed)), device)
    task = ROLES[role][0]
    ex = Split.load(model_dir(role, int(seed)) / "train_split.json").item_hashes()
    splits = build_eval_splits(ROLES["donor"][0], ex, args.seed_base, 1000, 100)
    names = ["target_iid", "control_iid"]
    base = evaluate_condition(model, None, None, splits, vocab, device, names)
    ad = AffineMap("adapter_ridge", z["W"], z["b"])
    graft = evaluate_condition(model, hook, AdapterGraft(GraftSpec("adapter", hook, None, 1.0, "all"), ad), splits, vocab, device, names)
    rep = {"utc": utc_now(), "run": args.run, "recipient": rkey, "hook": hook, "fresh_process": True, "splits": {k: v.hash for k, v in splits.items() if k in names}}
    for n in names:
        rep[n] = {"base": base[n].accuracy, "with_saved_adapter": graft[n].accuracy,
                  "ci95_with_saved_adapter": cluster_bootstrap_mean(graft[n].correct.astype(float), graft[n].clusters, 2000).to_dict(),
                  "omega_step_base": float(np.nanmean(omega_step_acc(splits[n], base[n]))) if n == "target_iid" else None,
                  "omega_step_with_saved_adapter": float(np.nanmean(omega_step_acc(splits[n], graft[n]))) if n == "target_iid" else None}
    (pdir / "fresh_process_eval.json").write_text(json.dumps(rep, indent=2))
    print(json.dumps(rep, indent=2))


if __name__ == "__main__":
    main()