"""Convert box+mark labeled tracks -> 3-PoinTr 2D dataset tree, with per-episode labels.

Input:  /data0/hongyi/reward_model_data/3pointr_data_push_box_2d_boxmark/<src>/tracks_labeled.npz
          {tracks (T,N,2) px@384, vis (T,N), labels (N,) 0=box/1=mark, n_box, n_mark}
Output: /data0/hongyi/reward_model_data/3pointr_data_push_box_boxmark/no_actions/push_box_boxmark/<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   <- box/mark, aligned with point order (box first, then mark)
          metadata.json  {point_dim:2, has_labels:true, n_box, n_mark}

Points/flow stay 2D (normalization untouched); the label is consumed by a label-embedding
in the model, NOT as a coordinate. Tracks are already 384x384 (videos_small) -> no rescale.
ignore_vis=True (trust CoTracker through-occlusion positions; box is occluded by hand/mark).
"""
import os, json, glob
from pathlib import Path
import numpy as np
import torch

SRC = "/data0/hongyi/reward_model_data/3pointr_data_push_box_2d_boxmark"
OUT = "/data0/hongyi/reward_model_data/3pointr_data_push_box_boxmark/no_actions/push_box_boxmark"


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", flush=True)
    written, minN = 0, 10**9
    summary = []
    for src_npz in eps:
        src = Path(src_npz).parent.name
        d = np.load(src_npz)
        tracks = d["tracks"].astype(np.float32)        # (T,N,2)
        labels = d["labels"].astype(np.int64)          # (N,)
        n_box, n_mark = int(d["n_box"]), int(d["n_mark"])
        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_box": n_box, "n_mark": n_mark, "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_box": n_box, "n_mark": n_mark})
        minN = min(minN, N)
        written += 1
    with open(os.path.join(os.path.dirname(OUT), "_boxmark_build_summary.json"), "w") as f:
        json.dump({"episodes": written, "min_N": minN, "items": summary}, f, indent=2)
    print(f"[conv] wrote {written} episodes -> {OUT}   min total points N = {minN}", flush=True)
    print(f"[conv] => set training --subsample <= {minN}", flush=True)


if __name__ == "__main__":
    main()