# Copyright (c) OpenMMLab. All rights reserved.
import argparse
import math
import os.path as osp

import mmengine

from mmocr.utils import dump_ocr_data


def parse_args():
    parser = argparse.ArgumentParser(
        description='Generate training and validation set of COCO Text v2 ')
    parser.add_argument('root_path', help='Root dir path of COCO Text v2')
    args = parser.parse_args()
    return args


def collect_cocotext_info(root_path, split, print_every=1000):
    """Collect the annotation information.

    The annotation format is as the following:
    {
        'anns':{
            '45346':{
                'mask': [468.9,286.7,468.9,295.2,493.0,295.8,493.0,287.2],
                'class': 'machine printed',
                'bbox': [468.9, 286.7, 24.1, 9.1], # x, y, w, h
                'image_id': 217925,
                'id': 45346,
                'language': 'english', # 'english' or 'not english'
                'area': 206.06,
                'utf8_string': 'New',
                'legibility': 'legible', # 'legible' or 'illegible'
            },
            ...
        }
        'imgs':{
            '540965':{
                'id': 540965,
                'set': 'train', # 'train' or 'val'
                'width': 640,
                'height': 360,
                'file_name': 'COCO_train2014_000000540965.jpg'
            },
            ...
        }
        'imgToAnns':{
            '540965': [],
            '260932': [63993, 63994, 63995, 63996, 63997, 63998, 63999],
            ...
        }
    }

    Args:
        root_path (str): Root path to the dataset
        split (str): Dataset split, which should be 'train' or 'val'
        print_every (int): Print log information per iter

    Returns:
        img_info (dict): The dict of the img and annotation information
    """

    annotation_path = osp.join(root_path, 'annotations/cocotext.v2.json')
    if not osp.exists(annotation_path):
        raise Exception(
            f'{annotation_path} not exists, please check and try again.')

    annotation = mmengine.load(annotation_path)

    img_infos = []
    for i, img_info in enumerate(annotation['imgs'].values()):
        if i > 0 and i % print_every == 0:
            print(f'{i}/{len(annotation["imgs"].values())}')

        if img_info['set'] == split:
            img_info['segm_file'] = annotation_path
            ann_ids = annotation['imgToAnns'][str(img_info['id'])]
            # Filter out images without text
            if len(ann_ids) == 0:
                continue
            anno_info = []
            for ann_id in ann_ids:
                ann = annotation['anns'][str(ann_id)]

                # Ignore illegible or non-English words
                iscrowd = 1 if ann['language'] == 'not english' or ann[
                    'legibility'] == 'illegible' else 0

                x, y, w, h = ann['bbox']
                x, y = max(0, math.floor(x)), max(0, math.floor(y))
                w, h = math.ceil(w), math.ceil(h)
                bbox = [x, y, w, h]
                segmentation = [max(0, int(x)) for x in ann['mask']]
                if len(segmentation) < 8 or len(segmentation) % 2 != 0:
                    segmentation = [x, y, x + w, y, x + w, y + h, x, y + h]
                anno = dict(
                    iscrowd=iscrowd,
                    category_id=1,
                    bbox=bbox,
                    area=ann['area'],
                    segmentation=[segmentation])
                anno_info.append(anno)
            img_info.update(anno_info=anno_info)
            img_infos.append(img_info)
    return img_infos


def main():
    args = parse_args()
    root_path = args.root_path
    print('Processing training set...')
    training_infos = collect_cocotext_info(root_path, 'train')
    dump_ocr_data(training_infos,
                  osp.join(root_path, 'instances_training.json'), 'textdet')
    print('Processing validation set...')
    val_infos = collect_cocotext_info(root_path, 'val')
    dump_ocr_data(val_infos, osp.join(root_path, 'instances_val.json'),
                  'textdet')
    print('Finish')


if __name__ == '__main__':
    main()
