import numpy as np
import pytest
import torch

from careamics.transforms import ImageRestorationTTA, XYFlip, XYRandomRotate90


@pytest.mark.parametrize(
    "shape",
    [
        # 2D
        (1, 1, 8, 8),
        (2, 1, 8, 8),
        (1, 3, 8, 8),
        (2, 3, 8, 8),
        # 3D
        (1, 1, 8, 8, 8),
        (2, 1, 8, 8, 8),
        (1, 3, 8, 8, 8),
        (2, 3, 8, 8, 8),
    ],
)
def test_forward(shape):
    """Test that the transformations are correct."""
    array = np.arange(np.prod(shape)).reshape(shape)
    tensor = torch.Tensor(array)

    tta = ImageRestorationTTA()
    augmented = tta.forward(tensor)
    assert len(augmented) == 8

    # check that the shape is the same
    for aug in augmented:
        assert aug.shape == tensor.shape

    # check that all transformed tensors are unique
    for i, aug1 in enumerate(augmented):
        for j, aug2 in enumerate(augmented):
            if i != j:
                assert not torch.allclose(aug1, aug2)


@pytest.mark.parametrize(
    "shape",
    [
        # 2D
        (1, 1, 8, 8),
        (1, 3, 8, 8),
        # 3D
        (1, 1, 8, 8, 8),
        (1, 3, 8, 8, 8),
    ],
)
def test_same_transforms(shape):
    """Check that all arrays produced by the rotation and flip
    transforms are generated by the TTA.

    Note that the transforms require no sample dimension.
    """
    array = np.arange(np.prod(shape)).reshape(shape)
    tensor = torch.Tensor(array)

    # apply forward transformation
    tta = ImageRestorationTTA()
    augmented = tta.forward(tensor)

    # transofrms
    rot = XYRandomRotate90(seed=42)
    flip = XYFlip(seed=42)

    for _ in range(100):
        rotated, *_ = rot(tensor)
        flipped, *_ = flip(rotated)

        # check that is in the augmented list
        assert any(torch.allclose(torch.Tensor(flipped), aug) for aug in augmented)


@pytest.mark.parametrize(
    "shape",
    [
        # 2D
        (1, 1, 8, 8),
        (2, 1, 8, 8),
        (1, 3, 8, 8),
        (2, 3, 8, 8),
        # 3D
        (1, 1, 8, 8, 8),
        (2, 1, 8, 8, 8),
        (1, 3, 8, 8, 8),
        (2, 3, 8, 8, 8),
    ],
)
def test_backward(shape):
    """Test that the backward transformation returns the original tensor."""
    array = np.arange(np.prod(shape)).reshape(shape)
    tensor = torch.Tensor(array)

    # apply forward transformation
    tta = ImageRestorationTTA()
    augmented = tta.forward(tensor)

    # apply backward transformation
    original = tta.backward(augmented)
    assert torch.allclose(tensor, original)
