#!/usr/bin/env python3
"""Minimal Robometer wrapper that resizes frames to fit in memory."""
from __future__ import annotations

import argparse
import os
from pathlib import Path
from typing import Optional, Tuple

import numpy as np
import torch
import cv2

# Must import unsloth first
import unsloth  # noqa

from robometer.data.dataset_types import ProgressSample, Trajectory
from robometer.evals.eval_server import compute_batch_outputs
from robometer.evals.eval_viz_utils import extract_frames
from robometer.utils.save import load_model_from_hf
from robometer.utils.setup_utils import setup_batch_collator


def resize_frames(frames: np.ndarray, max_side: int = 256) -> np.ndarray:
    """Resize frames so max side = max_side."""
    T, H, W, C = frames.shape
    scale = max_side / max(H, W)
    if scale < 1:
        new_h = (int(H * scale) // 28) * 28  # must be multiple of 28 for qwen-vl
        new_w = (int(W * scale) // 28) * 28
        out = np.zeros((T, new_h, new_w, C), dtype=np.uint8)
        for i in range(T):
            out[i] = cv2.resize(frames[i], (new_w, new_h), interpolation=cv2.INTER_AREA)
        return out
    return frames


_CACHED_MODEL = None


def load_robometer_once(model_path: str = "robometer/Robometer-4B", device: str = "cuda"):
    """Load model once globally."""
    global _CACHED_MODEL
    if _CACHED_MODEL is None:
        device = torch.device(device)
        exp_config, tokenizer, processor, reward_model = load_model_from_hf(
            model_path=model_path, device=device,
        )
        reward_model.eval()
        batch_collator = setup_batch_collator(processor, tokenizer, exp_config, is_eval=True)
        _CACHED_MODEL = dict(
            model=reward_model, tokenizer=tokenizer, processor=processor,
            batch_collator=batch_collator, exp_config=exp_config, device=device,
        )
    return _CACHED_MODEL


def compute_rewards(
    video_frames: np.ndarray,
    task: str,
    max_side: int = 256,
    state: Optional[dict] = None,
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Compute per-frame progress and success probability.

    Args:
        video_frames: (T, H, W, C) uint8
        task: task instruction string
        max_side: resize frames so max dim = max_side

    Returns:
        (progress, success): both arrays of length T
    """
    if state is None:
        state = load_robometer_once()

    frames = resize_frames(video_frames, max_side=max_side)
    T = int(frames.shape[0])

    traj = Trajectory(
        frames=frames,
        frames_shape=tuple(frames.shape),
        task=task,
        id="0",
        metadata={"subsequence_length": T},
        video_embeddings=None,
    )
    progress_sample = ProgressSample(trajectory=traj, sample_type="progress")
    batch = state["batch_collator"]([progress_sample])

    progress_inputs = batch["progress_inputs"]
    for key, value in progress_inputs.items():
        if hasattr(value, "to"):
            progress_inputs[key] = value.to(state["device"])

    loss_config = getattr(state["exp_config"], "loss", None)
    is_discrete = (
        getattr(loss_config, "progress_loss_type", "l2").lower() == "discrete"
        if loss_config else False
    )
    num_bins = (
        getattr(loss_config, "progress_discrete_bins", None)
        or getattr(state["exp_config"].model, "progress_discrete_bins", 10)
    )

    with torch.inference_mode():
        results = compute_batch_outputs(
            state["model"], state["tokenizer"], progress_inputs,
            sample_type="progress", is_discrete_mode=is_discrete, num_bins=num_bins,
        )

    progress_pred = results.get("progress_pred", [])
    progress_array = (
        np.array(progress_pred[0], dtype=np.float32)
        if progress_pred and len(progress_pred) > 0
        else np.array([], dtype=np.float32)
    )
    outputs_success = results.get("outputs_success", {})
    success_probs = outputs_success.get("success_probs", []) if outputs_success else []
    success_array = (
        np.array(success_probs[0], dtype=np.float32)
        if success_probs and len(success_probs) > 0
        else np.array([], dtype=np.float32)
    )

    return progress_array, success_array


def score_video_file(
    video_path: str,
    task: str,
    fps: float = 3.0,
    max_frames: int = 8,
    max_side: int = 256,
    state: Optional[dict] = None,
) -> Tuple[np.ndarray, np.ndarray]:
    frames = extract_frames(video_path, fps=fps, max_frames=max_frames)
    if frames is None or frames.size == 0:
        raise RuntimeError(f"Could not extract frames from {video_path}")
    return compute_rewards(frames, task, max_side=max_side, state=state)


if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("--video", required=True)
    p.add_argument("--task", required=True)
    p.add_argument("--model-path", default="robometer/Robometer-4B")
    p.add_argument("--fps", type=float, default=3.0)
    p.add_argument("--max-frames", type=int, default=8)
    p.add_argument("--max-side", type=int, default=256)
    p.add_argument("--out", default="/tmp/robometer_out")
    args = p.parse_args()

    print(f"Loading model from {args.model_path}...")
    state = load_robometer_once(args.model_path)

    print(f"Scoring video: {args.video}")
    progress, success = score_video_file(
        args.video, args.task,
        fps=args.fps, max_frames=args.max_frames, max_side=args.max_side,
        state=state,
    )

    Path(args.out).mkdir(parents=True, exist_ok=True)
    np.save(os.path.join(args.out, "progress.npy"), progress)
    np.save(os.path.join(args.out, "success.npy"), success)

    print(f"Progress: shape={progress.shape}, values={progress.tolist()}")
    print(f"Success:  shape={success.shape}, values={success.tolist()}")
    print(f"Mean progress: {float(np.mean(progress)):.3f}")
    print(f"Mean success:  {float(np.mean(success)):.3f}")