#!/usr/bin/env python3
"""Open-loop comparison of DreamZero and RAPTOR/FACT on the same T-Rex episodes, in physical units.

Written 2026-10-02 for this comparison; it is NOT part of the upstream dreamzero repo. It reuses
  * RAPTOR side   : ChunkBuilder / load_pipeline / _sample from scripts/eval/trex_ablation_eval.py
  * DreamZero side: GrootSimPolicy (upstream) + TrexEpisode / build_obs from scripts/inference_trex.py
  * FK            : joint58_to_eef62_batch from scripts/data/build_trex_track_force_v2.py — the same FK
                    that produced the dataset's ``observation.state_eef62`` / ``action.eef62_absolute``

Protocol. For every (episode, anchor) both models get ONE current frame per camera and the current
state (RAPTOR additionally its 16 past track frames and force history) and predict the next 0.8 s.
Predictions are compared with the recorded teleop targets on a common 20 Hz grid (16 steps):
    pos_err  [m]   per arm  = || p_pred - p_gt ||
    rot_err  [deg] per arm  = geodesic angle of R_pred^T R_gt
    hand_err [rad] per hand = RMSE over the 22 hand joints
DreamZero predicts 24 joint-space actions at 30 Hz; they are linearly interpolated onto the 20 Hz grid
and passed through FK, so arms are compared in the same end-effector space as RAPTOR.

Usage:
    python scripts/eval/openloop_compare.py run --model raptor    --checkpoint <dir> --out <dir>
    python scripts/eval/openloop_compare.py run --model dreamzero --checkpoint <dir> --out <dir>
    python scripts/eval/openloop_compare.py summarize --out <dir>
"""
from __future__ import annotations

import argparse
import json
import os
import sys
import time
from pathlib import Path

import numpy as np
import torch

ROOT = Path(__file__).resolve().parents[2]
for sub in ("scripts", "scripts/eval", "scripts/data"):
    sys.path.insert(0, str(ROOT / sub))

import trex_ablation_eval as tae  # noqa: E402  (RAPTOR helpers)
from groot.vla.model.trex_track_force.dataset import (  # noqa: E402
    eef62_delta_base,
    rotation_6d_to_matrix,
)
from groot.vla.model.trex_track_force.runtime import delta_base_to_absolute  # noqa: E402

STEPS = 16
RATE = 20.0
ARM = {"left": slice(0, 9), "right": slice(31, 40)}
HAND = {"left": slice(9, 31), "right": slice(40, 62)}
DZ_ACTION_KEYS = ("action.left_arm", "action.left_hand", "action.right_arm", "action.right_hand")


# ----------------------------------------------------------------------------- anchors / ground truth
def select_anchors(dataset_root: str, n_episodes: int, n_anchors: int, n_total: int = 100):
    episodes = sorted(set(int(round(e)) for e in np.linspace(0, n_total - 1, n_episodes)))
    anchors = []
    for ep in episodes:
        builder = tae.ChunkBuilder(dataset_root, ep)
        valid = builder.valid_anchor_times(1)
        ts = builder.timestamps
        valid = [t for t in valid if int(np.argmin(np.abs(ts - t))) + 24 <= len(ts) - 1]
        picks = np.linspace(0, len(valid) - 1, n_anchors + 2)[1:-1]
        anchors += [(ep, float(valid[int(round(i))])) for i in picks]
    return anchors


def episode_task(builder) -> str:
    try:
        return str(builder.episode.table.column("annotation.task")[0].as_py())
    except Exception:
        with open(Path(builder.root) / "meta" / "tasks.jsonl") as handle:
            return json.loads(handle.readline())["task"]


def ground_truth(builder, anchor: float):
    sel = tae._sample(builder.timestamps, anchor, range(STEPS), RATE)
    joints = tae._column(builder.episode, "action")
    return builder.action_abs[sel.indices].astype(np.float64), joints[sel.indices].astype(np.float64)


def load_or_make_anchors(args):
    path = Path(args.out) / "anchors.json"
    if path.exists():
        return [(int(e), float(t)) for e, t in json.load(open(path))["anchors"]]
    anchors = select_anchors(args.dataset_root, args.episodes, args.anchors)
    path.parent.mkdir(parents=True, exist_ok=True)
    json.dump({"dataset_root": args.dataset_root, "anchors": anchors}, open(path, "w"), indent=1)
    return anchors


