#!/usr/bin/env python3
"""Run Robometer progress scoring on the BenchVideo clips (no marks).

For each episode in a category's tasks.json:
  - load mp4 -> sample to <= max_frames at fps
  - frame_steps inference -> per-frame progress curve
  - save results_robometer.json + overview_robometer.png

Output: BenchVideo/eval_results/<Category>/<id>/
"""
import argparse
import json
import os
from pathlib import Path

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np

from benchmark_progress_mark_local import (
    RobometerLocalRunner,
    load_video_frames_with_indices,
)

BENCH_ROOT = Path("/home/vcj9002/jianshu/workspace/code_keliang/Current_Baseline/Benchmark_Videos/BenchVideo")
EVAL_ROOT = BENCH_ROOT / "eval_results"
MODEL_PATH = "/home/vcj9002/jianshu/workspace/code_keliang/Current_Baseline/Robometer/models/Robometer-4B"


def write_overview(out_path, frames, progress, title):
    fig, ax = plt.subplots(figsize=(11, 4.6))
    ax.plot(frames[: len(progress)], progress, marker="o", linewidth=1.5, markersize=2.8, color="#b9482f")
    ax.set_ylim(-0.05, 1.05)
    ax.set_xlim(0, max(5, int(max(frames) * 1.02)) if frames else 5)
    ax.set_xlabel("Original Frame"); ax.set_ylabel("Progress (normalized)")
    ax.set_title(title); ax.grid(True, linestyle="--", alpha=0.35)
    fig.tight_layout(); fig.savefig(out_path, dpi=180); plt.close(fig)


def run_category(runner, category, fps, max_frames, fallback):
    cat_dir = BENCH_ROOT / category
    data = json.loads((cat_dir / "tasks.json").read_text(encoding="utf-8"))
    out_root = EVAL_ROOT / category
    out_root.mkdir(parents=True, exist_ok=True)

    for ep in data["episodes"]:
        name = ep["episode"]; stem = name[:-4]
        task = (ep.get("tasks") or [fallback])[0] or fallback
        src = cat_dir / name
        ep_out = out_root / stem
        if (ep_out / "results_robometer.json").exists():
            print(f"[skip] {category}/{name} (results_robometer.json exists)")
            continue
        ep_out.mkdir(parents=True, exist_ok=True)

        print(f"\n[RUN] {category}/{name}  task={task!r}")
        frames, sampled_idx, total, native_fps = load_video_frames_with_indices(
            src, fps=fps, max_frames=max_frames, required_frames=[])
        progress, success = runner.compute_rewards_per_frame(
            frames, task, inference_mode="frame_steps",
            prefix_sample_frames=8, prefix_batch_size=1)
        progress = [float(x) for x in progress]

        results = {
            "category": category, "episode": name,
            "original_file": ep.get("original_file"), "task": task,
            "num_total_frames": total, "native_fps": native_fps, "sample_fps": fps,
            "num_sampled_frames": len(sampled_idx),
            "sampled_original_frame_indices": sampled_idx,
            "progress_curve": progress,
        }
        (ep_out / "results_robometer.json").write_text(
            json.dumps(results, ensure_ascii=False, indent=2), encoding="utf-8")
        write_overview(ep_out / "overview_robometer.png", sampled_idx, progress,
                       f"Robometer {category}/{stem}: {task[:50]}")
        rng = f"{min(progress):.2f}~{max(progress):.2f}" if progress else "empty"
        print(f"[OK] {ep_out}  (progress {rng})")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--category", default="all")
    ap.add_argument("--fps", type=float, default=2.0)
    ap.add_argument("--max-frames", type=int, default=128)
    ap.add_argument("--instruction-fallback", default="Complete the task shown in the video.")
    args = ap.parse_args()

    runner = RobometerLocalRunner(model_path=MODEL_PATH)
    cats = ["Count", "State", "Sequence"] if args.category == "all" else [args.category]
    for c in cats:
        run_category(runner, c, args.fps, args.max_frames, args.instruction_fallback)


if __name__ == "__main__":
    main()