"""Convert Stack-Can can+plate labeled tracks -> 3-PoinTr 2D dataset tree (label-conditioned).
Mirrors build_boxmark_3pointr.py.

Input:  /data0/hongyi/reward_model_data/3pointr_data_stack_can_2d_canplate/<src>/tracks_labeled.npz
          {tracks (T,N,2) px@384, vis (T,N), labels (N,) 0=can/1=plate, n_can, n_plate}
Output: /data0/hongyi/reward_model_data/3pointr_data_stack_can_canplate/no_actions/stack_can_canplate/<NNNN>/
          timestep_k/{k_pos.pth (N,2)=tracks[0], k_cum_flow.pth (N,2)=tracks[k]-tracks[0],
                      k_visibility.pth (N,)}
          labels.pth (N,) int64 (can first, then plate); metadata.json

EXCLUDES ep042 (false-positive plate: 23 pts, gap doesn't shrink). Tracks already 384x384.
"""
import os, json, glob
from pathlib import Path
import numpy as np
import torch

SRC = "/data0/hongyi/reward_model_data/3pointr_data_stack_can_2d_canplate"
OUT = "/data0/hongyi/reward_model_data/3pointr_data_stack_can_canplate/no_actions/stack_can_canplate"
EXCLUDE = {"tableHF_ep042"}     # wrong-object plate


def main():
    os.makedirs(OUT, exist_ok=True)
    eps = sorted(glob.glob(os.path.join(SRC, "tableHF_ep*", "tracks_labeled.npz")))
    print(f"[conv] {len(eps)} labeled-track episodes (excluding {sorted(EXCLUDE)})", flush=True)
    written, minN, skipped = 0, 10**9, []
    summary = []
    for src_npz in eps:
        src = Path(src_npz).parent.name
        if src in EXCLUDE:
            skipped.append(src); continue
        d = np.load(src_npz)
        tracks = d["tracks"].astype(np.float32)        # (T,N,2)
        labels = d["labels"].astype(np.int64)          # (N,) 0=can 1=plate
        n_can, n_plate = int(d["n_can"]), int(d["n_plate"])
        if n_can == 0 or n_plate == 0:                  # safety: only both-object demos
            skipped.append(src); continue
        T, N, _ = tracks.shape
        x0 = tracks[0]                                  # (N,2)
        out_dir = Path(OUT) / f"{written:04d}"
        for k in range(T):
            ts = out_dir / f"timestep_{k}"
            ts.mkdir(parents=True, exist_ok=True)
            torch.save(torch.from_numpy(x0.copy()), ts / f"{k}_pos.pth")
            torch.save(torch.from_numpy((tracks[k] - x0).copy()), ts / f"{k}_cum_flow.pth")
            torch.save(torch.from_numpy(np.ones(N, dtype=np.bool_)), ts / f"{k}_visibility.pth")
        torch.save(torch.from_numpy(labels.copy()), out_dir / "labels.pth")
        meta = {"T": int(T), "N": int(N), "point_dim": 2, "has_labels": True,
                "n_can": n_can, "n_plate": n_plate, "src": src, "frame_size": [384, 384],
                "intrinsics": {"camera": {"fx": 1.0, "fy": 1.0, "cx": 0.0, "cy": 0.0,
                                          "width": 384, "height": 384}}}
        with open(out_dir / "metadata.json", "w") as f:
            json.dump(meta, f, indent=2)
        summary.append({"idx": written, "src": src, "T": T, "N": N, "n_can": n_can, "n_plate": n_plate})
        minN = min(minN, N)
        written += 1
    with open(os.path.join(os.path.dirname(OUT), "_canplate_build_summary.json"), "w") as f:
        json.dump({"episodes": written, "min_N": minN, "skipped": skipped, "items": summary}, f, indent=2)
    print(f"[conv] wrote {written} episodes -> {OUT}   min total points N = {minN}", flush=True)
    print(f"[conv] skipped: {skipped}", flush=True)
    print(f"[conv] => set training --subsample <= {minN}", flush=True)


if __name__ == "__main__":
    main()