#!/usr/bin/env python3
"""Join the official T-Rex predictions (16 steps @ 30 Hz = 0.5 s) with the other models on the same anchors.
All models are compared on the overlapping window k = 0..10 of the 20 Hz grid (0.05 s .. 0.5 s).
    python scripts/eval/summarize_official.py
"""
import json, sys
from pathlib import Path
import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parent)); sys.argv = ["x"]
import trex_ablation_eval as tae
import openloop_compare as oc

OUT = Path("ablation_runs/openloop_compare"); K = 11
r = np.load(OUT / "pred_raptor.npz"); d = np.load(OUT / "pred_dreamzero.npz"); t = np.load(OUT / "../openloop_compare_trex/pred_raptor.npz")
o = np.load(OUT / "../openloop_compare_official/pred_official.npz")
A = json.load(open(OUT / "anchors.json")); anchors = [(int(e), float(tt)) for e, tt in A["anchors"]]
gt = r["gt"].astype(np.float64)[:, :K]
assert np.allclose(r["gt"], d["gt"]) and np.allclose(r["gt"], t["gt"])
assert all(int(m[0]) == e and abs(float(m[1]) - tt) < 1e-6 for m, (e, tt) in zip(o["meta"], anchors)), "official run used different anchors"

# official: interpolate the 30 Hz chunk onto the 20 Hz grid (k = 0..10 -> 0.05 .. 0.5 s after the anchor)
def to_grid(pred30, t30):
    out = np.empty((len(pred30), K, 62))
    for n in range(len(pred30)):
        anchor = anchors[n][1]; tg = anchor + np.arange(K) / 20.0
        out[n] = np.stack([np.interp(tg, t30[n], pred30[n][:, j]) for j in range(62)], -1)
    return out
cache = {}; hold = []
for ep, tt in anchors:
    if ep not in cache: b = tae.ChunkBuilder(A["dataset_root"], ep); cache[ep] = (b.timestamps, b.state)
    ts, st = cache[ep]; hold.append(np.repeat(st[int(np.argmin(np.abs(ts - tt)))][None], K, 0))
runs = {
    "RAPTOR full (joint refiner, 100k)": r["pred_cascade"][:, :K],
    "RAPTOR coarse only": r["pred_coarse_only"][:, :K],
    "T-Rex-style on our backbone (VQ force, 96k)": t["pred_cascade"][:, :K],
    "Official T-Rex midtrain (cascaded, tactile expert)": to_grid(o["pred_cascaded"], o["t30"]),
    "Official T-Rex midtrain (tactile expert disabled)": to_grid(o["pred_notactile"], o["t30"]),
    "DreamZero (ckpt-16000)": d["pred"][:, :K],
    "hold current pose": np.stack(hold),
}
H = np.r_[9:31, 40:62]; P = 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 = gt[:, K - 1, H] - gt[:, 0, H], gt[:, K - 1, P] - gt[:, 0, P]
mh = np.linalg.norm(gdh, axis=-1) >= np.quantile(np.linalg.norm(gdh, axis=-1), .75)
mp = np.linalg.norm(gdp, axis=-1) >= np.quantile(np.linalg.norm(gdp, axis=-1), .75)
lines = [f"# Open-loop comparison incl. official T-Rex — {len(anchors)} anchors, window 0.05–0.5 s (11 steps @ 20 Hz)", "",
         "| model | EEF pos err [cm] | EEF rot err [deg] | hand RMSE [rad] | pos @0.5 s | hand @0.5 s | hand disp cos (high-motion) | EEF disp cos (high-motion) | hand jitter ×GT |",
         "|---|---|---|---|---|---|---|---|---|"]
per = {}
gj = (np.diff(gt[:, :, H], 2, axis=1) ** 2).mean()
for name, p in runs.items():
    p = p.astype(np.float64); e = oc.errors(p, gt); per[name] = {m: e[m].mean(1) for m in ("pos", "rot", "hand")}
    dh, dp = p[:, K - 1, H] - p[:, 0, H], p[:, K - 1, P] - p[:, 0, P]
    jit = (np.diff(p[:, :, H], 2, axis=1) ** 2).mean() / gj
    lines.append(f"| {name} | {100 * e['pos'].mean():.2f} | {e['rot'].mean():.2f} | {e['hand'].mean():.4f} | {100 * e['pos'][:, -1].mean():.2f} | {e['hand'][:, -1].mean():.4f} | "
                 f"{np.nanmean(cos(dh, gdh)[mh]):.2f} | {np.nanmean(cos(dp, gdp)[mp]):.2f} | {jit:.0f} |")
def paired(a, b):
    s = []
    for m, u, f in (("pos", "cm", 100), ("rot", "deg", 1), ("hand", "rad", 1)):
        df = (per[a][m] - per[b][m]) * f; s.append(f"{m} {df.mean():+.3f} ± {df.std(ddof=1) / np.sqrt(len(df)):.3f} {u} ({100 * (df < 0).mean():.0f}% anchors)")
    return "; ".join(s)
ref = "RAPTOR full (joint refiner, 100k)"
lines += ["", "Paired differences vs RAPTOR full (negative = RAPTOR better):", ""]
for name in runs:
    if name != ref: lines.append(f"- **RAPTOR full − {name}**: {paired(ref, name)}")
lines += ["", f"GT motion in the window: EEF {100 * np.linalg.norm(gdp, axis=-1).mean():.2f} cm, hand {np.linalg.norm(gdh, axis=-1).mean():.3f} rad; "
          f"high-motion subsets = top quartile ({mh.sum()} / {mp.sum()} anchors).",
          f"Official T-Rex inference time {o['sec'].mean():.2f} s per anchor per variant.",
          "", "Notes: the official midtrain checkpoint is evaluated as released (no post-training on this data; the public T-Rex dataset is its own collection, "
          "so these episodes may or may not be in its midtrain corpus). Its head/wrist frames come from the 320x180 re-encoded copy, upscaled to 384x288. "
          "RAPTOR / T-Rex-style / DreamZero were trained on these episodes."]
text = "\n".join(lines); print(text)
(OUT / "../openloop_compare_official/summary_official.md").write_text(text + "\n")