#!/usr/bin/env python3
"""Diagnose why Stack-Can tableHF_ep002.mp4 (a SUCCESS) scored track_max=0.0.

Re-runs the EXACT runtime TrackFlowScorer used by compare_robo_trackflow.py
on ep002 (and ep000 as a healthy reference), verbose, a few times — SAM3 is
non-deterministic so we want to see how often the can/plate detection misses.
"""
import sys, os, numpy as np

RL = "/home/hongyi/PWM/policy_rl/rl"
for p in [RL, f"{RL}/old/online"]:
    if p not in sys.path:
        sys.path.insert(0, p)
from rewards import TrackFlowScorer

M = "/home/hongyi/PWM/reward_model/3pointr_models"
SC = "/home/hongyi/PWM/external/DiffSynth-Studio/data/robot_new/Stack-Can"
TF_GPU = os.environ.get("TF_GPU", "0")


def ckpt():
    for c in ("stack_can_canplate_150pts_384_flow_backbone_final.pth",
              "stack_can_canplate_150pts_384_flow_backbone.pth"):
        if os.path.exists(f"{M}/{c}"):
            return f"{M}/{c}"
    raise FileNotFoundError


tf = TrackFlowScorer(
    pointr_ckpt=ckpt(),
    prompt="red can", prompt_fallbacks=["can", "orange can"],
    mark_prompt="white round plate",
    mark_prompt_fallbacks=["white circle", "white disc", "white object on the table"],
    subprocess_gpu=TF_GPU, verbose=True)

NRUNS = int(os.environ.get("NRUNS", "3"))
for ep in ("ep002", "ep000"):
    v = f"{SC}/tableHF_{ep}.mp4"
    print(f"\n########## {ep}  {v} ##########", flush=True)
    for i in range(NRUNS):
        tp = tf.score(v, mode="progress_gated", k_pixel=50.0)
        tp = np.asarray(tp)
        print(f">>> {ep} run{i}: track_max={tp.max():.4f}  nonzero={int((tp>0).sum())}/{tp.size}",
              flush=True)