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(
    image: mx.array,
    size: Sequence[int],
    interpolation: str = "linear",
) -> mx.array:
    interpolations = {
        "linear": cv2.INTER_LINEAR,
        "nearest": cv2.INTER_NEAREST,
        "cubic": cv2.INTER_CUBIC,
        "area": cv2.INTER_AREA,
        "lanczos": cv2.INTER_LANCZOS4,
    }
    if interpolation not in interpolations:
        raise ValueError(f"interpolation must be one of {list(interpolations)}")
    if len(size) != 2:
        raise ValueError("size must be (height, width)")

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

    array = np.asarray(image)
    if array.ndim not in (3, 4):
        raise ValueError("image must have shape [H, W, C] or [B, H, W, C]")

    flag = interpolations[interpolation]
    if array.ndim == 4:
        out = np.stack(
            [cv2.resize(frame, (width, height), interpolation=flag) for frame in array],
            axis=0,
        )
    else:
        out = cv2.resize(array, (width, height), interpolation=flag)
        if out.ndim == 2:
            out = out[:, :, np.newaxis]

    return mx.array(out, dtype=image.dtype)


def center_crop(image: mx.array, size: Sequence[int]) -> mx.array:
    if len(size) != 2:
        raise ValueError("size must be (height, width)")

    crop_h, crop_w = int(size[0]), int(size[1])
    array = np.asarray(image)
    if array.ndim not in (3, 4):
        raise ValueError("image must have shape [H, W, C] or [B, H, W, C]")

    if array.ndim == 4:
        height, width = array.shape[1], array.shape[2]
    else:
        height, width = array.shape[0], array.shape[1]

    if crop_h > height or crop_w > width:
        raise ValueError(
            f"Crop size ({crop_h}, {crop_w}) exceeds image size ({height}, {width})"
        )

    top = (height - crop_h) // 2
    left = (width - crop_w) // 2

    if array.ndim == 4:
        out = array[:, top : top + crop_h, left : left + crop_w, :]
    else:
        out = array[top : top + crop_h, left : left + crop_w, :]

    return mx.array(out, dtype=image.dtype)


def pad(
    image: mx.array,
    padding: int | Sequence[int],
    fill: int | float = 0,
    mode: str = "constant",
) -> mx.array:
    if isinstance(padding, int):
        top = bottom = left = right = padding
    else:
        if len(padding) != 4:
            raise ValueError("padding must be int or (top, bottom, left, right)")
        top, bottom, left, right = [int(value) for value in padding]

    array = np.asarray(image)
    if array.ndim not in (3, 4):
        raise ValueError("image must have shape [H, W, C] or [B, H, W, C]")

    modes = {
        "constant": "constant",
        "reflect": "reflect",
        "edge": "edge",
    }
    if mode not in modes:
        raise ValueError(f"mode must be one of {list(modes)}")

    kwargs = {"constant_values": fill} if mode == "constant" else {}
    if array.ndim == 4:
        pad_width = ((0, 0), (top, bottom), (left, right), (0, 0))
    else:
        pad_width = ((top, bottom), (left, right), (0, 0))

    out = np.pad(array, pad_width, mode=modes[mode], **kwargs)
    return mx.array(out, dtype=image.dtype)


def horizontal_flip(image: mx.array) -> mx.array:
    array = np.asarray(image)
    if array.ndim not in (3, 4):
        raise ValueError("image must have shape [H, W, C] or [B, H, W, C]")
    axis = 2 if array.ndim == 4 else 1
    return mx.array(np.flip(array, axis=axis).copy(), dtype=image.dtype)


def to_float(image: mx.array, scale: bool = True) -> mx.array:
    result = image.astype(mx.float32)
    if scale:
        result = result / 255.0
    return result


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

    img = image.astype(mx.float32)
    if image.dtype == mx.uint8:
        img = img / 255.0

    if img.ndim == 4:
        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)
    elif img.ndim == 3:
        mean_arr = mx.array(mean, dtype=mx.float32).reshape(1, 1, 3)
        std_arr = mx.array(std, dtype=mx.float32).reshape(1, 1, 3)
    else:
        raise ValueError("image must have shape [H, W, C] or [B, H, W, C]")

    return (img - mean_arr) / std_arr


def to_chw(image: mx.array) -> mx.array:
    if image.ndim == 3:
        return mx.transpose(image, (2, 0, 1))
    if image.ndim == 4:
        return mx.transpose(image, (0, 3, 1, 2))
    raise ValueError("image must have shape [H, W, C] or [B, H, W, C]")


def to_hwc(image: mx.array) -> mx.array:
    if image.ndim == 3:
        return mx.transpose(image, (1, 2, 0))
    if image.ndim == 4:
        return mx.transpose(image, (0, 2, 3, 1))
    raise ValueError("image must have shape [C, H, W] or [B, C, H, W]")
