#!/usr/bin/env python3
"""Probe NEW plate prompts on the can-only Stack-Can demos (plate previously missed).
For each (prompt), report how many of the probed demos it detects at ANY anchor."""
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

ROOT = "/home/hongyi/PWM/external/DiffSynth-Studio/data/robot_new/Stack-Can"
# subset of the 34 can-only demos (representative)
DEMOS = ["ep002","ep007","ep010","ep011","ep012","ep018","ep019","ep020",
         "ep031","ep042","ep061","ep073","ep080","ep091","ep092"]
PROMPTS = ["white round plate","round plate","white dish","dish","white disc",
           "white circle","ceramic plate","dinner plate","white bowl","saucer",
           "plate on the table","white object on the table","empty plate","white tray"]
ANCHORS = [0,16,32,48]
OUT="/tmp/stackcan_plate_v2"; 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)
    return 0 if m is None else int(np.asarray(m[0]).sum())

# cache squashed frames per demo
frames_cache={}
for d in DEMOS:
    fr=[cv2.resize(f,(384,384),interpolation=cv2.INTER_AREA) for f in load_video(f"{ROOT}/tableHF_{d}.mp4",None)]
    frames_cache[d]=fr

hits={p:[] for p in PROMPTS}   # demos detected (at any anchor)
for p in PROMPTS:
    for d in DEMOS:
        fr=frames_cache[d]
        best=max((ground(fr,a,p) for a in ANCHORS if a<len(fr)), default=0)
        if best>50:
            hits[p].append(d)
    print(f"{p!r:32s} -> {len(hits[p])}/{len(DEMOS)} demos  {hits[p]}", flush=True)

print("\n=== ranking (most demos recovered) ===", flush=True)
for p,h in sorted(hits.items(), key=lambda kv:-len(kv[1])):
    print(f"  {len(h):2d}/{len(DEMOS)}  {p!r}", flush=True)