import itertools

import re

import PIL.Image
import pytest
import torch

from common_utils import (
    assert_equal,
    DEFAULT_EXTRA_DIMS,
    make_bounding_box,
    make_detection_mask,
    make_image,
    make_images,
    make_segmentation_mask,
    make_video,
    make_videos,
)

from prototype_common_utils import make_label, make_one_hot_labels

from torchvision.datapoints import BoundingBox, BoundingBoxFormat, Image, Mask, Video
from torchvision.prototype import datapoints, transforms
from torchvision.transforms.v2._utils import _convert_fill_arg
from torchvision.transforms.v2.functional import InterpolationMode, pil_to_tensor, to_image_pil
from torchvision.transforms.v2.utils import check_type, is_simple_tensor

BATCH_EXTRA_DIMS = [extra_dims for extra_dims in DEFAULT_EXTRA_DIMS if extra_dims]


def parametrize(transforms_with_inputs):
    return pytest.mark.parametrize(
        ("transform", "input"),
        [
            pytest.param(
                transform,
                input,
                id=f"{type(transform).__name__}-{type(input).__module__}.{type(input).__name__}-{idx}",
            )
            for transform, inputs in transforms_with_inputs
            for idx, input in enumerate(inputs)
        ],
    )


@parametrize(
    [
        (
            transform,
            [
                dict(inpt=inpt, one_hot_label=one_hot_label)
                for inpt, one_hot_label in itertools.product(
                    itertools.chain(
                        make_images(extra_dims=BATCH_EXTRA_DIMS, dtypes=[torch.float]),
                        make_videos(extra_dims=BATCH_EXTRA_DIMS, dtypes=[torch.float]),
                    ),
                    make_one_hot_labels(extra_dims=BATCH_EXTRA_DIMS, dtypes=[torch.float]),
                )
            ],
        )
        for transform in [
            transforms.RandomMixup(alpha=1.0),
            transforms.RandomCutmix(alpha=1.0),
        ]
    ]
)
def test_mixup_cutmix(transform, input):
    transform(input)

    input_copy = dict(input)
    input_copy["path"] = "/path/to/somewhere"
    input_copy["num"] = 1234
    transform(input_copy)

    # Check if we raise an error if sample contains bbox or mask or label
    err_msg = "does not support PIL images, bounding boxes, masks and plain labels"
    input_copy = dict(input)
    for unsup_data in [
        make_label(),
        make_bounding_box(format="XYXY"),
        make_detection_mask(),
        make_segmentation_mask(),
    ]:
        input_copy["unsupported"] = unsup_data
        with pytest.raises(TypeError, match=err_msg):
            transform(input_copy)


