#!/usr/bin/env python3
"""Track TWO objects on the Stack-Can teleop demos, saving combined point tracks
WITH per-point object labels (0=can, 1=plate).  Mirrors build_boxmark_tracks.py
(the box+mark push-box pipeline).

NOTE on prompts: the user asked for "Green cup" + "White plate on the table".
A SAM3 probe showed the manipulated object is a red/orange CAN, not a cup
("Green cup"/"cup"/"mug" -> nothing), and that grounding is FLAKY across demos
and across anchor frames.  The robust recipe (from _probe_stackcan_robust.py over
ep000..ep004):
  - CAN  : "red can" (a0 on most demos; identical to "can"/"orange can")
  - PLATE: "white round plate" grounds on ALL anchors/demos, far better than
           "White plate on the table" (which is per-anchor flaky).
We try a prompt-variant list x several anchor frames and take the first that
grounds, then CoTracker tracks the FULL video from that frame (CoTracker3 is
bidirectional, so tracks cover all T frames regardless of anchor).

Source: robot_new/Stack-Can/tableHF_ep{NNN}.mp4 — the 100 success teleop demos,
native 1280x720.  Each frame is anisotropically SQUASHed to 384x384 before
tracking, matching the WM / 3-PoinTr domain (the "tableHF 1280x720 squash" gotcha).

Per episode -> <OUT>/<src>/tracks_labeled.npz with:
  tracks (T,N,2) float32, vis (T,N) bool, labels (N,) int [0=can,1=plate],
  n_can, n_plate, episode_index, src, H, W
Also writes a combined overlay mp4 for the first --n_overlay episodes (can=green,
plate=blue).  Loads SAM3 + CoTracker ONCE and loops.
"""
import sys, os, json, glob, argparse
from pathlib import Path
import numpy as np
import cv2
import imageio

sys.path.insert(0, "/home/hongyi/PWM/external")
import sam3.model_builder as _mb
from sam3.model_builder import build_sam3_video_predictor
from cotracker.predictor import CoTrackerPredictor
from track_objects_in_video import load_video, sam3_propagate, run_cotracker, COTRACKER_CKPT

DEMO_GLOB = "/home/hongyi/PWM/external/DiffSynth-Studio/data/robot_new/Stack-Can/tableHF_ep[0-9][0-9][0-9].mp4"
OUT_ROOT_DEFAULT = "/data0/hongyi/reward_model_data/3pointr_data_stack_can_2d_canplate"
# (tag/label-name, ordered SAM3 prompt variants, overlay BGR color).  label_id = list index.
PROMPTS = [
    ("can",   ["red can", "can", "orange can"],                            (0, 255, 0)),
    # Plate is hard: "white round plate" misses ~1/3 of demos. Probe (v2) on the
    # failed set found shape-specific "white circle"/"white disc" recover most, and
    # "white object on the table" recovers all 15/15 (generic but table-constrained).
    # Order: specific-shape first (least false-positive risk), generic catch-all last.
    ("plate", ["white round plate", "white circle", "white disc", "round plate",
               "White plate on the table", "white plate", "white object on the table"], (255, 0, 0)),
]
ANCHORS = [0, 16, 8, 32]   # grounding-clip start frames, tried in order (0 first = full coverage)
GROUND_N = 16              # grounding-clip length (the probe used 16)
SQUASH = 384               # square target the WM / 3-PoinTr operate in


