#!/usr/bin/env python3
"""
Run RBM inference locally: load a checkpoint from HuggingFace and compute per-frame progress
and success for a video (or .npy/.npz frames) and task instruction. Writes rewards .npy,
success-probs .npy, and a progress/success plot. Requires the robometer package.

Example:
  python scripts/example_inference_local.py \\
    --model-path aliangdw/qwen4b_pref_prog_succ_8_frames_all_part2 \\
    --video /path/to/video.mp4 \\
    --task "Pick up the red block and place it in the bin"
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np

from benchmark_progress_mark_local import (
    DEFAULT_INFERENCE_MODE,
    DEFAULT_MAX_FRAMES,
    DEFAULT_MIN_FRAMES,
    DEFAULT_PREFIX_BATCH_SIZE,
    DEFAULT_PREFIX_SAMPLE_FRAMES,
    RobometerLocalRunner,
    build_frame_retry_schedule,
    is_cuda_oom_error,
    load_all_video_frames,
    sample_video_frames_with_indices,
)
from robometer.evals.eval_viz_utils import create_combined_progress_success_plot


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Run RBM inference locally: load model from HuggingFace and compute per-frame progress and success.",
        epilog="Outputs: <out>.npy (rewards), <out>_success_probs.npy, <out>_progress_success.png",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument("--model-path", default="../../models/Robometer-4B", help="HuggingFace model id or local checkpoint path")
    parser.add_argument("--video", default="example_videos/soar_put_green_stick_in_brown_bowl.mp4", help="Video path/URL or .npy/.npz with frames (T,H,W,C)")
    parser.add_argument("--task", default="Put green stick in brown bowl", help="Task instruction for the trajectory")
    parser.add_argument("--fps", type=float, default=1.0, help="FPS when sampling from video (default: 1.0)")
    parser.add_argument("--max-frames", type=int, default=DEFAULT_MAX_FRAMES, help="Max frames to extract from video (default: 128)")
    parser.add_argument("--min-frames", type=int, default=DEFAULT_MIN_FRAMES, help="Minimum retry frame budget after OOM (default: 32)")
    parser.add_argument(
        "--inference-mode",
        choices=["frame_steps", "whole"],
        default=DEFAULT_INFERENCE_MODE,
        help="frame_steps matches benchmark behavior; whole does a single forward pass on the full sampled trajectory",
    )
    parser.add_argument("--prefix-sample-frames", type=int, default=DEFAULT_PREFIX_SAMPLE_FRAMES, help="Frames per prefix in frame_steps mode (default: 4)")
    parser.add_argument("--prefix-batch-size", type=int, default=DEFAULT_PREFIX_BATCH_SIZE, help="Batch size for prefix inference in frame_steps mode (default: 1)")
    parser.add_argument(
        "--adaptive-max-frames",
        dest="adaptive_max_frames",
        action="store_true",
        help="On CUDA OOM in whole mode, retry with a smaller frame budget",
    )
    parser.add_argument(
        "--no-adaptive-max-frames",
        dest="adaptive_max_frames",
        action="store_false",
        help="Disable frame-budget retry on CUDA OOM",
    )
    parser.add_argument(
        "--success-threshold",
        type=float,
        default=0.5,
        help="Threshold for binary success in plot (default: 0.5)",
    )
    parser.add_argument("--out", default=None, help="Output path for rewards .npy (default: <video_stem>_rewards.npy)")
    parser.set_defaults(adaptive_max_frames=True)
    args = parser.parse_args()

    video_path = Path(args.video)
    
    # Create output directory: scripts/outputs/{video_name}/
    if args.out is not None:
        out_path = Path(args.out)
    else:
        output_dir = Path(__file__).parent / "outputs" / video_path.stem
        output_dir.mkdir(parents=True, exist_ok=True)
        out_path = output_dir / f"{video_path.stem}_rewards.npy"

    runner = RobometerLocalRunner(model_path=args.model_path)

    all_frames, native_fps = load_all_video_frames(video_path)
    retry_schedule = (
        build_frame_retry_schedule(args.max_frames, args.min_frames, bool(args.adaptive_max_frames))
        if args.inference_mode == "whole"
        else [int(args.max_frames)]
    )

    frames = None
    rewards = None
    success_probs = None
    used_max_frames = retry_schedule[0]

    for attempt_idx, frame_budget in enumerate(retry_schedule, start=1):
        frames, _ = sample_video_frames_with_indices(
            all_frames,
            native_fps=native_fps,
            fps=float(args.fps),
            max_frames=int(frame_budget),
            required_frames=[],
        )
        print(
            f"Loaded {len(all_frames)} total frames; sampled {len(frames)} frames at fps={float(args.fps):g} "
            f"(max_frames={int(frame_budget)}, try {attempt_idx}/{len(retry_schedule)})"
        )
        try:
            rewards, success_probs = runner.compute_rewards_per_frame(
                video_frames=frames,
                task=args.task,
                inference_mode=args.inference_mode,
                prefix_sample_frames=int(args.prefix_sample_frames),
                prefix_batch_size=int(args.prefix_batch_size),
            )
            used_max_frames = int(frame_budget)
            break
        except RuntimeError as exc:
            if args.inference_mode != "whole" or not is_cuda_oom_error(exc) or attempt_idx == len(retry_schedule):
                raise
            next_budget = retry_schedule[attempt_idx]
            print(
                f"[OOM] whole inference hit CUDA OOM at max_frames={int(frame_budget)}; "
                f"retrying with max_frames={int(next_budget)}"
            )
            runner.reload_model()

    if rewards is None or success_probs is None or frames is None:
        raise RuntimeError("Robometer inference did not produce outputs.")

    # Save results (directory already created)
    np.save(str(out_path), rewards)
    success_path = out_path.with_name(out_path.stem + "_success_probs.npy")
    np.save(str(success_path), success_probs)

    show_success = success_probs.size > 0 and success_probs.size == rewards.size
    success_binary = (success_probs > float(args.success_threshold)).astype(np.int32) if show_success else None
    fig = create_combined_progress_success_plot(
        progress_pred=rewards,
        num_frames=int(frames.shape[0]),
        success_binary=success_binary,
        success_probs=success_probs if show_success else None,
        success_labels=None,
        title=f"Progress/Success — {video_path.name}",
    )
    plot_path = out_path.with_name(out_path.stem + "_progress_success.png")
    fig.savefig(str(plot_path), dpi=200)
    plt.close(fig)

    summary = {
        "video": str(video_path),
        "num_frames": int(frames.shape[0]),
        "inference_mode": args.inference_mode,
        "max_frames_used": int(used_max_frames),
        "model_path": args.model_path,
        "out_rewards": str(out_path),
        "out_success_probs": str(success_path),
        "out_plot": str(plot_path),
        "reward_min": float(np.min(rewards)) if rewards.size else None,
        "reward_max": float(np.max(rewards)) if rewards.size else None,
        "reward_mean": float(np.mean(rewards)) if rewards.size else None,
    }
    print(json.dumps(summary, indent=2))


if __name__ == "__main__":
    main()