"""Generic hold-tail extender for ANY 3-PoinTr dataset (single- or multi-object).

Appends H timesteps that REPEAT the final cum_flow (pos is constant = x0, visibility
copied from the last frame), so the flow model learns to HOLD at the final pose instead
of relaxing the points back toward frame-0 near the end. Task-agnostic: operates on the
existing `no_actions/<task>/<NNNN>/timestep_k/{k_pos,k_cum_flow,k_visibility}.pth` tree,
copies labels.pth + metadata.json, and bumps T → T+H. Per-episode T is read from the
data (variable-length demos handled).

  python extend_3pointr_holdtail.py \
    --src /data0/.../3pointr_data_open_drawer_2d_wb/no_actions/open_drawer \
    --dst /data0/.../3pointr_data_open_drawer_2d_wb_hold30/no_actions/open_drawer --hold_tail 30
"""
import os, json, glob, shutil, argparse
from pathlib import Path
import torch


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--src", required=True, help="no_actions/<task> dir of the existing dataset")
    ap.add_argument("--dst", required=True, help="no_actions/<task> dir to write the extended dataset")
    ap.add_argument("--hold_tail", type=int, default=30)
    a = ap.parse_args()
    H = a.hold_tail
    os.makedirs(a.dst, exist_ok=True)
    eps = sorted(glob.glob(os.path.join(a.src, "[0-9][0-9][0-9][0-9]")))
    print(f"[extend] {len(eps)} episodes  hold_tail={H}  {a.src} -> {a.dst}", flush=True)
    minN = 10**9
    for ep in eps:
        name = Path(ep).name
        dst_ep = Path(a.dst) / name
        if dst_ep.exists():
            shutil.rmtree(dst_ep)
        shutil.copytree(ep, dst_ep)                       # original timesteps + labels + metadata
        T = len(glob.glob(os.path.join(dst_ep, "timestep_*")))
        x0  = torch.load(dst_ep / "timestep_0" / "0_pos.pth", weights_only=False)
        cf  = torch.load(dst_ep / f"timestep_{T-1}" / f"{T-1}_cum_flow.pth", weights_only=False)
        vis = torch.load(dst_ep / f"timestep_{T-1}" / f"{T-1}_visibility.pth", weights_only=False)
        for k in range(T, T + H):
            ts = dst_ep / f"timestep_{k}"
            ts.mkdir(parents=True, exist_ok=True)
            torch.save(x0,  ts / f"{k}_pos.pth")
            torch.save(cf,  ts / f"{k}_cum_flow.pth")     # held = final displacement
            torch.save(vis, ts / f"{k}_visibility.pth")
        mp = dst_ep / "metadata.json"
        meta = json.loads(mp.read_text()) if mp.exists() else {}
        meta["T_tracked"] = int(T); meta["hold_tail"] = int(H); meta["T"] = int(T + H)
        mp.write_text(json.dumps(meta, indent=2))
        minN = min(minN, int(x0.shape[0]))
    print(f"[extend] done: {len(eps)} eps, each +{H} held frames.  min_N = {minN}  "
          f"(set --subsample <= {minN})", flush=True)


if __name__ == "__main__":
    main()