class TestSimpleCopyPaste:
    def create_fake_image(self, mocker, image_type):
        if image_type == PIL.Image.Image:
            return PIL.Image.new("RGB", (32, 32), 123)
        return mocker.MagicMock(spec=image_type)

    def test__extract_image_targets_assertion(self, mocker):
        transform = transforms.SimpleCopyPaste()

        flat_sample = [
            # images, batch size = 2
            self.create_fake_image(mocker, Image),
            # labels, bboxes, masks
            mocker.MagicMock(spec=datapoints.Label),
            mocker.MagicMock(spec=BoundingBox),
            mocker.MagicMock(spec=Mask),
            # labels, bboxes, masks
            mocker.MagicMock(spec=BoundingBox),
            mocker.MagicMock(spec=Mask),
        ]

        with pytest.raises(TypeError, match="requires input sample to contain equal sized list of Images"):
            transform._extract_image_targets(flat_sample)

    @pytest.mark.parametrize("image_type", [Image, PIL.Image.Image, torch.Tensor])
    @pytest.mark.parametrize("label_type", [datapoints.Label, datapoints.OneHotLabel])
    def test__extract_image_targets(self, image_type, label_type, mocker):
        transform = transforms.SimpleCopyPaste()

        flat_sample = [
            # images, batch size = 2
            self.create_fake_image(mocker, image_type),
            self.create_fake_image(mocker, image_type),
            # labels, bboxes, masks
            mocker.MagicMock(spec=label_type),
            mocker.MagicMock(spec=BoundingBox),
            mocker.MagicMock(spec=Mask),
            # labels, bboxes, masks
            mocker.MagicMock(spec=label_type),
            mocker.MagicMock(spec=BoundingBox),
            mocker.MagicMock(spec=Mask),
        ]

        images, targets = transform._extract_image_targets(flat_sample)

        assert len(images) == len(targets) == 2
        if image_type == PIL.Image.Image:
            torch.testing.assert_close(images[0], pil_to_tensor(flat_sample[0]))
            torch.testing.assert_close(images[1], pil_to_tensor(flat_sample[1]))
        else:
            assert images[0] == flat_sample[0]
            assert images[1] == flat_sample[1]

        for target in targets:
            for key, type_ in [
                ("boxes", BoundingBox),
                ("masks", Mask),
                ("labels", label_type),
            ]:
                assert key in target
                assert isinstance(target[key], type_)
                assert target[key] in flat_sample

    @pytest.mark.parametrize("label_type", [datapoints.Label, datapoints.OneHotLabel])
    def test__copy_paste(self, label_type):
        image = 2 * torch.ones(3, 32, 32)
        masks = torch.zeros(2, 32, 32)
        masks[0, 3:9, 2:8] = 1
        masks[1, 20:30, 20:30] = 1
        labels = torch.tensor([1, 2])
        blending = True
        resize_interpolation = InterpolationMode.BILINEAR
        antialias = None
        if label_type == datapoints.OneHotLabel:
            labels = torch.nn.functional.one_hot(labels, num_classes=5)
        target = {
            "boxes": BoundingBox(
                torch.tensor([[2.0, 3.0, 8.0, 9.0], [20.0, 20.0, 30.0, 30.0]]), format="XYXY", spatial_size=(32, 32)
            ),
            "masks": Mask(masks),
            "labels": label_type(labels),
        }

        paste_image = 10 * torch.ones(3, 32, 32)
        paste_masks = torch.zeros(2, 32, 32)
        paste_masks[0, 13:19, 12:18] = 1
        paste_masks[1, 15:19, 1:8] = 1
        paste_labels = torch.tensor([3, 4])
        if label_type == datapoints.OneHotLabel:
            paste_labels = torch.nn.functional.one_hot(paste_labels, num_classes=5)
        paste_target = {
            "boxes": BoundingBox(
                torch.tensor([[12.0, 13.0, 19.0, 18.0], [1.0, 15.0, 8.0, 19.0]]), format="XYXY", spatial_size=(32, 32)
            ),
            "masks": Mask(paste_masks),
            "labels": label_type(paste_labels),
        }

        transform = transforms.SimpleCopyPaste()
        random_selection = torch.tensor([0, 1])
        output_image, output_target = transform._copy_paste(
            image, target, paste_image, paste_target, random_selection, blending, resize_interpolation, antialias
        )

        assert output_image.unique().tolist() == [2, 10]
        assert output_target["boxes"].shape == (4, 4)
        torch.testing.assert_close(output_target["boxes"][:2, :], target["boxes"])
        torch.testing.assert_close(output_target["boxes"][2:, :], paste_target["boxes"])

        expected_labels = torch.tensor([1, 2, 3, 4])
        if label_type == datapoints.OneHotLabel:
            expected_labels = torch.nn.functional.one_hot(expected_labels, num_classes=5)
        torch.testing.assert_close(output_target["labels"], label_type(expected_labels))

        assert output_target["masks"].shape == (4, 32, 32)
        torch.testing.assert_close(output_target["masks"][:2, :], target["masks"])
        torch.testing.assert_close(output_target["masks"][2:, :], paste_target["masks"])


