#!/usr/bin/env python3
"""For the flaky demos, try each prompt-variant grounded at several anchor frames,
report which (prompt, anchor) succeed.  Goal: a robust grounding recipe."""
import sys, os
from pathlib import Path
import numpy as np, cv2, 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 track_objects_in_video import load_video, sam3_propagate

DEMOS = ["tableHF_ep000", "tableHF_ep001", "tableHF_ep002", "tableHF_ep003", "tableHF_ep004"]
ROOT = "/home/hongyi/PWM/external/DiffSynth-Studio/data/robot_new/Stack-Can"
CAN_P = ["red can", "can", "orange can"]
PLATE_P = ["White plate on the table", "white plate", "plate", "white round plate"]
ANCHORS = [0, 8, 16]           # start frame of the 16-frame grounding clip
OUT = "/tmp/stackcan_robust"; os.makedirs(OUT, exist_ok=True)

bpe = str(Path(_mb.__file__).parent / "assets" / "bpe_simple_vocab_16e6.txt.gz")
pred = build_sam3_video_predictor(bpe_path=bpe, gpus_to_use=[0], apply_temporal_disambiguation=False)

def ground(frames, anchor, prompt):
    clip = frames[anchor:anchor+16]
    if len(clip) < 4:
        return -1
    tmp = os.path.join(OUT, "clip.mp4"); imageio.mimwrite(tmp, clip, fps=15, macro_block_size=1)
    m, n = sam3_propagate(pred, tmp, prompt)
    if m is None:
        return 0
    return int(np.asarray(m[0]).sum())

for d in DEMOS:
    vid = f"{ROOT}/{d}.mp4"
    frames = [cv2.resize(f, (384, 384), interpolation=cv2.INTER_AREA) for f in load_video(vid, None)]
    print(f"\n=== {d} ({len(frames)} frames) ===", flush=True)
    for name, plist in [("CAN", CAN_P), ("PLATE", PLATE_P)]:
        row = []
        for p in plist:
            cells = []
            for a in ANCHORS:
                px = ground(frames, a, p)
                cells.append(f"a{a}={px}")
            row.append(f"{p!r}: " + " ".join(cells))
        print(f"  {name}:", flush=True)
        for r in row:
            print("    ", r, flush=True)
print("\ndone", flush=True)