from __future__ import annotations

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

import mlx_vision as mv


def _make_image(height: int = 32, width: int = 48, channels: int = 3) -> mx.array:
    return mx.array(
        np.random.randint(0, 255, (height, width, channels), dtype=np.uint8)
    )


def _make_batch(
    batch: int = 2, height: int = 32, width: int = 48, channels: int = 3
) -> mx.array:
    return mx.array(
        np.random.randint(0, 255, (batch, height, width, channels), dtype=np.uint8)
    )


def test_resize_single() -> None:
    image = _make_image()
    out = mv.resize(image, (16, 24))
    assert out.shape == (16, 24, 3)


def test_resize_batch() -> None:
    batch = _make_batch()
    out = mv.resize(batch, (16, 24))
    assert out.shape == (2, 16, 24, 3)


def test_center_crop() -> None:
    image = _make_image(32, 48)
    out = mv.center_crop(image, (16, 24))
    assert out.shape == (16, 24, 3)


def test_pad_and_layout_roundtrip() -> None:
    image = _make_image(12, 10)
    padded = mv.pad(image, (1, 2, 3, 4), fill=7)
    chw = mv.to_chw(padded)
    hwc = mv.to_hwc(chw)
    assert padded.shape == (15, 17, 3)
    assert hwc.shape == padded.shape


def test_horizontal_flip() -> None:
    arr = np.zeros((4, 6, 3), dtype=np.uint8)
    arr[:, 0, 0] = 255
    image = mx.array(arr)
    out = mv.horizontal_flip(image)
    assert int(np.asarray(out)[:, -1, 0].mean()) == 255


def test_normalize_shape() -> None:
    image = _make_image()
    out = mv.normalize(image)
    assert out.dtype == mx.float32
    assert out.shape == image.shape


def test_invalid_resize_mode() -> None:
    with pytest.raises(ValueError):
        mv.resize(_make_image(), (8, 8), interpolation="bad")