class TestFixedSizeCrop:
    def test__get_params(self, mocker):
        crop_size = (7, 7)
        batch_shape = (10,)
        spatial_size = (11, 5)

        transform = transforms.FixedSizeCrop(size=crop_size)

        flat_inputs = [
            make_image(size=spatial_size, color_space="RGB"),
            make_bounding_box(format=BoundingBoxFormat.XYXY, spatial_size=spatial_size, batch_dims=batch_shape),
        ]
        params = transform._get_params(flat_inputs)

        assert params["needs_crop"]
        assert params["height"] <= crop_size[0]
        assert params["width"] <= crop_size[1]

        assert (
            isinstance(params["is_valid"], torch.Tensor)
            and params["is_valid"].dtype is torch.bool
            and params["is_valid"].shape == batch_shape
        )

        assert params["needs_pad"]
        assert any(pad > 0 for pad in params["padding"])

    @pytest.mark.parametrize("needs", list(itertools.product((False, True), repeat=2)))
    def test__transform(self, mocker, needs):
        fill_sentinel = 12
        padding_mode_sentinel = mocker.MagicMock()

        transform = transforms.FixedSizeCrop((-1, -1), fill=fill_sentinel, padding_mode=padding_mode_sentinel)
        transform._transformed_types = (mocker.MagicMock,)
        mocker.patch("torchvision.prototype.transforms._geometry.has_any", return_value=True)

        needs_crop, needs_pad = needs
        top_sentinel = mocker.MagicMock()
        left_sentinel = mocker.MagicMock()
        height_sentinel = mocker.MagicMock()
        width_sentinel = mocker.MagicMock()
        is_valid = mocker.MagicMock() if needs_crop else None
        padding_sentinel = mocker.MagicMock()
        mocker.patch(
            "torchvision.prototype.transforms._geometry.FixedSizeCrop._get_params",
            return_value=dict(
                needs_crop=needs_crop,
                top=top_sentinel,
                left=left_sentinel,
                height=height_sentinel,
                width=width_sentinel,
                is_valid=is_valid,
                padding=padding_sentinel,
                needs_pad=needs_pad,
            ),
        )

        inpt_sentinel = mocker.MagicMock()

        mock_crop = mocker.patch("torchvision.prototype.transforms._geometry.F.crop")
        mock_pad = mocker.patch("torchvision.prototype.transforms._geometry.F.pad")
        transform(inpt_sentinel)

        if needs_crop:
            mock_crop.assert_called_once_with(
                inpt_sentinel,
                top=top_sentinel,
                left=left_sentinel,
                height=height_sentinel,
                width=width_sentinel,
            )
        else:
            mock_crop.assert_not_called()

        if needs_pad:
            # If we cropped before, the input to F.pad is no longer inpt_sentinel. Thus, we can't use
            # `MagicMock.assert_called_once_with` and have to perform the checks manually
            mock_pad.assert_called_once()
            args, kwargs = mock_pad.call_args
            if not needs_crop:
                assert args[0] is inpt_sentinel
            assert args[1] is padding_sentinel
            fill_sentinel = _convert_fill_arg(fill_sentinel)
            assert kwargs == dict(fill=fill_sentinel, padding_mode=padding_mode_sentinel)
        else:
            mock_pad.assert_not_called()

    def test__transform_culling(self, mocker):
        batch_size = 10
        spatial_size = (10, 10)

        is_valid = torch.randint(0, 2, (batch_size,), dtype=torch.bool)
        mocker.patch(
            "torchvision.prototype.transforms._geometry.FixedSizeCrop._get_params",
            return_value=dict(
                needs_crop=True,
                top=0,
                left=0,
                height=spatial_size[0],
                width=spatial_size[1],
                is_valid=is_valid,
                needs_pad=False,
            ),
        )

        bounding_boxes = make_bounding_box(
            format=BoundingBoxFormat.XYXY, spatial_size=spatial_size, batch_dims=(batch_size,)
        )
        masks = make_detection_mask(size=spatial_size, batch_dims=(batch_size,))
        labels = make_label(extra_dims=(batch_size,))

        transform = transforms.FixedSizeCrop((-1, -1))
        mocker.patch("torchvision.prototype.transforms._geometry.has_any", return_value=True)

        output = transform(
            dict(
                bounding_boxes=bounding_boxes,
                masks=masks,
                labels=labels,
            )
        )

        assert_equal(output["bounding_boxes"], bounding_boxes[is_valid])
        assert_equal(output["masks"], masks[is_valid])
        assert_equal(output["labels"], labels[is_valid])

    def test__transform_bounding_box_clamping(self, mocker):
        batch_size = 3
        spatial_size = (10, 10)

        mocker.patch(
            "torchvision.prototype.transforms._geometry.FixedSizeCrop._get_params",
            return_value=dict(
                needs_crop=True,
                top=0,
                left=0,
                height=spatial_size[0],
                width=spatial_size[1],
                is_valid=torch.full((batch_size,), fill_value=True),
                needs_pad=False,
            ),
        )

        bounding_box = make_bounding_box(
            format=BoundingBoxFormat.XYXY, spatial_size=spatial_size, batch_dims=(batch_size,)
        )
        mock = mocker.patch("torchvision.prototype.transforms._geometry.F.clamp_bounding_box")

        transform = transforms.FixedSizeCrop((-1, -1))
        mocker.patch("torchvision.prototype.transforms._geometry.has_any", return_value=True)

        transform(bounding_box)

        mock.assert_called_once()


