from __future__ import annotations

from typing import Sequence

import cv2
import mlx.core as mx
import numpy as np

_IMAGENET_MEAN = (0.485, 0.456, 0.406)
_IMAGENET_STD = (0.229, 0.224, 0.225)


def resize(frames: mx.array, size: Sequence[int]) -> mx.array:
    """Resize a THWC clip to (height, width)."""
    if len(size) != 2:
        raise ValueError("size must be a sequence of (height, width)")

    height, width = int(size[0]), int(size[1])
    if height <= 0 or width <= 0:
        raise ValueError("resize dimensions must be positive")

    clip = np.asarray(frames)
    if clip.ndim != 4:
        raise ValueError("frames must have shape [T, H, W, C]")

    resized = [
        cv2.resize(frame, (width, height), interpolation=cv2.INTER_LINEAR)
        for frame in clip
    ]
    return mx.array(np.stack(resized, axis=0), dtype=frames.dtype)


def to_float(frames: mx.array, scale: bool = True) -> mx.array:
    """Cast frames to float32, optionally scaling from [0, 255] to [0, 1]."""
    result = frames.astype(mx.float32)
    if scale:
        result = result / 255.0
    return result


def normalize(
    frames: mx.array,
    mean: Sequence[float] = _IMAGENET_MEAN,
    std: Sequence[float] = _IMAGENET_STD,
) -> mx.array:
    """Normalize a THWC clip channel-wise."""
    if len(mean) != 3 or len(std) != 3:
        raise ValueError("mean and std must contain 3 channel values")

    clip = frames
    if clip.dtype != mx.float32:
        clip = to_float(clip, scale=True)

    scale_hint = float(mx.max(clip).item()) if clip.size > 0 else 0.0
    if scale_hint > 1.5:
        clip = clip / 255.0

    mean_arr = mx.array(mean, dtype=mx.float32).reshape(1, 1, 1, 3)
    std_arr = mx.array(std, dtype=mx.float32).reshape(1, 1, 1, 3)
    return (clip - mean_arr) / std_arr
