import numpy as np

from volumentations import (
    Compose,
    ElasticTransform,
    Flip,
    RandomCropFromBorders,
    RandomRotate90,
    RandomScale,
    Resize,
    Rotate,
)
from volumentations.augmentations import functional as F


def test_resize_skimage_backend_preserves_float32():
    image = np.random.rand(16, 8, 8, 5).astype(np.float32)

    resized = F.resize(image, new_shape=(32, 16, 16), interpolation=1, resize_type=0)

    assert resized.shape == (32, 16, 16, 5)
    assert resized.dtype == np.float32
    assert np.isfinite(resized).all()


def test_resize_nearest_neighbor_preserves_bool_masks():
    mask = np.zeros((8, 6, 4, 1), dtype=np.bool_)
    mask[2:6, 1:5, 1:3, 0] = True

    resized = F.resize(mask, new_shape=(16, 12, 10), interpolation=0, resize_type=0)

    assert resized.shape == (16, 12, 10, 1)
    assert resized.dtype == np.bool_
    assert np.isfinite(resized.astype(np.float32)).all()
    assert np.all(np.isin(np.unique(resized), [False, True]))


def test_rescale_scales_all_three_spatial_axes():
    image = np.random.rand(10, 8, 6).astype(np.float32)

    scaled = F.rescale(image, scale=1.5, interpolation=1)

    assert scaled.shape == (15, 12, 9)
    assert scaled.dtype == np.float32
    assert np.isfinite(scaled).all()


def test_numpy2_smoke_pipeline_preserves_float32_and_mask_values():
    aug = Compose(
        [
            Rotate(p=1.0),
            RandomCropFromBorders(crop_value=0.1, p=1.0),
            ElasticTransform((0, 0.25), interpolation=2, p=1.0),
            Resize((32, 16, 16), interpolation=1, resize_type=0, always_apply=True, p=1.0),
            Flip(0, p=1.0),
            Flip(1, p=1.0),
            Flip(2, p=1.0),
            RandomRotate90((1, 2), p=1.0),
        ],
        p=1.0,
        targets=[["image", "pmap", "mask"]],
    )

    image = np.random.rand(16, 8, 8, 42).astype(np.float32)
    pmap = np.random.rand(16, 8, 8, 4).astype(np.float32)
    mask = (np.random.rand(16, 8, 8, 1) > 0.5).astype(np.float32)

    out = aug(image=image, pmap=pmap, mask=mask)

    assert out["image"].dtype == np.float32
    assert out["pmap"].dtype == np.float32
    assert out["mask"].dtype == np.float32
    assert out["image"].shape == (32, 16, 16, 42)
    assert out["pmap"].shape == (32, 16, 16, 4)
    assert out["mask"].shape == (32, 16, 16, 1)
    assert np.isfinite(out["image"]).all()
    assert np.isfinite(out["pmap"]).all()
    assert np.isfinite(out["mask"]).all()
    assert np.all(np.isin(np.unique(out["mask"]), [0.0, 1.0]))


def test_random_scale_preserves_mask_dtype_and_finiteness():
    aug = Compose(
        [RandomScale(scale_limit=(1.25, 1.25), always_apply=True, p=1.0)],
        p=1.0,
        targets=[["image", "mask"]],
    )

    image = np.random.rand(12, 10, 8).astype(np.float32)
    mask = (np.random.rand(12, 10, 8) > 0.5)

    out = aug(image=image, mask=mask)

    assert out["image"].dtype == np.float32
    assert out["mask"].dtype == np.float32
    assert out["image"].shape == (15, 12, 10)
    assert out["mask"].shape == (15, 12, 10)
    assert np.isfinite(out["image"]).all()
    assert np.isfinite(out["mask"]).all()
    assert np.all(np.isin(np.unique(out["mask"]), [0.0, 1.0]))