class TestLabelToOneHot:
    def test__transform(self):
        categories = ["apple", "pear", "pineapple"]
        labels = datapoints.Label(torch.tensor([0, 1, 2, 1]), categories=categories)
        transform = transforms.LabelToOneHot()
        ohe_labels = transform(labels)
        assert isinstance(ohe_labels, datapoints.OneHotLabel)
        assert ohe_labels.shape == (4, 3)
        assert ohe_labels.categories == labels.categories == categories


class TestPermuteDimensions:
    @pytest.mark.parametrize(
        ("dims", "inverse_dims"),
        [
            (
                {Image: (2, 1, 0), Video: None},
                {Image: (2, 1, 0), Video: None},
            ),
            (
                {Image: (2, 1, 0), Video: (1, 2, 3, 0)},
                {Image: (2, 1, 0), Video: (3, 0, 1, 2)},
            ),
        ],
    )
    def test_call(self, dims, inverse_dims):
        sample = dict(
            image=make_image(),
            bounding_box=make_bounding_box(format=BoundingBoxFormat.XYXY),
            video=make_video(),
            str="str",
            int=0,
        )

        transform = transforms.PermuteDimensions(dims)
        transformed_sample = transform(sample)

        for key, value in sample.items():
            value_type = type(value)
            transformed_value = transformed_sample[key]

            if check_type(value, (Image, is_simple_tensor, Video)):
                if transform.dims.get(value_type) is not None:
                    assert transformed_value.permute(inverse_dims[value_type]).equal(value)
                assert type(transformed_value) == torch.Tensor
            else:
                assert transformed_value is value

    @pytest.mark.filterwarnings("error")
    def test_plain_tensor_call(self):
        tensor = torch.empty((2, 3, 4))
        transform = transforms.PermuteDimensions(dims=(1, 2, 0))

        assert transform(tensor).shape == (3, 4, 2)

    @pytest.mark.parametrize("other_type", [Image, Video])
    def test_plain_tensor_warning(self, other_type):
        with pytest.warns(UserWarning, match=re.escape("`torch.Tensor` will *not* be transformed")):
            transforms.PermuteDimensions(dims={torch.Tensor: (0, 1), other_type: (1, 0)})


