import torch
from torchvision.transforms import Compose
import pytest
import random
import numpy as np
import warnings
from _assert_utils import assert_equal

try:
    from scipy import stats
except ImportError:
    stats = None


with warnings.catch_warnings(record=True):
    warnings.simplefilter("always")
    import torchvision.transforms._transforms_video as transforms


class TestVideoTransforms():

    def test_random_crop_video(self):
        numFrames = random.randint(4, 128)
        height = random.randint(10, 32) * 2
        width = random.randint(10, 32) * 2
        oheight = random.randint(5, (height - 2) / 2) * 2
        owidth = random.randint(5, (width - 2) / 2) * 2
        clip = torch.randint(0, 256, (numFrames, height, width, 3), dtype=torch.uint8)
        result = Compose([
            transforms.ToTensorVideo(),
            transforms.RandomCropVideo((oheight, owidth)),
        ])(clip)
        assert result.size(2) == oheight
        assert result.size(3) == owidth

        transforms.RandomCropVideo((oheight, owidth)).__repr__()

    def test_random_resized_crop_video(self):
        numFrames = random.randint(4, 128)
        height = random.randint(10, 32) * 2
        width = random.randint(10, 32) * 2
        oheight = random.randint(5, (height - 2) / 2) * 2
        owidth = random.randint(5, (width - 2) / 2) * 2
        clip = torch.randint(0, 256, (numFrames, height, width, 3), dtype=torch.uint8)
        result = Compose([
            transforms.ToTensorVideo(),
            transforms.RandomResizedCropVideo((oheight, owidth)),
        ])(clip)
        assert result.size(2) == oheight
        assert result.size(3) == owidth

        transforms.RandomResizedCropVideo((oheight, owidth)).__repr__()

    def test_center_crop_video(self):
        numFrames = random.randint(4, 128)
        height = random.randint(10, 32) * 2
        width = random.randint(10, 32) * 2
        oheight = random.randint(5, (height - 2) / 2) * 2
        owidth = random.randint(5, (width - 2) / 2) * 2

        clip = torch.ones((numFrames, height, width, 3), dtype=torch.uint8) * 255
        oh1 = (height - oheight) // 2
        ow1 = (width - owidth) // 2
        clipNarrow = clip[:, oh1:oh1 + oheight, ow1:ow1 + owidth, :]
        clipNarrow.fill_(0)
        result = Compose([
            transforms.ToTensorVideo(),
            transforms.CenterCropVideo((oheight, owidth)),
        ])(clip)

        msg = "height: " + str(height) + " width: " \
            + str(width) + " oheight: " + str(oheight) + " owidth: " + str(owidth)
        assert result.sum().item() == 0, msg

        oheight += 1
        owidth += 1
        result = Compose([
            transforms.ToTensorVideo(),
            transforms.CenterCropVideo((oheight, owidth)),
        ])(clip)
        sum1 = result.sum()

        msg = "height: " + str(height) + " width: " \
            + str(width) + " oheight: " + str(oheight) + " owidth: " + str(owidth)
        assert sum1.item() > 1, msg

        oheight += 1
        owidth += 1
        result = Compose([
            transforms.ToTensorVideo(),
            transforms.CenterCropVideo((oheight, owidth)),
        ])(clip)
        sum2 = result.sum()

        msg = "height: " + str(height) + " width: " \
            + str(width) + " oheight: " + str(oheight) + " owidth: " + str(owidth)
        assert sum2.item() > 1, msg
        assert sum2.item() > sum1.item(), msg

    @pytest.mark.skipif(stats is None, reason='scipy.stats is not available')
    @pytest.mark.parametrize('channels', [1, 3])
    def test_normalize_video(self, channels):
        def samples_from_standard_normal(tensor):
            p_value = stats.kstest(list(tensor.view(-1)), 'norm', args=(0, 1)).pvalue
            return p_value > 0.0001

        random_state = random.getstate()
        random.seed(42)

        numFrames = random.randint(4, 128)
        height = random.randint(32, 256)
        width = random.randint(32, 256)
        mean = random.random()
        std = random.random()
        clip = torch.normal(mean, std, size=(channels, numFrames, height, width))
        mean = [clip[c].mean().item() for c in range(channels)]
        std = [clip[c].std().item() for c in range(channels)]
        normalized = transforms.NormalizeVideo(mean, std)(clip)
        assert samples_from_standard_normal(normalized)
        random.setstate(random_state)

        # Checking the optional in-place behaviour
        tensor = torch.rand((3, 128, 16, 16))
        tensor_inplace = transforms.NormalizeVideo((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)(tensor)
        assert_equal(tensor, tensor_inplace)

        transforms.NormalizeVideo((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True).__repr__()

    def test_to_tensor_video(self):
        numFrames, height, width = 64, 4, 4
        trans = transforms.ToTensorVideo()

        with pytest.raises(TypeError):
            trans(np.random.rand(numFrames, height, width, 1).tolist())
            trans(torch.rand((numFrames, height, width, 1), dtype=torch.float))

        with pytest.raises(ValueError):
            trans(torch.ones((3, numFrames, height, width, 3), dtype=torch.uint8))
            trans(torch.ones((height, width, 3), dtype=torch.uint8))
            trans(torch.ones((width, 3), dtype=torch.uint8))
            trans(torch.ones((3), dtype=torch.uint8))

        trans.__repr__()

    @pytest.mark.skipif(stats is None, reason='scipy.stats not available')
    def test_random_horizontal_flip_video(self):
        random_state = random.getstate()
        random.seed(42)
        clip = torch.rand((3, 4, 112, 112), dtype=torch.float)
        hclip = clip.flip((-1))

        num_samples = 250
        num_horizontal = 0
        for _ in range(num_samples):
            out = transforms.RandomHorizontalFlipVideo()(clip)
            if torch.all(torch.eq(out, hclip)):
                num_horizontal += 1

        p_value = stats.binom_test(num_horizontal, num_samples, p=0.5)
        random.setstate(random_state)
        assert p_value > 0.0001

        num_samples = 250
        num_horizontal = 0
        for _ in range(num_samples):
            out = transforms.RandomHorizontalFlipVideo(p=0.7)(clip)
            if torch.all(torch.eq(out, hclip)):
                num_horizontal += 1

        p_value = stats.binom_test(num_horizontal, num_samples, p=0.7)
        random.setstate(random_state)
        assert p_value > 0.0001

        transforms.RandomHorizontalFlipVideo().__repr__()


if __name__ == '__main__':
    pytest.main([__file__])