# ----------------------------------------------------------------------------- RAPTOR
def run_raptor(args, anchors):
    from groot.vla.model.trex_track_force.vla import TrexTrackForceVLA

    device = torch.device("cuda")
    ckpt = Path(args.checkpoint)
    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)
    modes = [("cascade", True), ("coarse_only", False)] if policy.config.use_force else [("coarse_only", False)]
    builders: dict[int, object] = {}
    preds = {m: [] for m, _ in modes}
    gts, gtj, times = [], [], []
    for i, (ep, anchor) in enumerate(anchors):
        if ep not in builders:
            builders[ep] = tae.ChunkBuilder(args.dataset_root, ep)
        builder = builders[ep]
        gt_eef, gt_joint = ground_truth(builder, anchor)
        raw = builder.build(anchor, blocks=1, prompt=episode_task(builder), with_deform=False,
                            history_only_video=True)
        reference = np.asarray(raw["state.eef62"], np.float64).reshape(-1, 62)[0]
        if i == 0:  # round-trip check of the delta-base math used below
            rt = delta_base_to_absolute(reference, eef62_delta_base(reference, gt_eef))
            print(f"delta-base round trip max |err| = {np.abs(rt - gt_eef).max():.2e}", flush=True)
        raw.pop("action.eef62")
        collated = transform_eval(tae.unsqueeze_dict_values(dict(raw)))
        batch = {}
        for key, value in collated.items():
            if torch.is_tensor(value):
                value = value.to(device=device, dtype=torch.bfloat16) if value.is_floating_point() else value.to(device)
            batch[key] = value
        t0 = time.perf_counter()
        for mode, refine in modes:
            with torch.inference_mode():
                result = policy.sample(batch, seed=500 + i, run_force_refinement=refine)
            pred_norm = result["action_pred"][0, :STEPS, :62].float().cpu().numpy()
            delta = builder.stats.denormalize_action(pred_norm)
            preds[mode].append(delta_base_to_absolute(reference, delta))
        times.append(time.perf_counter() - t0)
        gts.append(gt_eef); gtj.append(gt_joint)
        print(f"[raptor] {i + 1}/{len(anchors)} ep{ep} t={anchor:.2f}s  {times[-1]:.1f}s", flush=True)
    out = {f"pred_{m}": np.stack(v) for m, v in preds.items()}
    np.savez(Path(args.out) / "pred_raptor.npz", gt=np.stack(gts), gt_joint=np.stack(gtj),
             time=np.asarray(times), **out)


# ----------------------------------------------------------------------------- DreamZero
def run_dreamzero(args, anchors):
    import torch.distributed as dist
    import inference_trex as it
    from build_trex_track_force_v2 import joint58_to_eef62_batch
    from groot.vla.data.schema import EmbodimentTag
    from groot.vla.model.n1_5.sim_policy import GrootSimPolicy

    if not dist.is_initialized():
        os.environ.setdefault("MASTER_ADDR", "localhost")
        os.environ.setdefault("MASTER_PORT", str(29500 + os.getpid() % 1000))
        dist.init_process_group(backend="gloo", world_size=1, rank=0)
    policy = GrootSimPolicy(embodiment_tag=EmbodimentTag.TREX, model_path=args.checkpoint, device="cuda:0")
    height, width = it.get_expected_video_resolution(policy)
    print(f"DreamZero input video {height}x{width}", flush=True)
    builders: dict[int, object] = {}
    episodes: dict[int, object] = {}
    preds, preds_joint, gts, gtj, times = [], [], [], [], []
    for i, (ep, anchor) in enumerate(anchors):
        if ep not in builders:
            builders[ep] = tae.ChunkBuilder(args.dataset_root, ep)
            episodes[ep] = it.TrexEpisode(args.dataset_root, ep)
        builder, episode = builders[ep], episodes[ep]
        gt_eef, gt_joint = ground_truth(builder, anchor)
        ts = builder.timestamps
        row = int(np.argmin(np.abs(ts - anchor)))
        obs = it.build_obs(episode, [row], episode_task(builder), height, width)  # 1 frame -> causal reset
        torch.manual_seed(500 + i)
        t0 = time.perf_counter()
        with torch.inference_mode():
            result, _ = policy.lazy_joint_forward_causal(it.Batch(obs=obs))
        times.append(time.perf_counter() - t0)
        act = result.act
        joints24 = np.concatenate(
            [np.asarray(act[k], dtype=np.float64).reshape(-1, len(range(*{"action.left_arm": (0, 7), "action.left_hand": (7, 29), "action.right_arm": (29, 36), "action.right_hand": (36, 58)}[k])))
             for k in DZ_ACTION_KEYS], axis=-1)
        n = joints24.shape[0]
        t_pred = ts[row:row + n]
        if len(t_pred) < n:  # episode tail: extend with the nominal 30 Hz spacing
            t_pred = np.concatenate([t_pred, t_pred[-1] + (np.arange(n - len(t_pred)) + 1) / 30.0])
        t_grid = anchor + np.arange(STEPS) / RATE
        joints16 = np.stack([np.interp(t_grid, t_pred, joints24[:, d]) for d in range(58)], axis=-1)
        preds.append(joint58_to_eef62_batch(joints16))
        preds_joint.append(joints16)
        gts.append(gt_eef); gtj.append(gt_joint)
        print(f"[dreamzero] {i + 1}/{len(anchors)} ep{ep} t={anchor:.2f}s  {times[-1]:.1f}s  "
              f"(model returned {n} steps)", flush=True)
    np.savez(Path(args.out) / "pred_dreamzero.npz", gt=np.stack(gts), gt_joint=np.stack(gtj),
             time=np.asarray(times), pred=np.stack(preds), pred_joint=np.stack(preds_joint))