class TestTransposeDimensions:
    @pytest.mark.parametrize(
        "dims",
        [
            (-1, -2),
            {Image: (1, 2), Video: None},
        ],
    )
    def test_call(self, dims):
        sample = dict(
            image=make_image(),
            bounding_box=make_bounding_box(format=BoundingBoxFormat.XYXY),
            video=make_video(),
            str="str",
            int=0,
        )

        transform = transforms.TransposeDimensions(dims)
        transformed_sample = transform(sample)

        for key, value in sample.items():
            value_type = type(value)
            transformed_value = transformed_sample[key]

            transposed_dims = transform.dims.get(value_type)
            if check_type(value, (Image, is_simple_tensor, Video)):
                if transposed_dims is not None:
                    assert transformed_value.transpose(*transposed_dims).equal(value)
                assert type(transformed_value) == torch.Tensor
            else:
                assert transformed_value is value

    @pytest.mark.filterwarnings("error")
    def test_plain_tensor_call(self):
        tensor = torch.empty((2, 3, 4))
        transform = transforms.TransposeDimensions(dims=(0, 2))

        assert transform(tensor).shape == (4, 3, 2)

    @pytest.mark.parametrize("other_type", [Image, Video])
    def test_plain_tensor_warning(self, other_type):
        with pytest.warns(UserWarning, match=re.escape("`torch.Tensor` will *not* be transformed")):
            transforms.TransposeDimensions(dims={torch.Tensor: (0, 1), other_type: (1, 0)})


import importlib.machinery
import importlib.util
from pathlib import Path


def import_transforms_from_references(reference):
    HERE = Path(__file__).parent
    PROJECT_ROOT = HERE.parent

    loader = importlib.machinery.SourceFileLoader(
        "transforms", str(PROJECT_ROOT / "references" / reference / "transforms.py")
    )
    spec = importlib.util.spec_from_loader("transforms", loader)
    module = importlib.util.module_from_spec(spec)
    loader.exec_module(module)
    return module


det_transforms = import_transforms_from_references("detection")


def test_fixed_sized_crop_against_detection_reference():
    def make_datapoints():
        size = (600, 800)
        num_objects = 22

        pil_image = to_image_pil(make_image(size=size, color_space="RGB"))
        target = {
            "boxes": make_bounding_box(spatial_size=size, format="XYXY", batch_dims=(num_objects,), dtype=torch.float),
            "labels": make_label(extra_dims=(num_objects,), categories=80),
            "masks": make_detection_mask(size=size, num_objects=num_objects, dtype=torch.long),
        }

        yield (pil_image, target)

        tensor_image = torch.Tensor(make_image(size=size, color_space="RGB"))
        target = {
            "boxes": make_bounding_box(spatial_size=size, format="XYXY", batch_dims=(num_objects,), dtype=torch.float),
            "labels": make_label(extra_dims=(num_objects,), categories=80),
            "masks": make_detection_mask(size=size, num_objects=num_objects, dtype=torch.long),
        }

        yield (tensor_image, target)

        datapoint_image = make_image(size=size, color_space="RGB")
        target = {
            "boxes": make_bounding_box(spatial_size=size, format="XYXY", batch_dims=(num_objects,), dtype=torch.float),
            "labels": make_label(extra_dims=(num_objects,), categories=80),
            "masks": make_detection_mask(size=size, num_objects=num_objects, dtype=torch.long),
        }

        yield (datapoint_image, target)

    t = transforms.FixedSizeCrop((1024, 1024), fill=0)
    t_ref = det_transforms.FixedSizeCrop((1024, 1024), fill=0)

    for dp in make_datapoints():
        # We should use prototype transform first as reference transform performs inplace target update
        torch.manual_seed(12)
        output = t(dp)

        torch.manual_seed(12)
        expected_output = t_ref(*dp)

        assert_equal(expected_output, output)
