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

import numpy as np
import pytest
import torchvision.transforms as TF
from mmdet.core import BitmapMasks, PolygonMasks
from PIL import Image

import mmocr.datasets.pipelines.transforms as transforms


@mock.patch('%s.transforms.np.random.random_sample' % __name__)
@mock.patch('%s.transforms.np.random.randint' % __name__)
def test_random_crop_instances(mock_randint, mock_sample):

    img_gt = np.array([[0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 1, 1, 1],
                       [0, 0, 1, 1, 1], [0, 0, 1, 1, 1]])
    # test target is bigger than img size in sample_offset
    mock_sample.side_effect = [1]
    rci = transforms.RandomCropInstances(6, instance_key='gt_kernels')
    (i, j) = rci.sample_offset(img_gt, (5, 5))
    assert i == 0
    assert j == 0

    # test the second branch in sample_offset

    rci = transforms.RandomCropInstances(3, instance_key='gt_kernels')
    mock_sample.side_effect = [1]
    mock_randint.side_effect = [1, 2]
    (i, j) = rci.sample_offset(img_gt, (5, 5))
    assert i == 1
    assert j == 2

    mock_sample.side_effect = [1]
    mock_randint.side_effect = [1, 2]
    rci = transforms.RandomCropInstances(5, instance_key='gt_kernels')
    (i, j) = rci.sample_offset(img_gt, (5, 5))
    assert i == 0
    assert j == 0

    # test the first bracnh is sample_offset

    rci = transforms.RandomCropInstances(3, instance_key='gt_kernels')
    mock_sample.side_effect = [0.1]
    mock_randint.side_effect = [1, 1]
    (i, j) = rci.sample_offset(img_gt, (5, 5))
    assert i == 1
    assert j == 1

    # test crop_img(img, offset, target_size)

    img = img_gt
    offset = [0, 0]
    target = [6, 6]
    crop = rci.crop_img(img, offset, target)
    assert np.allclose(img, crop[0])
    assert np.allclose(crop[1], [0, 0, 5, 5])

    target = [3, 2]
    crop = rci.crop_img(img, offset, target)
    assert np.allclose(np.array([[0, 0], [0, 0], [0, 0]]), crop[0])
    assert np.allclose(crop[1], [0, 0, 2, 3])

    # test crop_bboxes
    canvas_box = np.array([2, 3, 5, 5])
    bboxes = np.array([[2, 3, 4, 4], [0, 0, 1, 1], [1, 2, 4, 4],
                       [0, 0, 10, 10]])
    kept_bboxes, kept_idx = rci.crop_bboxes(bboxes, canvas_box)
    assert np.allclose(kept_bboxes,
                       np.array([[0, 0, 2, 1], [0, 0, 2, 1], [0, 0, 3, 2]]))
    assert kept_idx == [0, 2, 3]

    bboxes = np.array([[10, 10, 11, 11], [0, 0, 1, 1]])
    kept_bboxes, kept_idx = rci.crop_bboxes(bboxes, canvas_box)
    assert kept_bboxes.size == 0
    assert kept_bboxes.shape == (0, 4)
    assert len(kept_idx) == 0

    # test __call__
    rci = transforms.RandomCropInstances(3, instance_key='gt_kernels')
    results = {}
    gt_kernels = [img_gt, img_gt.copy()]
    results['gt_kernels'] = BitmapMasks(gt_kernels, 5, 5)
    results['img'] = img_gt.copy()
    results['mask_fields'] = ['gt_kernels']
    mock_sample.side_effect = [0.1]
    mock_randint.side_effect = [1, 1]
    output = rci(results)
    target = np.array([[0, 0, 0], [0, 1, 1], [0, 1, 1]])
    assert output['img_shape'] == (3, 3)

    assert np.allclose(output['img'], target)

    assert np.allclose(output['gt_kernels'].masks[0], target)
    assert np.allclose(output['gt_kernels'].masks[1], target)


@mock.patch('%s.transforms.np.random.random_sample' % __name__)
def test_scale_aspect_jitter(mock_random):
    img_scale = [(3000, 1000)]  # unused
    ratio_range = (0.5, 1.5)
    aspect_ratio_range = (1, 1)
    multiscale_mode = 'value'
    long_size_bound = 2000
    short_size_bound = 640
    resize_type = 'long_short_bound'
    keep_ratio = False
    jitter = transforms.ScaleAspectJitter(
        img_scale=img_scale,
        ratio_range=ratio_range,
        aspect_ratio_range=aspect_ratio_range,
        multiscale_mode=multiscale_mode,
        long_size_bound=long_size_bound,
        short_size_bound=short_size_bound,
        resize_type=resize_type,
        keep_ratio=keep_ratio)
    mock_random.side_effect = [0.5]

    # test sample_from_range

    result = jitter.sample_from_range([100, 200])
    assert result == 150

    # test _random_scale
    results = {}
    results['img'] = np.zeros((4000, 1000))
    mock_random.side_effect = [0.5, 1]
    jitter._random_scale(results)
    # scale1 0.5， scale2=1 scale =0.5  650/1000, w, h
    # print(results['scale'])
    assert results['scale'] == (650, 2600)