# ----------------------------------------------------------------------------- metrics
def errors(pred: np.ndarray, gt: np.ndarray) -> dict[str, np.ndarray]:
    """pred, gt: [N, 16, 62] -> per-anchor, per-step error arrays."""
    out = {}
    for side, sl in ARM.items():
        p, g = pred[..., sl], gt[..., sl]
        out[f"pos_{side}"] = np.linalg.norm(p[..., :3] - g[..., :3], axis=-1)
        rp = rotation_6d_to_matrix(p[..., 3:9]); rg = rotation_6d_to_matrix(g[..., 3:9])
        tr = np.einsum("...ji,...ji->...", rp, rg)  # trace(Rp^T Rg)
        out[f"rot_{side}"] = np.degrees(np.arccos(np.clip((tr - 1.0) / 2.0, -1.0, 1.0)))
    for side, sl in HAND.items():
        out[f"hand_{side}"] = np.sqrt(((pred[..., sl] - gt[..., sl]) ** 2).mean(-1))
    out["pos"] = 0.5 * (out["pos_left"] + out["pos_right"])
    out["rot"] = 0.5 * (out["rot_left"] + out["rot_right"])
    out["hand"] = 0.5 * (out["hand_left"] + out["hand_right"])
    return out


def summarize(args):
    out = Path(args.out)
    runs = {}
    r = np.load(out / "pred_raptor.npz")
    for key in r.files:
        if key.startswith("pred_"):
            runs["RAPTOR " + key[5:]] = (r[key], r["gt"], r["time"])
    d = np.load(out / "pred_dreamzero.npz")
    runs["DreamZero"] = (d["pred"], d["gt"], d["time"])
    assert np.allclose(r["gt"], d["gt"]), "both runs must use the same anchors"
    for spec in (args.extra or []):                      # NAME=path.npz[:key]
        name, _, rest = spec.partition("="); path, _, key = rest.partition(":")
        x = np.load(path); key = key or ("pred_cascade" if "pred_cascade" in x.files else "pred")
        assert np.allclose(x["gt"], r["gt"]), f"{path} was evaluated on different anchors"
        runs[name] = (x[key], x["gt"], x["time"])
    # baseline: hold the current measured pose for the whole chunk
    anchors_list = json.load(open(out / "anchors.json"))["anchors"]; root = json.load(open(out / "anchors.json"))["dataset_root"]
    hold = []; cache = {}
    for ep, t in anchors_list:
        if ep not in cache:
            b_ = tae.ChunkBuilder(root, ep); cache[ep] = (b_.timestamps, b_.state)
        ts, st = cache[ep]; hold.append(np.repeat(st[int(np.argmin(np.abs(ts - t)))][None], STEPS, 0))
    runs["hold current pose"] = (np.stack(hold).astype(np.float64), r["gt"], np.zeros(len(hold)))
    n_anchors, n_steps = r["gt"].shape[:2]
    anchors = json.load(open(out / "anchors.json"))["anchors"]
    n_eps = len({e for e, _ in anchors})

    table = {}
    per_anchor = {}
    for name, (pred, gt, tm) in runs.items():
        e = errors(pred, gt)
        row = {}
        for m in ("pos", "rot", "hand", "pos_left", "pos_right", "rot_left", "rot_right", "hand_left", "hand_right"):
            row[m] = float(e[m].mean())
            row[f"{m}@k0"] = float(e[m][:, 0].mean()); row[f"{m}@k8"] = float(e[m][:, 8].mean()); row[f"{m}@k15"] = float(e[m][:, 15].mean())
            row[f"{m}_curve"] = e[m].mean(0).tolist()
        row["sec_per_anchor"] = float(np.mean(tm))
        table[name] = row
        per_anchor[name] = {m: e[m].mean(1) for m in ("pos", "rot", "hand")}
    # extra: DreamZero's own joint-space arm error (not comparable to RAPTOR, for reference)
    dj = d["pred_joint"]; gj = d["gt_joint"]
    arm_idx = list(range(0, 7)) + list(range(29, 36))
    table["DreamZero"]["arm_joint_rmse_rad"] = float(np.sqrt(((dj[..., arm_idx] - gj[..., arm_idx]) ** 2).mean(-1)).mean())

    def paired(a, b):
        diff = per_anchor[a]["pos"] * 0 + 0  # placeholder for shape
        res = {}
        for m, unit in (("pos", "m"), ("rot", "deg"), ("hand", "rad")):
            diff = per_anchor[a][m] - per_anchor[b][m]
            res[m] = {"mean": float(diff.mean()), "sem": float(diff.std(ddof=1) / np.sqrt(len(diff))),
                      "win_rate_a": float((diff < 0).mean()), "unit": unit}
        return res
    comparisons = {}
    if "RAPTOR cascade" in runs:
        comparisons["RAPTOR cascade - DreamZero"] = paired("RAPTOR cascade", "DreamZero")
        comparisons["RAPTOR cascade - RAPTOR coarse_only"] = paired("RAPTOR cascade", "RAPTOR coarse_only")
    comparisons["RAPTOR coarse_only - DreamZero"] = paired("RAPTOR coarse_only", "DreamZero")
    for name in runs:
        if name not in ("RAPTOR cascade", "RAPTOR coarse_only", "DreamZero") and "RAPTOR cascade" in runs:
            comparisons[f"RAPTOR cascade - {name}"] = paired("RAPTOR cascade", name)

    lines = [f"# Open-loop comparison on {args.out}", "",
             f"{n_anchors} anchors from {n_eps} episodes of trex_small_force; 16-step (0.8 s, 20 Hz) chunk; one conditioning frame per camera.", "",
             "| model | EEF pos err [cm] | EEF rot err [deg] | hand RMSE [rad] | pos@0.05s | pos@0.45s | pos@0.8s | hand@0.05s | hand@0.8s | s/anchor |",
             "|---|---|---|---|---|---|---|---|---|---|"]
    for name, row in table.items():
        lines.append(f"| {name} | {100 * row['pos']:.2f} | {row['rot']:.2f} | {row['hand']:.4f} | {100 * row['pos@k0']:.2f} | {100 * row['pos@k8']:.2f} | {100 * row['pos@k15']:.2f} | {row['hand@k0']:.4f} | {row['hand@k15']:.4f} | {row['sec_per_anchor']:.1f} |")
    lines += ["", "Paired differences over anchors (negative = first model better):", ""]
    for name, res in comparisons.items():
        parts = [f"{m}: {v['mean'] * (100 if m == 'pos' else 1):+.3f} ± {v['sem'] * (100 if m == 'pos' else 1):.3f} {'cm' if m == 'pos' else v['unit']} (first better on {100 * v['win_rate_a']:.0f}% of anchors)" for m, v in res.items()]
        lines.append(f"- **{name}**: " + "; ".join(parts))
    lines += ["", f"DreamZero arm error in its own joint space: {table['DreamZero']['arm_joint_rmse_rad']:.4f} rad RMSE (reference only).",
              "", "Caveat: both checkpoints saw these episodes in training (trex_small is a subset of trex_full); this measures fit on the training distribution, not generalization."]
    text = "\n".join(lines)
    print(text)
    (out / "summary.md").write_text(text + "\n")
    json.dump({"table": table, "comparisons": comparisons, "n_anchors": n_anchors, "n_episodes": n_eps},
              open(out / "summary.json", "w"), indent=1)