def ground_and_track(predictor, cot, full_frames, prompt_variants, stride, retries=2):
    """Ground SAM3 on a short clip (prompt-variant x anchor sweep), then CoTracker
    over the FULL video.  Returns (tracks (T,N,2), vis (T,N), used_prompt, used_frame)
    or None if nothing grounds."""
    n = len(full_frames)
    tmp = f"/tmp/stackcan_ground_{os.getpid()}.mp4"
    for anchor in ANCHORS:
        if anchor >= n:
            continue
        clip = full_frames[anchor:anchor + GROUND_N]
        if len(clip) < 4:
            continue
        imageio.mimwrite(tmp, clip, fps=15, macro_block_size=1)
        for prompt in prompt_variants:
            masks = None
            for _ in range(retries):
                m, _n = sam3_propagate(predictor, tmp, prompt)
                if m is not None:
                    masks = m
                    break
            if masks is None:
                continue
            masks_arr = np.stack([masks[i] for i in range(len(clip))], axis=0)
            nz = [i for i in range(len(clip)) if masks_arr[i].any()]
            if not nz:
                continue
            ci = nz[0]
            tr, vis = run_cotracker(cot, full_frames, masks_arr[ci], stride,
                                    init_frame_idx=anchor + ci)
            if tr is not None:
                return tr, vis, prompt, anchor + ci
    return None


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out_root", default=OUT_ROOT_DEFAULT)
    ap.add_argument("--stride", type=int, default=3)         # box+mark used 3
    ap.add_argument("--limit", type=int, default=0)          # >0 = first N demos
    ap.add_argument("--n_overlay", type=int, default=5)
    ap.add_argument("--demos", default="")                   # comma ep ids (e.g. ep002,ep007) = re-track only these
    args = ap.parse_args()

    os.makedirs(args.out_root, exist_ok=True)
    demos = sorted(glob.glob(DEMO_GLOB))
    if args.demos:
        want = {e.strip() for e in args.demos.split(",") if e.strip()}
        demos = [d for d in demos if Path(d).stem.replace("tableHF_", "") in want]
    if args.limit > 0:
        demos = demos[:args.limit]
    print(f"[build] {len(demos)} Stack-Can demos -> {args.out_root}  (stride={args.stride})", flush=True)

    bpe = str(Path(_mb.__file__).parent / "assets" / "bpe_simple_vocab_16e6.txt.gz")
    print("[build] loading SAM3 ...", flush=True)
    predictor = build_sam3_video_predictor(bpe_path=bpe, gpus_to_use=[0],
                                           apply_temporal_disambiguation=False)
    print("[build] loading CoTracker3 ...", flush=True)
    cot = CoTrackerPredictor(checkpoint=COTRACKER_CKPT).cuda()

    manifest = []
    for k, vid in enumerate(demos):
        src = Path(vid).stem                                  # tableHF_epNNN
        ep_idx = int(src.split("ep")[-1])
        outdir = os.path.join(args.out_root, src)
        os.makedirs(outdir, exist_ok=True)
        # native frames -> squash to SQUASH x SQUASH (anisotropic, WM-domain match)
        frames = [cv2.resize(f, (SQUASH, SQUASH), interpolation=cv2.INTER_AREA)
                  for f in load_video(vid, None)]
        H, W = frames[0].shape[:2]

        res = {tag: ground_and_track(predictor, cot, frames, variants, args.stride)
               for tag, variants, _ in PROMPTS}
        if all(v is None for v in res.values()):
            print(f"[warn] {src}: NO objects detected — skipping", flush=True)
            manifest.append({"src": src, "episode_index": ep_idx, "status": "no_objects"})
            continue

        # Build labeled combined tracks. Order = PROMPTS order (can=0, plate=1).
        track_list, vis_list, label_list, counts, used = [], [], [], {}, {}
        for label_id, (tag, _, _) in enumerate(PROMPTS):
            r = res[tag]
            if r is None:
                counts[tag] = 0
                print(f"[warn] {src}: {tag.upper()} not detected", flush=True)
                continue
            tr, vs, used_prompt, used_frame = r
            n = tr.shape[1]
            track_list.append(tr.astype(np.float32))
            vis_list.append(vs)
            label_list.append(np.full(n, label_id, np.int64))
            counts[tag] = n
            used[tag] = f"{used_prompt}@f{used_frame}"
        T = track_list[0].shape[0]
        tracks = np.concatenate(track_list, axis=1).astype(np.float32)   # (T, N, 2)
        vis = np.concatenate(vis_list, axis=1)
        labels = np.concatenate(label_list)                             # (N,) 0=can 1=plate

        np.savez(os.path.join(outdir, "tracks_labeled.npz"),
                 tracks=tracks, vis=vis, labels=labels,
                 n_can=counts.get("can", 0), n_plate=counts.get("plate", 0),
                 episode_index=ep_idx, src=src, H=H, W=W)

        # can<->plate centroid gap (sanity: should shrink as the can lands on the plate)
        d0 = dl = -1.0
        if res["can"] is not None and res["plate"] is not None:
            can, plate = res["can"][0], res["plate"][0]
            d0 = float(np.linalg.norm(can[0].mean(0) - plate[0].mean(0)))
            dl = float(np.linalg.norm(can[-1].mean(0) - plate[-1].mean(0)))
        manifest.append({"src": src, "episode_index": ep_idx, "status": "ok",
                         "n_can": counts.get("can", 0), "n_plate": counts.get("plate", 0),
                         "T": int(T), "can_plate_d0": round(d0, 1), "can_plate_dlast": round(dl, 1),
                         "used": used})
        print(f"[{k+1}/{len(demos)}] {src}: can={counts.get('can',0)} "
              f"plate={counts.get('plate',0)} T={T} gap {d0:.0f}->{dl:.0f}px "
              f"({used.get('can','-')} | {used.get('plate','-')})", flush=True)

        if k < args.n_overlay:
            ov = os.path.join(outdir, "combined_overlay.mp4")
            vw = imageio.get_writer(ov, fps=15, macro_block_size=1)
            for t in range(T):
                im = frames[t].copy()
                for label_id, (tag, _, color) in enumerate(PROMPTS):
                    r = res[tag]
                    if r is None:
                        continue
                    tr = r[0]
                    for pi in range(tr.shape[1]):
                        x, y = tr[t, pi]
                        cv2.circle(im, (int(x), int(y)), 1, color, -1)
                vw.append_data(im)
            vw.close()

    with open(os.path.join(args.out_root, "manifest.json"), "w") as f:
        json.dump(manifest, f, indent=2)
    ok = sum(1 for m in manifest if m.get("status") == "ok")
    both = sum(1 for m in manifest if m.get("n_can", 0) > 0 and m.get("n_plate", 0) > 0)
    print(f"[build] DONE: {ok} ok ({both} with BOTH objects), "
          f"{len(manifest)-ok} failed. -> {args.out_root}", flush=True)


if __name__ == "__main__":
    main()