"""Image resize/pad and state-selection helpers used to build the pi0.5 training data.

The functions below are copied VERBATIM from the benchmark repo's
task1_isaacsim/baseline/preprocess.py (the rest of that file is the offline
dataset-conversion tool and is not needed at inference time). They are what
`reference/policy_adapter_sim.py` imports as `from preprocess import ...`.

Use this exact code for live images: bilinear + antialias resizing gives
slightly different pixels in OpenCV, PIL and torchvision, and the policy was
trained on the output of THIS implementation.

    from preprocess import resize_with_pad
    img224 = resize_with_pad(rgb_uint8_hwc)   # (H, W, 3) uint8 RGB -> (224, 224, 3) uint8 RGB

Requires numpy and torch.
"""

from __future__ import annotations

import numpy as np
import torch

# Name-based state projections. "full" is a no-op (identity); "compact"
# selects each arm's measured joint positions + gripper joint value, in
# the same order as the franka_fr3_duo embodiment's ai_data_contract.
COMPACT_STATE_NAMES = (
    [f"franka_robot_left_measured_joint_states_left_fr3v2_joint{i}" for i in range(1, 8)]
    + ["franka_robot_left_joint_states_robotiq_85_left_knuckle_joint"]
    + [f"franka_robot_right_measured_joint_states_right_fr3v2_joint{i}" for i in range(1, 8)]
    + ["franka_robot_right_joint_states_robotiq_85_left_knuckle_joint"]
)


def _even(n: int) -> int:
    return n - (n % 2)


def fit_resolution(src: tuple[int, int], target: tuple[int, int]) -> tuple[int, int]:
    """Largest (w, h) that fits inside `target` while keeping `src`'s aspect ratio."""
    sw, sh = src
    tw, th = target
    scale = min(tw / sw, th / sh)
    return max(2, _even(round(sw * scale))), max(2, _even(round(sh * scale)))


def pad_to_size(img: np.ndarray, size: tuple[int, int]) -> np.ndarray:
    """Center `img` (H,W,3 uint8) on a black canvas of `size` = (W, H)."""
    w, h = size
    ih, iw = img.shape[:2]
    if (iw, ih) == (w, h):
        return img
    canvas = np.zeros((h, w, 3), dtype=img.dtype)
    y0, x0 = (h - ih) // 2, (w - iw) // 2
    canvas[y0 : y0 + ih, x0 : x0 + iw] = img
    return canvas


def _resize_to_uint8(img: torch.Tensor, width: int, height: int) -> np.ndarray:
    """(3,H,W) float32 in [0,1] - how LeRobotDataset decodes video - to (H,W,3) uint8."""
    if (int(img.shape[2]), int(img.shape[1])) != (width, height):
        img = torch.nn.functional.interpolate(
            img.unsqueeze(0), size=(height, width), mode="bilinear", align_corners=False, antialias=True
        )[0]
    arr = (img.clamp(0.0, 1.0) * 255.0).round().to(torch.uint8).permute(1, 2, 0).numpy()
    return np.ascontiguousarray(arr)


# --- Added for the handoff (not in the original file): the same pipeline for live uint8 frames. ---
def resize_with_pad(img: np.ndarray, size: tuple[int, int] = (224, 224)) -> np.ndarray:
    """(H, W, 3) uint8 RGB -> (size[1], size[0], 3) uint8 RGB, exactly as the training images were made."""
    h, w = img.shape[:2]
    inner_w, inner_h = fit_resolution((w, h), size)
    tensor = torch.from_numpy(np.ascontiguousarray(img)).permute(2, 0, 1).float() / 255.0
    return pad_to_size(_resize_to_uint8(tensor, inner_w, inner_h), size)