# Copyright (c) OpenMMLab. All rights reserved.
import unittest
from itertools import chain, permutations

import numpy as np
import torch
from shapely.geometry import MultiPolygon, Polygon

from mmocr.utils import (boundary_iou, crop_polygon, offset_polygon, poly2bbox,
                         poly2shapely, poly_intersection, poly_iou,
                         poly_make_valid, poly_union, polys2shapely,
                         rescale_polygon, rescale_polygons, shapely2poly,
                         sort_points, sort_vertex, sort_vertex8)


class TestPolygonUtils(unittest.TestCase):

    def test_crop_polygon(self):
        # polygon cross box
        polygon = np.array([20., -10., 40., 10., 10., 40., -10., 20.])
        crop_box = np.array([0., 0., 60., 60.])
        target_poly_cropped = np.array(
            [10, 40, 0, 30, 0, 10, 10, 0, 30, 0, 40, 10])
        poly_cropped = crop_polygon(polygon, crop_box)
        self.assertTrue(
            poly2shapely(poly_cropped).equals(
                poly2shapely(target_poly_cropped)))

        # polygon inside box
        polygon = np.array([0., 0., 30., 0., 30., 30., 0., 30.])
        crop_box = np.array([0., 0., 60., 60.])
        target_poly_cropped = polygon
        poly_cropped = crop_polygon(polygon, crop_box)
        self.assertTrue(
            poly2shapely(poly_cropped).equals(
                poly2shapely(target_poly_cropped)))

        # polygon outside box
        polygon = np.array([0., 0., 30., 0., 30., 30., 0., 30.])
        crop_box = np.array([80., 80., 90., 90.])
        poly_cropped = crop_polygon(polygon, crop_box)
        self.assertEqual(poly_cropped, None)

        # polygon and box are overlapped at a point
        polygon = np.array([0., 0., 10., 0., 10., 10., 0., 10.])
        crop_box = np.array([10., 10., 20., 20.])
        poly_cropped = crop_polygon(polygon, crop_box)
        self.assertEqual(poly_cropped, None)

    def test_rescale_polygon(self):
        scale_factor = (0.3, 0.4)

        with self.assertRaises(AssertionError):
            polygons = [0, 0, 1, 0, 1, 1, 0]
            rescale_polygon(polygons, scale_factor)

        polygons = [0, 0, 1, 0, 1, 1, 0, 1]
        self.assertTrue(
            np.allclose(
                rescale_polygon(polygons, scale_factor, mode='div'),
                np.array([0, 0, 1 / 0.3, 0, 1 / 0.3, 1 / 0.4, 0, 1 / 0.4])))
        self.assertTrue(
            np.allclose(
                rescale_polygon(polygons, scale_factor, mode='mul'),
                np.array([0, 0, 0.3, 0, 0.3, 0.4, 0, 0.4])))

    def test_rescale_polygons(self):
        polygons = [
            np.array([0, 0, 1, 0, 1, 1, 0, 1]),
            np.array([1, 1, 2, 1, 2, 2, 1, 2])
        ]
        scale_factor = (0.5, 0.5)
        self.assertTrue(
            np.allclose(
                rescale_polygons(polygons, scale_factor, mode='div'), [
                    np.array([0, 0, 2, 0, 2, 2, 0, 2]),
                    np.array([2, 2, 4, 2, 4, 4, 2, 4])
                ]))
        self.assertTrue(
            np.allclose(
                rescale_polygons(polygons, scale_factor, mode='mul'), [
                    np.array([0, 0, 0.5, 0, 0.5, 0.5, 0, 0.5]),
                    np.array([0.5, 0.5, 1, 0.5, 1, 1, 0.5, 1])
                ]))

        polygons = np.array([[0, 0, 1, 0, 1, 1, 0, 1],
                             [1, 1, 2, 1, 2, 2, 1, 2]])
        scale_factor = (0.5, 0.5)
        self.assertTrue(
            np.allclose(
                rescale_polygons(polygons, scale_factor, mode='div'),
                np.array([[0, 0, 2, 0, 2, 2, 0, 2], [2, 2, 4, 2, 4, 4, 2,
                                                     4]])))
        self.assertTrue(
            np.allclose(
                rescale_polygons(polygons, scale_factor, mode='mul'),
                np.array([[0, 0, 0.5, 0, 0.5, 0.5, 0, 0.5],
                          [0.5, 0.5, 1, 0.5, 1, 1, 0.5, 1]])))

        polygons = [torch.Tensor([0, 0, 1, 0, 1, 1, 0, 1])]
        scale_factor = (0.3, 0.4)
        self.assertTrue(
            np.allclose(
                rescale_polygons(polygons, scale_factor, mode='div'),
                [np.array([0, 0, 1 / 0.3, 0, 1 / 0.3, 1 / 0.4, 0, 1 / 0.4])]))
        self.assertTrue(
            np.allclose(
                rescale_polygons(polygons, scale_factor, mode='mul'),
                [np.array([0, 0, 0.3, 0, 0.3, 0.4, 0, 0.4])]))

    def test_poly2bbox(self):
        # test np.array
        polygon = np.array([0, 0, 1, 0, 1, 1, 0, 1])
        self.assertTrue(np.all(poly2bbox(polygon) == np.array([0, 0, 1, 1])))
        # test list
        polygon = [0, 0, 1, 0, 1, 1, 0, 1]
        self.assertTrue(np.all(poly2bbox(polygon) == np.array([0, 0, 1, 1])))
        # test tensor
        polygon = torch.Tensor([0, 0, 1, 0, 1, 1, 0, 1])
        self.assertTrue(np.all(poly2bbox(polygon) == np.array([0, 0, 1, 1])))

    def test_poly2shapely(self):
        polygon = Polygon([[0, 0], [1, 0], [1, 1], [0, 1]])
        # test np.array
        poly = np.array([0, 0, 1, 0, 1, 1, 0, 1])
        self.assertEqual(poly2shapely(poly), polygon)
        # test list
        poly = [0, 0, 1, 0, 1, 1, 0, 1]
        self.assertEqual(poly2shapely(poly), polygon)
        # test tensor
        poly = torch.Tensor([0, 0, 1, 0, 1, 1, 0, 1])
        self.assertEqual(poly2shapely(poly), polygon)
        # test invalid
        poly = [0, 0, 1]
        with self.assertRaises(AssertionError):
            poly2shapely(poly)
        poly = [0, 0, 1, 0, 1, 1, 0, 1, 1]
        with self.assertRaises(AssertionError):
            poly2shapely(poly)

    def test_polys2shapely(self):
        polygons = [
            Polygon([[0, 0], [1, 0], [1, 1], [0, 1]]),
            Polygon([[1, 0], [1, 1], [0, 1], [0, 0]])
        ]
        # test np.array
        polys = np.array([[0, 0, 1, 0, 1, 1, 0, 1], [1, 0, 1, 1, 0, 1, 0, 0]])
        self.assertEqual(polys2shapely(polys), polygons)
        # test list
        polys = [[0, 0, 1, 0, 1, 1, 0, 1], [1, 0, 1, 1, 0, 1, 0, 0]]
        self.assertEqual(polys2shapely(polys), polygons)
        # test tensor
        polys = torch.Tensor([[0, 0, 1, 0, 1, 1, 0, 1],
                              [1, 0, 1, 1, 0, 1, 0, 0]])
        self.assertEqual(polys2shapely(polys), polygons)
        # test invalid
        polys = [0, 0, 1]
        with self.assertRaises(AssertionError):
            polys2shapely(polys)
        polys = [0, 0, 1, 0, 1, 1, 0, 1, 1]
        with self.assertRaises(AssertionError):
            polys2shapely(polys)

    def test_shapely2poly(self):
        polygon = Polygon([[0., 0.], [1., 0.], [1., 1.], [0., 1.]])
        poly = np.array([0., 0., 1., 0., 1., 1., 0., 1., 0., 0.])
        self.assertTrue(poly2shapely(poly).equals(polygon))
        self.assertTrue(isinstance(shapely2poly(polygon), np.ndarray))

    def test_poly_make_valid(self):
        poly = Polygon([[0, 0], [1, 1], [1, 0], [0, 1]])
        self.assertFalse(poly.is_valid)
        poly = poly_make_valid(poly)
        self.assertTrue(poly.is_valid)
        # invalid input
        with self.assertRaises(AssertionError):
            poly_make_valid([0, 0, 1, 1, 1, 0, 0, 1])
        poly = Polygon([[337, 441], [326, 386], [334, 397], [342, 412],
                        [296, 382], [317, 366], [324, 427], [315, 413],
                        [308, 400], [349, 419], [337, 441]])
        self.assertFalse(poly.is_valid)
        poly = poly_make_valid(poly)
        self.assertTrue(poly.is_valid)

    def test_poly_intersection(self):

        # test unsupported type
        with self.assertRaises(AssertionError):
            poly_intersection(0, 1)

        # test non-overlapping polygons
        points = [0, 0, 0, 1, 1, 1, 1, 0]
        points1 = [10, 20, 30, 40, 50, 60, 70, 80]
        points2 = [0, 0, 0, 0, 0, 0, 0, 0]  # Invalid polygon
        points3 = [0, 0, 0, 1, 1, 0, 1, 1]  # Self-intersected polygon
        points4 = [0.5, 0, 1.5, 0, 1.5, 1, 0.5, 1]
        poly = poly2shapely(points)
        poly1 = poly2shapely(points1)
        poly2 = poly2shapely(points2)
        poly3 = poly2shapely(points3)
        poly4 = poly2shapely(points4)

        area_inters = poly_intersection(poly, poly1)
        self.assertEqual(area_inters, 0.)

        # test overlapping polygons
        area_inters = poly_intersection(poly, poly)
        self.assertEqual(area_inters, 1)
        area_inters = poly_intersection(poly, poly4)
        self.assertEqual(area_inters, 0.5)

        # test invalid polygons
        self.assertEqual(poly_intersection(poly2, poly2), 0)
        self.assertEqual(poly_intersection(poly3, poly3, invalid_ret=1), 1)
        self.assertEqual(
            poly_intersection(poly3, poly3, invalid_ret=None), 0.25)

        # test poly return
        _, poly = poly_intersection(poly, poly4, return_poly=True)
        self.assertTrue(isinstance(poly, Polygon))
        _, poly = poly_intersection(
            poly3, poly3, invalid_ret=None, return_poly=True)
        self.assertTrue(isinstance(poly, Polygon))
        _, poly = poly_intersection(
            poly2, poly3, invalid_ret=1, return_poly=True)
        self.assertTrue(poly is None)

    def test_poly_union(self):

        # test unsupported type
        with self.assertRaises(AssertionError):
            poly_union(0, 1)

        # test non-overlapping polygons

        points = [0, 0, 0, 1, 1, 1, 1, 0]
        points1 = [2, 2, 2, 3, 3, 3, 3, 2]
        points2 = [0, 0, 0, 0, 0, 0, 0, 0]  # Invalid polygon
        points3 = [0, 0, 0, 1, 1, 0, 1, 1]  # Self-intersected polygon
        points4 = [0.5, 0.5, 1, 0, 1, 1, 0.5, 0.5]
        poly = poly2shapely(points)
        poly1 = poly2shapely(points1)
        poly2 = poly2shapely(points2)
        poly3 = poly2shapely(points3)
        poly4 = poly2shapely(points4)

        assert poly_union(poly, poly1) == 2

        # test overlapping polygons
        assert poly_union(poly, poly) == 1

        # test invalid polygons
        self.assertEqual(poly_union(poly2, poly2), 0)
        self.assertEqual(poly_union(poly3, poly3, invalid_ret=1), 1)

        # The return value depends on the implementation of the package
        self.assertEqual(poly_union(poly3, poly3, invalid_ret=None), 0.25)
        self.assertEqual(poly_union(poly2, poly3), 0.25)
        self.assertEqual(poly_union(poly3, poly4), 0.5)

        # test poly return
        _, poly = poly_union(poly, poly1, return_poly=True)
        self.assertTrue(isinstance(poly, MultiPolygon))
        _, poly = poly_union(poly3, poly3, return_poly=True)
        self.assertTrue(isinstance(poly, Polygon))
        _, poly = poly_union(poly2, poly3, invalid_ret=0, return_poly=True)
        self.assertTrue(poly is None)

    def test_poly_iou(self):
        # test unsupported type
        with self.assertRaises(AssertionError):
            poly_iou([1], [2])

        points = [0, 0, 0, 1, 1, 1, 1, 0]
        points1 = [10, 20, 30, 40, 50, 60, 70, 80]
        points2 = [0, 0, 0, 0, 0, 0, 0, 0]  # Invalid polygon
        points3 = [0, 0, 0, 1, 1, 0, 1, 1]  # Self-intersected polygon

        poly = poly2shapely(points)
        poly1 = poly2shapely(points1)
        poly2 = poly2shapely(points2)
        poly3 = poly2shapely(points3)

        self.assertEqual(poly_iou(poly, poly1), 0)

        # test overlapping polygons
        self.assertEqual(poly_iou(poly, poly), 1)

        # test invalid polygons
        self.assertEqual(poly_iou(poly2, poly2), 0)
        self.assertEqual(poly_iou(poly3, poly3, zero_division=1), 1)
        self.assertEqual(poly_iou(poly2, poly3), 0)

    def test_offset_polygon(self):
        # usual case
        polygons = np.array([0, 0, 0, 1, 1, 1, 1, 0], dtype=np.float32)
        expanded_polygon = offset_polygon(polygons, 1)
        self.assertTrue(
            poly2shapely(expanded_polygon).equals(
                poly2shapely(
                    np.array(
                        [2, 0, 2, 1, 1, 2, 0, 2, -1, 1, -1, 0, 0, -1, 1,
                         -1]))))

        # Overshrunk polygon doesn't exist
        shrunk_polygon = offset_polygon(polygons, -10)
        self.assertEqual(len(shrunk_polygon), 0)

        # When polygon is shrunk into two polygons, it is regarded as invalid
        # and an empty array is returned.
        polygons = np.array([0, 0, 0, 3, 1, 2, 2, 3, 2, 0, 1, 1],
                            dtype=np.float32)
        shrunk = offset_polygon(polygons, -1)
        self.assertEqual(len(shrunk), 0)

    def test_boundary_iou(self):
        points = [0, 0, 0, 1, 1, 1, 1, 0]
        points1 = [10, 20, 30, 40, 50, 60, 70, 80]
        points2 = [0, 0, 0, 0, 0, 0, 0, 0]  # Invalid polygon
        points3 = [0, 0, 0, 1, 1, 0, 1, 1]  # Self-intersected polygon

        self.assertEqual(boundary_iou(points, points1), 0)

        # test overlapping boundaries
        self.assertEqual(boundary_iou(points, points), 1)

        # test invalid boundaries
        self.assertEqual(boundary_iou(points2, points2), 0)
        self.assertEqual(boundary_iou(points3, points3, zero_division=1), 1)
        self.assertEqual(boundary_iou(points2, points3), 0)

    def test_sort_points(self):
        points = np.array([[1, 1], [0, 0], [1, -1], [2, -2], [0, 2], [1, 1],
                           [0, 1], [-1, 1], [-1, -1]])
        target = np.array([[-1, -1], [0, 0], [-1, 1], [0, 1], [0, 2], [1, 1],
                           [1, 1], [2, -2], [1, -1]])
        self.assertTrue(np.allclose(target, sort_points(points)))

        points = np.array([[1, 1], [1, -1], [-1, 1], [-1, -1]])
        target = np.array([[-1, -1], [-1, 1], [1, 1], [1, -1]])
        self.assertTrue(np.allclose(target, sort_points(points)))

        points = [[1, 1], [1, -1], [-1, 1], [-1, -1]]
        self.assertTrue(np.allclose(target, sort_points(points)))

        points = [[0.5, 0.3], [1, 0.5], [-0.5, 0.8], [-0.1, 1]]
        target = [[-0.5, 0.8], [-0.1, 1], [1, 0.5], [0.5, 0.3]]
        self.assertTrue(np.allclose(target, sort_points(points)))

        points = [[0.5, 3], [0.1, -0.2], [-0.5, -0.3], [-0.7, 3.1]]
        target = [[-0.5, -0.3], [-0.7, 3.1], [0.5, 3], [0.1, -0.2]]
        self.assertTrue(np.allclose(target, sort_points(points)))

        points = [[1, 0.8], [0.8, -1], [1.8, 0.5], [1.9, -0.6], [-0.5, 2],
                  [-1, 1.8], [-2, 0.7], [-1.6, -0.2], [-1, -0.5]]
        target = [[-1, -0.5], [-1.6, -0.2], [-2, 0.7], [-1, 1.8], [-0.5, 2],
                  [1, 0.8], [1.8, 0.5], [1.9, -0.6], [0.8, -1]]
        self.assertTrue(np.allclose(target, sort_points(points)))

        # concave polygon may failed
        points = [[1, 0], [-1, 0], [0, 0], [0, -1], [0.25, 1], [0.75, 1],
                  [-0.25, 1], [-0.75, 1]]
        target = [[-1, 0], [-0.75, 1], [-0.25, 1], [0, 0], [0.25, 1],
                  [0.75, 1], [1, 0], [0, -1]]
        self.assertFalse(np.allclose(target, sort_points(points)))

        with self.assertRaises(AssertionError):
            sort_points([1, 2])

    def test_sort_vertex(self):
        dummy_points_x = [20, 20, 120, 120]
        dummy_points_y = [20, 40, 40, 20]

        expect_points_x = [20, 120, 120, 20]
        expect_points_y = [20, 20, 40, 40]

        with self.assertRaises(AssertionError):
            sort_vertex([], dummy_points_y)
        with self.assertRaises(AssertionError):
            sort_vertex(dummy_points_x, [])

        for perm in set(permutations([0, 1, 2, 3])):
            points_x = [dummy_points_x[i] for i in perm]
            points_y = [dummy_points_y[i] for i in perm]
            ordered_points_x, ordered_points_y = sort_vertex(
                points_x, points_y)

            self.assertTrue(np.allclose(ordered_points_x, expect_points_x))
            self.assertTrue(np.allclose(ordered_points_y, expect_points_y))

    def test_sort_vertex8(self):
        dummy_points_x = [21, 21, 122, 122]
        dummy_points_y = [21, 39, 39, 21]

        expect_points = [21, 21, 122, 21, 122, 39, 21, 39]

        for perm in set(permutations([0, 1, 2, 3])):
            points_x = [dummy_points_x[i] for i in perm]
            points_y = [dummy_points_y[i] for i in perm]
            points = list(chain.from_iterable(zip(points_x, points_y)))
            ordered_points = sort_vertex8(points)

            self.assertTrue(np.allclose(ordered_points, expect_points))