# Copyright (c) OpenMMLab. All rights reserved.
import os.path as osp
import tempfile
from unittest import TestCase

import mmengine

from mmocr.datasets.icdar_dataset import IcdarDataset


class TestIcdarDataset(TestCase):

    def _create_dummy_icdar_json(self, json_name):
        image_1 = {
            'id': 0,
            'width': 640,
            'height': 640,
            'file_name': 'fake_name.jpg',
        }
        image_2 = {
            'id': 1,
            'width': 640,
            'height': 640,
            'file_name': 'fake_name1.jpg',
        }

        annotation_1 = {
            'id': 1,
            'image_id': 0,
            'category_id': 0,
            'area': 400,
            'bbox': [50, 60, 20, 20],
            'iscrowd': 0,
            'segmentation': [[50, 60, 70, 60, 70, 80, 50, 80]]
        }

        annotation_2 = {
            'id': 2,
            'image_id': 0,
            'category_id': 0,
            'area': 900,
            'bbox': [100, 120, 30, 30],
            'iscrowd': 0,
            'segmentation': [[100, 120, 130, 120, 120, 150, 100, 150]]
        }

        annotation_3 = {
            'id': 3,
            'image_id': 0,
            'category_id': 0,
            'area': 1600,
            'bbox': [150, 160, 40, 40],
            'iscrowd': 1,
            'segmentation': [[150, 160, 190, 160, 190, 200, 150, 200]]
        }

        annotation_4 = {
            'id': 4,
            'image_id': 0,
            'category_id': 0,
            'area': 10000,
            'bbox': [250, 260, 100, 100],
            'iscrowd': 1,
            'segmentation': [[250, 260, 350, 260, 350, 360, 250, 360]]
        }
        annotation_5 = {
            'id': 5,
            'image_id': 1,
            'category_id': 0,
            'area': 10000,
            'bbox': [250, 260, 100, 100],
            'iscrowd': 1,
            'segmentation': [[250, 260, 350, 260, 350, 360, 250, 360]]
        }
        annotation_6 = {
            'id': 6,
            'image_id': 1,
            'category_id': 0,
            'area': 0,
            'bbox': [0, 0, 0, 0],
            'iscrowd': 1,
            'segmentation': [[250, 260, 350, 260, 350, 360, 250, 360]]
        }
        annotation_7 = {
            'id': 7,
            'image_id': 1,
            'category_id': 2,
            'area': 10000,
            'bbox': [250, 260, 100, 100],
            'iscrowd': 1,
            'segmentation': [[250, 260, 350, 260, 350, 360, 250, 360]]
        }
        annotation_8 = {
            'id': 8,
            'image_id': 1,
            'category_id': 0,
            'area': 10000,
            'bbox': [250, 260, 100, 100],
            'iscrowd': 1,
            'segmentation': [[250, 260, 350, 260, 350, 360, 250, 360]]
        }

        categories = [{
            'id': 0,
            'name': 'text',
            'supercategory': 'text',
        }]

        fake_json = {
            'images': [image_1, image_2],
            'annotations': [
                annotation_1, annotation_2, annotation_3, annotation_4,
                annotation_5, annotation_6, annotation_7, annotation_8
            ],
            'categories':
            categories
        }
        self.metainfo = dict(classes=('text'))
        mmengine.dump(fake_json, json_name)

    def test_icdar_dataset(self):
        tmp_dir = tempfile.TemporaryDirectory()
        # create dummy data
        fake_json_file = osp.join(tmp_dir.name, 'fake_data.json')
        self._create_dummy_icdar_json(fake_json_file)

        # test initialization
        dataset = IcdarDataset(
            ann_file=fake_json_file,
            data_prefix=dict(img_path='imgs'),
            metainfo=self.metainfo,
            pipeline=[])
        self.assertEqual(dataset.metainfo['classes'], self.metainfo['classes'])
        dataset.full_init()
        self.assertEqual(len(dataset), 2)
        self.assertEqual(len(dataset.load_data_list()), 2)

        # test load_data_list
        anno = dataset.load_data_list()[0]
        self.assertEqual(len(anno['instances']), 4)
        self.assertTrue('ignore' in anno['instances'][0])
        self.assertTrue('bbox' in anno['instances'][0])
        self.assertEqual(anno['instances'][0]['bbox_label'], 0)
        self.assertTrue('polygon' in anno['instances'][0])
        tmp_dir.cleanup()