@mock.patch('%s.transforms.np.random.random_sample' % __name__)
def test_random_rotate(mock_random):

    mock_random.side_effect = [0.5, 0]
    results = {}
    img = np.random.rand(5, 5)
    results['img'] = img.copy()
    results['mask_fields'] = ['masks']
    gt_kernels = [results['img'].copy()]
    results['masks'] = BitmapMasks(gt_kernels, 5, 5)

    rotater = transforms.RandomRotateTextDet()

    results = rotater(results)
    assert np.allclose(results['img'], img)
    assert np.allclose(results['masks'].masks, img)


def test_color_jitter():
    img = np.ones((64, 256, 3), dtype=np.uint8)
    results = {'img': img}

    pt_official_color_jitter = TF.ColorJitter()
    output1 = pt_official_color_jitter(img)

    color_jitter = transforms.ColorJitter()
    output2 = color_jitter(results)

    assert np.allclose(output1, output2['img'])


def test_affine_jitter():
    img = np.ones((64, 256, 3), dtype=np.uint8)
    results = {'img': img}

    pt_official_affine_jitter = TF.RandomAffine(degrees=0)
    output1 = pt_official_affine_jitter(Image.fromarray(img))

    affine_jitter = transforms.AffineJitter(
        degrees=0,
        translate=None,
        scale=None,
        shear=None,
        resample=False,
        fillcolor=0)
    output2 = affine_jitter(results)

    assert np.allclose(np.array(output1), output2['img'])


def test_random_scale():
    h, w, c = 100, 100, 3
    img = np.ones((h, w, c), dtype=np.uint8)
    results = {'img': img, 'img_shape': (h, w, c)}

    polygon = np.array([0., 0., 0., 10., 10., 10., 10., 0.])

    results['gt_masks'] = PolygonMasks([[polygon]], *(img.shape[:2]))
    results['mask_fields'] = ['gt_masks']

    size = 100
    scale = (2., 2.)
    random_scaler = transforms.RandomScaling(size=size, scale=scale)

    results = random_scaler(results)

    out_img = results['img']
    out_poly = results['gt_masks'].masks[0][0]
    gt_poly = polygon * 2

    assert np.allclose(out_img.shape, (2 * h, 2 * w, c))
    assert np.allclose(out_poly, gt_poly)


@mock.patch('%s.transforms.np.random.randint' % __name__)
def test_random_crop_flip(mock_randint):
    img = np.ones((10, 10, 3), dtype=np.uint8)
    img[0, 0, :] = 0
    results = {'img': img, 'img_shape': img.shape}

    polygon = np.array([0., 0., 0., 10., 10., 10., 10., 0.])

    results['gt_masks'] = PolygonMasks([[polygon]], *(img.shape[:2]))
    results['gt_masks_ignore'] = PolygonMasks([], *(img.shape[:2]))
    results['mask_fields'] = ['gt_masks', 'gt_masks_ignore']

    crop_ratio = 1.1
    iter_num = 3
    random_crop_fliper = transforms.RandomCropFlip(
        crop_ratio=crop_ratio, iter_num=iter_num)

    # test crop_target
    pad_ratio = 0.1
    h, w = img.shape[:2]
    pad_h = int(h * pad_ratio)
    pad_w = int(w * pad_ratio)
    all_polys = results['gt_masks'].masks
    h_axis, w_axis = random_crop_fliper.generate_crop_target(
        img, all_polys, pad_h, pad_w)

    assert np.allclose(h_axis, (0, 11))
    assert np.allclose(w_axis, (0, 11))

    # test __call__
    polygon = np.array([1., 1., 1., 9., 9., 9., 9., 1.])
    results['gt_masks'] = PolygonMasks([[polygon]], *(img.shape[:2]))
    results['gt_masks_ignore'] = PolygonMasks([[polygon]], *(img.shape[:2]))

    mock_randint.side_effect = [0, 1, 2]
    results = random_crop_fliper(results)

    out_img = results['img']
    out_poly = results['gt_masks'].masks[0][0]
    gt_img = img
    gt_poly = polygon

    assert np.allclose(out_img, gt_img)
    assert np.allclose(out_poly, gt_poly)


