"""Frame preprocessing identical to ``dataset.LiberoHDF5Dataset._to_tensor``.

Verbatim port of the reference transform. LIBERO stores every camera (including the wrist
camera) rotated by 180 degrees; the reference dataset rights them with ``x[:, ::-1, ::-1]``
applied to the ``[n, H, W, C]`` array *before* any other op.

The example PNGs shipped with this Space were exported from the demonstrations with that
rotation already applied, so the app passes ``rotate=False`` for them; set it to ``True`` (the
reference default) for raw, as-stored LIBERO frames.
"""

from __future__ import annotations

from typing import Sequence, Tuple

import numpy as np
import torch
import torch.nn.functional as F

IMAGE_SIZE: Tuple[int, int] = (128, 128)
ROTATE_180: bool = True


def to_tensor(raw: np.ndarray, rotate: bool = ROTATE_180) -> torch.Tensor:
    """``[H, W, C]`` or ``[n, H, W, C]`` -> ``[C, H', W']`` / ``[n, C, H', W']`` float32 in [0, 1]."""
    x = np.asarray(raw)
    single = x.ndim == 3
    if single:
        x = x[None]
    # the rotate must happen *before* every other op, exactly like the dataset does
    if rotate:
        x = x[:, ::-1, ::-1]
    out = torch.from_numpy(np.ascontiguousarray(x)).permute(0, 3, 1, 2).to(torch.float32)
    if x.dtype == np.uint8:
        out = out / 255.0
    if IMAGE_SIZE is not None and tuple(out.shape[-2:]) != IMAGE_SIZE:
        out = F.interpolate(out, size=IMAGE_SIZE, mode="bilinear", align_corners=False, antialias=True)
    return out[0] if single else out


def to_views(frames: Sequence[np.ndarray], rotate: bool = ROTATE_180) -> torch.Tensor:
    """One frame per camera -> ``[V, C, H, W]`` in the order given (the checkpoint's ``camera_keys``)."""
    return torch.stack([to_tensor(f, rotate=rotate) for f in frames], dim=0)