#!/usr/bin/env python3
"""Probe: which input protocol makes DreamZero (dreamzero/vla) produce sane actions on T-Rex data?

A  single current frame at an arbitrary anchor (what openloop_compare.py did)
B  sequential causal rollout from frame 0 with teacher-forced observations: chunk 0 = frame 0,
   chunk j = frames [a-23, a-16, a-8, a] with a = 24*j (the deployment path of inference_trex.py)
C  gt_video_action_pred: 25 real frames [a .. a+24] as (teacher-forced) video
For every protocol the 24 predicted joint actions are scored against the recorded targets
action[a .. a+23]: RMSE per group at k=0 and over the chunk, plus the copy-current-state baseline.
"""
import os, sys, time, json
from pathlib import Path
import numpy as np, torch

ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT / "scripts"))
import inference_trex as it
import torch.distributed as dist
from groot.vla.data.schema import EmbodimentTag
from groot.vla.model.n1_5.sim_policy import GrootSimPolicy

G = {"left_arm": slice(0, 7), "left_hand": slice(7, 29), "right_arm": slice(29, 36), "right_hand": slice(36, 58)}
KEYS = ("action.left_arm", "action.left_hand", "action.right_arm", "action.right_hand")
DIM = {"action.left_arm": 7, "action.left_hand": 22, "action.right_arm": 7, "action.right_hand": 22}


def act58(result):
    a = result.act
    return np.concatenate([np.asarray(a[k], dtype=np.float64).reshape(-1, DIM[k]) for k in KEYS], -1)


def score(pred, gt, state):
    out = {}
    n = min(len(pred), len(gt))
    for g, s in G.items():
        out[g] = (float(np.sqrt(((pred[0, s] - gt[0, s]) ** 2).mean())), float(np.sqrt(((pred[:n, s] - gt[:n, s]) ** 2).mean())),
                  float(np.sqrt(((state[s] - gt[:n, s]) ** 2).mean())))  # k0, chunk, copy-state baseline (chunk)
    return out


def fmt(sc):
    return "  ".join(f"{g}: k0 {v[0]:.3f} / chunk {v[1]:.3f} (hold {v[2]:.3f})" for g, v in sc.items())


def reset(policy):
    ah = policy.trained_model.action_head
    ah.current_start_frame = 0
    ah.language = None


def main():
    root = sys.argv[1] if len(sys.argv) > 1 else "data/trex_small_force"
    episodes = [int(x) for x in (sys.argv[2] if len(sys.argv) > 2 else "0,45").split(",")]
    chunks = [int(x) for x in (sys.argv[3] if len(sys.argv) > 3 else "1,2,3,4,5").split(",")]
    if not dist.is_initialized():
        os.environ.setdefault("MASTER_ADDR", "localhost"); os.environ.setdefault("MASTER_PORT", str(29600 + os.getpid() % 300))
        dist.init_process_group(backend="gloo", world_size=1, rank=0)
    policy = GrootSimPolicy(embodiment_tag=EmbodimentTag.TREX,
                            model_path="checkpoints/dreamzero_trex_wan22_lora/checkpoint-16000", device="cuda:0")
    H, W = it.get_expected_video_resolution(policy)
    ah = policy.trained_model.action_head
    if os.environ.get("DIT_ALL_STEPS") == "1" and hasattr(ah, "dit_step_mask"):
        ah.dit_step_mask = [True] * len(ah.dit_step_mask)
        print(f"DIT_ALL_STEPS: running all {len(ah.dit_step_mask)} denoising steps", flush=True)
    if os.environ.get("PROBE_ONLY_A") == "1":
        chunks_only_a = True
    else:
        chunks_only_a = False
    print(f"video {H}x{W}  num_frames={getattr(ah,'num_frames',None)} num_frame_per_block={getattr(ah,'num_frame_per_block',None)} action_horizon={getattr(ah,'action_horizon',None)}", flush=True)
    results = {}
    for ep in episodes:
        episode = it.TrexEpisode(root, ep)
        task = episode.get_task(0)
        actions = np.stack([np.array(episode.table.column("action")[i].as_py()) for i in range(episode.length)])
        states = np.stack([episode.get_state(i) for i in range(episode.length)])
        chunks_ep = [j for j in chunks if 24 * j + 24 <= episode.length]
        # ---- A: single frame at the chunk boundary
        for j in chunks_ep:
            a = 24 * j
            reset(policy)
            with torch.inference_mode():
                res, _ = policy.lazy_joint_forward_causal(it.Batch(obs=it.build_obs(episode, [a], task, H, W)))
            sc = score(act58(res), actions[a:a + 24], states[a])
            results[(ep, j, "A_single")] = sc; print(f"ep{ep} chunk{j} A_single  {fmt(sc)}", flush=True)
        # ---- B: sequential causal rollout from frame 0 (teacher-forced frames)
        if chunks_only_a:
            continue
        reset(policy)
        sched = it.build_chunk_schedule(episode.length, max(chunks_ep) + 1)
        for j, frames in enumerate(sched):
            with torch.inference_mode():
                res, _ = policy.lazy_joint_forward_causal(it.Batch(obs=it.build_obs(episode, frames, task, H, W)))
            a = frames[-1]
            if j in chunks_ep or j == 0:
                sc = score(act58(res), actions[a:a + 24], states[a])
                results[(ep, j, "B_sequential")] = sc; print(f"ep{ep} chunk{j} B_seq     frames={frames}  {fmt(sc)}  start_frame={ah.current_start_frame}", flush=True)
        # ---- C: teacher-forced future video
        for j in chunks_ep[:2]:
            a = 24 * j
            try:
                reset(policy)
                obs = it.build_obs(episode, list(range(a, a + 25)), task, H, W)
                with torch.inference_mode():
                    res, _ = policy.gt_video_action_pred(it.Batch(obs=obs))
                sc = score(act58(res), actions[a:a + 24], states[a])
                results[(ep, j, "C_gtvideo")] = sc; print(f"ep{ep} chunk{j} C_gtvideo {fmt(sc)}", flush=True)
            except Exception as e:
                print(f"ep{ep} chunk{j} C_gtvideo FAILED: {type(e).__name__}: {str(e)[:200]}", flush=True)
                break
    print("\n==== mean over probed chunks (k0 RMSE / chunk RMSE), rad")
    for proto in ("A_single", "B_sequential", "C_gtvideo"):
        rows = [v for (e, j, p), v in results.items() if p == proto and j > 0]
        if not rows: continue
        print(f"{proto:13s} " + "  ".join(f"{g}: {np.mean([r[g][0] for r in rows]):.3f}/{np.mean([r[g][1] for r in rows]):.3f}" for g in G))
    rows = [v for (e, j, p), v in results.items() if p == "A_single"]
    print(f"{'hold-state':13s} " + "  ".join(f"{g}: {np.mean([r[g][2] for r in rows]):.3f}" for g in G))
    json.dump({f"{e}|{j}|{p}": v for (e, j, p), v in results.items()}, open(os.environ.get("PROBE_OUT","ablation_runs/openloop_compare/dz_probe.json"), "w"), indent=1)


if __name__ == "__main__":
    main()