"""Trace Anything preprocessing: scripts/infer.py ``_load_images`` (ByteDance-Seed/TraceAnything @ 54677b5e) behind the
transformers image-processor interface: portrait frames rotated to landscape, longer side resized to 512 (PIL bilinear),
every frame cropped to the first frame's size floored to multiples of 16, normalised to [-1, 1]; frame i of N gets the
time step i / (N - 1). infer.py reads files with cv2 (BGR -> RGB); a path is read the same way here, a PIL image / RGB
uint8 array is used as RGB. (infer.py also keeps every (N // 39)-th frame of scenes with more than 40 frames; that
subsampling is left to the caller.)"""

import cv2
import numpy as np
import torch
import torchvision.transforms as tvf
from PIL import Image
from transformers.feature_extraction_utils import BatchFeature
from transformers.image_processing_utils import BaseImageProcessor


def _resize_long_side(pil: Image.Image, long: int = 512) -> Image.Image:  # infer.py
    w, h = pil.size
    if w >= h:
        return pil.resize((long, int(h * long / w)), Image.BILINEAR)
    else:
        return pil.resize((int(w * long / h), long), Image.BILINEAR)


class TraceAnythingImageProcessor(BaseImageProcessor):
    model_input_names = ["pixel_values", "time_steps"]

    def __init__(self, image_size: int = 512, **kwargs):
        super().__init__(**kwargs)
        self.image_size = image_size

    def preprocess(self, images, image_size: int | None = None, return_tensors=None, **kwargs) -> BatchFeature:
        """images: the frames of one video, in order (paths, PIL images or RGB uint8 arrays). pixel_values
        (1, N, 3, H, W) in [-1, 1]; time_steps (1, N) float64 (i / (N - 1), 0 for one frame)."""
        long = image_size or self.image_size
        # body of infer.py _load_images over the given frames
        tfm = tvf.Compose([tvf.ToTensor(), tvf.Normalize((0.5,)*3, (0.5,)*3)])
        frames, times, target = [], [], None
        for i, image in enumerate(images):
            if isinstance(image, str):
                arr = cv2.imread(image)
                if arr is None:
                    raise FileNotFoundError(f"Failed to read {image}")
                pil = Image.fromarray(cv2.cvtColor(arr, cv2.COLOR_BGR2RGB))
            else:
                pil = image.convert("RGB") if isinstance(image, Image.Image) else Image.fromarray(np.asarray(image))

            W0, H0 = pil.size
            if H0 > W0:  # portrait -> landscape
                pil = pil.transpose(Image.Transpose.ROTATE_90)

            pil = _resize_long_side(pil, long)
            if target is None:
                H, W = pil.size[1], pil.size[0]
                target = (H - H % 16, W - W % 16)
            Ht, Wt = target
            pil = pil.crop((0, 0, Wt, Ht))

            frames.append(tfm(pil))
            times.append(i / (len(images) - 1) if len(images) > 1 else 0.0)
        data = {"pixel_values": torch.stack(frames)[None], "time_steps": torch.tensor([times], dtype=torch.float64)}
        return BatchFeature(data, tensor_type=return_tensors)


__all__ = ["TraceAnythingImageProcessor"]