# ----------------------------------------------------------------------------- main
def main():
    parser = argparse.ArgumentParser()
    sub = parser.add_subparsers(dest="cmd", required=True)
    run = sub.add_parser("run")
    run.add_argument("--model", choices=["raptor", "dreamzero"], required=True)
    run.add_argument("--checkpoint", required=True)
    run.add_argument("--dataset-root", default="data/trex_small_force")
    run.add_argument("--episodes", type=int, default=12)
    run.add_argument("--anchors", type=int, default=8)
    run.add_argument("--out", required=True)
    run.add_argument("--dry-run", action="store_true", help="only select anchors and check ground truth")
    summ = sub.add_parser("summarize")
    summ.add_argument("--out", required=True)
    summ.add_argument("--extra", action="append", help="NAME=path/pred.npz[:key] evaluated on the same anchors")
    args = parser.parse_args()
    if args.cmd == "summarize":
        return summarize(args)
    anchors = load_or_make_anchors(args)
    print(f"{len(anchors)} anchors from {len({e for e, _ in anchors})} episodes", flush=True)
    if args.dry_run:
        builder = tae.ChunkBuilder(args.dataset_root, anchors[0][0])
        gt_eef, gt_joint = ground_truth(builder, anchors[0][1])
        assert np.allclose(gt_eef[:, 9:31], gt_joint[:, 7:29]) and np.allclose(gt_eef[:, 40:62], gt_joint[:, 36:58])
        print("ground-truth hand dims consistent between eef62 and joint58; dry run OK")
        return
    (run_raptor if args.model == "raptor" else run_dreamzero)(args, anchors)


if __name__ == "__main__":
    main()