# Copyright (c) OpenMMLab. All rights reserved.
import copy

import numpy as np
import pytest

from mmgen.datasets.pipelines import Crop, FixedCrop


class TestAugmentations(object):

    @classmethod
    def setup_class(cls):
        cls.results = dict()
        cls.img_gt = np.random.rand(256, 128, 3).astype(np.float32)
        cls.img_lq = np.random.rand(64, 32, 3).astype(np.float32)

        cls.results = dict(
            lq=cls.img_lq,
            gt=cls.img_gt,
            scale=4,
            lq_path='fake_lq_path',
            gt_path='fake_gt_path')

        cls.results['img'] = np.random.rand(256, 256, 3).astype(np.float32)

        cls.results['img_a'] = np.random.rand(286, 286, 3).astype(np.float32)
        cls.results['img_b'] = np.random.rand(286, 286, 3).astype(np.float32)

    @staticmethod
    def check_crop(result_img_shape, result_bbox):
        crop_w = result_bbox[2] - result_bbox[0]
        """Check if the result_bbox is in correspond to result_img_shape."""
        crop_h = result_bbox[3] - result_bbox[1]
        crop_shape = (crop_h, crop_w)
        return result_img_shape == crop_shape

    @staticmethod
    def check_crop_around_semi(alpha):
        return ((alpha > 0) & (alpha < 255)).any()

    @staticmethod
    def check_keys_contain(result_keys, target_keys):
        """Check if all elements in target_keys is in result_keys."""
        return set(target_keys).issubset(set(result_keys))

    def test_crop(self):
        with pytest.raises(TypeError):
            Crop(['img'], (0.23, 0.1))

        # test center crop
        results = copy.deepcopy(self.results)
        center_crop = Crop(['img'], crop_size=(128, 128), random_crop=False)
        results = center_crop(results)
        assert results['img_crop_bbox'] == [64, 64, 128, 128]
        assert np.array_equal(self.results['img'][64:192, 64:192, :],
                              results['img'])

        # test random crop
        results = copy.deepcopy(self.results)
        random_crop = Crop(['img'], crop_size=(128, 128), random_crop=True)
        results = random_crop(results)
        assert 0 <= results['img_crop_bbox'][0] <= 128
        assert 0 <= results['img_crop_bbox'][1] <= 128
        assert results['img_crop_bbox'][2] == 128
        assert results['img_crop_bbox'][3] == 128

        # test random crop for lager size than the original shape
        results = copy.deepcopy(self.results)
        random_crop = Crop(['img'], crop_size=(512, 512), random_crop=True)
        results = random_crop(results)
        assert np.array_equal(self.results['img'], results['img'])
        assert str(random_crop) == (
            random_crop.__class__.__name__ +
            "keys=['img'], crop_size=(512, 512), random_crop=True")

    def test_fixed_crop(self):
        with pytest.raises(TypeError):
            FixedCrop(['img_a', 'img_b'], (0.23, 0.1))
        with pytest.raises(TypeError):
            FixedCrop(['img_a', 'img_b'], (256, 256), (0, 0.1))

        # test shape consistency
        results = copy.deepcopy(self.results)
        fixed_crop = FixedCrop(['img_a', 'img'], crop_size=(128, 128))
        with pytest.raises(ValueError):
            results = fixed_crop(results)

        # test given pos crop
        results = copy.deepcopy(self.results)
        given_pos_crop = FixedCrop(['img_a', 'img_b'],
                                   crop_size=(256, 256),
                                   crop_pos=(1, 1))
        results = given_pos_crop(results)
        assert results['img_a_crop_bbox'] == [1, 1, 256, 256]
        assert results['img_b_crop_bbox'] == [1, 1, 256, 256]
        assert np.array_equal(self.results['img_a'][1:257, 1:257, :],
                              results['img_a'])
        assert np.array_equal(self.results['img_b'][1:257, 1:257, :],
                              results['img_b'])

        # test given pos crop if pos > suitable pos
        results = copy.deepcopy(self.results)
        given_pos_crop = FixedCrop(['img_a', 'img_b'],
                                   crop_size=(256, 256),
                                   crop_pos=(280, 280))
        results = given_pos_crop(results)
        assert results['img_a_crop_bbox'] == [280, 280, 6, 6]
        assert results['img_b_crop_bbox'] == [280, 280, 6, 6]
        assert np.array_equal(self.results['img_a'][280:, 280:, :],
                              results['img_a'])
        assert np.array_equal(self.results['img_b'][280:, 280:, :],
                              results['img_b'])
        assert str(given_pos_crop) == (
            given_pos_crop.__class__.__name__ +
            "keys=['img_a', 'img_b'], crop_size=(256, 256), " +
            'crop_pos=(280, 280)')

        # test random initialized fixed crop
        results = copy.deepcopy(self.results)
        random_fixed_crop = FixedCrop(['img_a', 'img_b'],
                                      crop_size=(256, 256),
                                      crop_pos=None)
        results = random_fixed_crop(results)
        assert 0 <= results['img_a_crop_bbox'][0] <= 30
        assert 0 <= results['img_a_crop_bbox'][1] <= 30
        assert results['img_a_crop_bbox'][2] == 256
        assert results['img_a_crop_bbox'][3] == 256
        x_offset, y_offset, crop_w, crop_h = results['img_a_crop_bbox']
        assert x_offset == results['img_b_crop_bbox'][0]
        assert y_offset == results['img_b_crop_bbox'][1]
        assert crop_w == results['img_b_crop_bbox'][2]
        assert crop_h == results['img_b_crop_bbox'][3]
        assert np.array_equal(
            self.results['img_a'][y_offset:y_offset + crop_h,
                                  x_offset:x_offset + crop_w, :],
            results['img_a'])
        assert np.array_equal(
            self.results['img_b'][y_offset:y_offset + crop_h,
                                  x_offset:x_offset + crop_w, :],
            results['img_b'])

        # test given pos crop for lager size than the original shape
        results = copy.deepcopy(self.results)
        given_pos_crop = FixedCrop(['img_a', 'img_b'],
                                   crop_size=(512, 512),
                                   crop_pos=(1, 1))
        results = given_pos_crop(results)
        assert results['img_a_crop_bbox'] == [1, 1, 285, 285]
        assert results['img_b_crop_bbox'] == [1, 1, 285, 285]
        assert np.array_equal(self.results['img_a'][1:, 1:, :],
                              results['img_a'])
        assert np.array_equal(self.results['img_b'][1:, 1:, :],
                              results['img_b'])
        assert str(given_pos_crop) == (
            given_pos_crop.__class__.__name__ +
            "keys=['img_a', 'img_b'], crop_size=(512, 512), crop_pos=(1, 1)")

        # test random initialized fixed crop for lager size
        # than the original shape
        results = copy.deepcopy(self.results)
        random_fixed_crop = FixedCrop(['img_a', 'img_b'],
                                      crop_size=(512, 512),
                                      crop_pos=None)
        results = random_fixed_crop(results)
        assert results['img_a_crop_bbox'] == [0, 0, 286, 286]
        assert results['img_b_crop_bbox'] == [0, 0, 286, 286]
        assert np.array_equal(self.results['img_a'], results['img_a'])
        assert np.array_equal(self.results['img_b'], results['img_b'])
        assert str(random_fixed_crop) == (
            random_fixed_crop.__class__.__name__ +
            "keys=['img_a', 'img_b'], crop_size=(512, 512), crop_pos=None")
