#!/usr/bin/env python3
"""Multi-seed open-loop test: is the RAPTOR cascade's error bias or sampling noise?
Samples S action chunks per anchor (different seeds), scores each sample, the per-anchor mean of
samples, and the hold-current-pose baseline, on the anchors of ablation_runs/openloop_compare.
    python scripts/eval/raptor_multiseed.py <checkpoint> <out_npz> [seeds=8]
"""
import json, sys, time
from pathlib import Path
import numpy as np, torch

ROOT = Path(__file__).resolve().parents[2]
for sub in ("scripts", "scripts/eval"):
    sys.path.insert(0, str(ROOT / sub))
import trex_ablation_eval as tae
import openloop_compare as oc
from groot.vla.model.trex_track_force.runtime import delta_base_to_absolute
from groot.vla.model.trex_track_force.vla import TrexTrackForceVLA

ckpt, out_npz = Path(sys.argv[1]), Path(sys.argv[2]); S = int(sys.argv[3]) if len(sys.argv) > 3 else 8
A = json.load(open("ablation_runs/openloop_compare/anchors.json")); root = A["dataset_root"]; anchors = [(int(e), float(t)) for e, t in A["anchors"]]
device = torch.device("cuda")
model = TrexTrackForceVLA.load_lora(str(ckpt)).to(device=device, dtype=torch.bfloat16).eval(); policy = model.action_head
transform_eval, _ = tae.load_pipeline(ckpt, training=False)
builders = {}; preds, gts, holds = [], [], []
for i, (ep, anchor) in enumerate(anchors):
    if ep not in builders: builders[ep] = tae.ChunkBuilder(root, ep)
    b = builders[ep]; gt_eef, _ = oc.ground_truth(b, anchor)
    raw = b.build(anchor, blocks=1, prompt=oc.episode_task(b), with_deform=False, history_only_video=True)
    reference = np.asarray(raw["state.eef62"], np.float64).reshape(-1, 62)[0]; raw.pop("action.eef62")
    collated = transform_eval(tae.unsqueeze_dict_values(dict(raw))); batch = {}
    for k, v in collated.items():
        if torch.is_tensor(v): v = v.to(device=device, dtype=torch.bfloat16) if v.is_floating_point() else v.to(device)
        batch[k] = v
    t0 = time.perf_counter(); samples = []
    for s in range(S):
        with torch.inference_mode():
            res = policy.sample(batch, seed=500 + i + 1000 * (s + 1), run_force_refinement=True)
        samples.append(delta_base_to_absolute(reference, b.stats.denormalize_action(res["action_pred"][0, :16, :62].float().cpu().numpy())))
    preds.append(np.stack(samples)); gts.append(gt_eef); holds.append(np.repeat(reference[None], 16, 0))
    print(f"[{i + 1}/{len(anchors)}] ep{ep} t={anchor:.2f}s  {S} samples in {time.perf_counter() - t0:.1f}s", flush=True)
P = np.stack(preds); G = np.stack(gts); Hd = np.stack(holds)       # P: [N,S,16,62]
np.savez(out_npz, pred=P, gt=G, hold=Hd)
e_single = oc.errors(P[:, 0].astype(np.float64), G); e_mean = oc.errors(P.mean(1).astype(np.float64), G); e_hold = oc.errors(Hd, G)
spread = {m: float(np.mean([oc.errors(P[:, s].astype(np.float64), P.mean(1).astype(np.float64))[m].mean() for s in range(S)])) for m in ("pos", "rot", "hand")}
H = np.r_[9:31, 40:62]; Pp = np.r_[0:3, 31:34]
cos = lambda a, b: ((a * b).sum(-1) / (np.linalg.norm(a, axis=-1) * np.linalg.norm(b, axis=-1) + 1e-9))
gdh, gdp = G[:, 15, H] - G[:, 0, H], G[:, 15, Pp] - G[:, 0, Pp]; M = P.mean(1)
print("\n## Multi-seed test on %s (%d anchors x %d samples)\n" % (ckpt, len(anchors), S))
print("| prediction | EEF pos err [cm] | EEF rot err [deg] | hand RMSE [rad] | hand disp cos | EEF disp cos |"); print("|---|---|---|---|---|---|")
for name, e, pr in (("single sample (seed 0)", e_single, P[:, 0]), (f"mean of {S} samples", e_mean, M), ("hold current pose", e_hold, Hd)):
    dh, dp = pr[:, 15, H] - pr[:, 0, H], pr[:, 15, Pp] - pr[:, 0, Pp]
    print(f"| {name} | {100 * e['pos'].mean():.2f} | {e['rot'].mean():.2f} | {e['hand'].mean():.4f} | {np.nanmean(cos(dh, gdh)):.2f} | {np.nanmean(cos(dp, gdp)):.2f} |")
print(f"\nsample-to-sample spread around the per-anchor mean: pos {100 * spread['pos']:.2f} cm, rot {spread['rot']:.2f} deg, hand {spread['hand']:.4f} rad "
      f"(compare: GT motion within a chunk {100 * np.linalg.norm(gdp, axis=-1).mean():.2f} cm)")