@mock.patch('%s.transforms.np.random.random_sample' % __name__)
@mock.patch('%s.transforms.np.random.randint' % __name__)
def test_random_crop_poly_instances(mock_randint, mock_sample):
    results = {}
    img = np.zeros((30, 30, 3))
    poly_masks = PolygonMasks([[
        np.array([5., 5., 25., 5., 25., 10., 5., 10.])
    ], [np.array([5., 20., 25., 20., 25., 25., 5., 25.])]], 30, 30)
    results['img'] = img
    results['gt_masks'] = poly_masks
    results['gt_masks_ignore'] = PolygonMasks([], 30, 30)
    results['mask_fields'] = ['gt_masks', 'gt_masks_ignore']
    results['gt_labels'] = [1, 1]
    rcpi = transforms.RandomCropPolyInstances(
        instance_key='gt_masks', crop_ratio=1.0, min_side_ratio=0.3)

    # test sample_crop_box(img_size, results)
    mock_randint.side_effect = [0, 0, 0, 0, 30, 0, 0, 0, 15]
    crop_box = rcpi.sample_crop_box((30, 30), results)
    assert np.allclose(np.array(crop_box), np.array([0, 0, 30, 15]))

    # test __call__
    mock_randint.side_effect = [0, 0, 0, 0, 30, 0, 15, 0, 30]
    mock_sample.side_effect = [0.1]
    output = rcpi(results)
    target = np.array([5., 5., 25., 5., 25., 10., 5., 10.])
    assert len(output['gt_masks']) == 1
    assert len(output['gt_masks_ignore']) == 0
    assert np.allclose(output['gt_masks'].masks[0][0], target)
    assert output['img'].shape == (15, 30, 3)

    # test __call__ with blank instace_key masks
    mock_randint.side_effect = [0, 0, 0, 0, 30, 0, 15, 0, 30]
    mock_sample.side_effect = [0.1]
    rcpi = transforms.RandomCropPolyInstances(
        instance_key='gt_masks_ignore', crop_ratio=1.0, min_side_ratio=0.3)
    results['img'] = img
    results['gt_masks'] = poly_masks
    output = rcpi(results)
    assert len(output['gt_masks']) == 2
    assert np.allclose(output['gt_masks'].masks[0][0], poly_masks.masks[0][0])
    assert np.allclose(output['gt_masks'].masks[1][0], poly_masks.masks[1][0])
    assert output['img'].shape == (30, 30, 3)


@mock.patch('%s.transforms.np.random.random_sample' % __name__)
def test_random_rotate_poly_instances(mock_sample):
    results = {}
    img = np.zeros((30, 30, 3))
    poly_masks = PolygonMasks(
        [[np.array([10., 10., 20., 10., 20., 20., 10., 20.])]], 30, 30)
    results['img'] = img
    results['gt_masks'] = poly_masks
    results['mask_fields'] = ['gt_masks']
    rrpi = transforms.RandomRotatePolyInstances(rotate_ratio=1.0, max_angle=90)

    mock_sample.side_effect = [0., 1.]
    output = rrpi(results)
    assert np.allclose(output['gt_masks'].masks[0][0],
                       np.array([10., 20., 10., 10., 20., 10., 20., 20.]))
    assert output['img'].shape == (30, 30, 3)


@mock.patch('%s.transforms.np.random.random_sample' % __name__)
def test_square_resize_pad(mock_sample):
    results = {}
    img = np.zeros((15, 30, 3))
    polygon = np.array([10., 5., 20., 5., 20., 10., 10., 10.])
    poly_masks = PolygonMasks([[polygon]], 15, 30)
    results['img'] = img
    results['gt_masks'] = poly_masks
    results['mask_fields'] = ['gt_masks']
    srp = transforms.SquareResizePad(target_size=40, pad_ratio=0.5)

    # test resize with padding
    mock_sample.side_effect = [0.]
    output = srp(results)
    target = 4. / 3 * polygon
    target[1::2] += 10.
    assert np.allclose(output['gt_masks'].masks[0][0], target)
    assert output['img'].shape == (40, 40, 3)

    # test resize to square without padding
    results['img'] = img
    results['gt_masks'] = poly_masks
    mock_sample.side_effect = [1.]
    output = srp(results)
    target = polygon.copy()
    target[::2] *= 4. / 3
    target[1::2] *= 8. / 3
    assert np.allclose(output['gt_masks'].masks[0][0], target)
    assert output['img'].shape == (40, 40, 3)


def test_pyramid_rescale():
    img = np.random.randint(0, 256, size=(128, 100, 3), dtype=np.uint8)
    x = {'img': copy.deepcopy(img)}
    f = transforms.PyramidRescale()
    results = f(x)
    assert results['img'].shape == (128, 100, 3)

    # Test invalid inputs
    with pytest.raises(AssertionError):
        transforms.PyramidRescale(base_shape=(128))
    with pytest.raises(AssertionError):
        transforms.PyramidRescale(base_shape=128)
    with pytest.raises(AssertionError):
        transforms.PyramidRescale(factor=[])
    with pytest.raises(AssertionError):
        transforms.PyramidRescale(randomize_factor=[])
    with pytest.raises(AssertionError):
        f({})

    # Test factor = 0
    f_derandomized = transforms.PyramidRescale(
        factor=0, randomize_factor=False)
    results = f_derandomized({'img': copy.deepcopy(img)})
    assert np.all(results['img'] == img)