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

import mmengine
import numpy as np

from mmocr.utils import bezier2polygon, sort_points

# The default dictionary used by CurvedSynthText
dict95 = [
    ' ', '!', '"', '#', '$', '%', '&', '\'', '(', ')', '*', '+', ',', '-', '.',
    '/', '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', ':', ';', '<', '=',
    '>', '?', '@', 'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L',
    'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z', '[',
    '\\', ']', '^', '_', '`', 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j',
    'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y',
    'z', '{', '|', '}', '~'
]
UNK = len(dict95)
EOS = UNK + 1


def digit2text(rec):
    res = []
    for d in rec:
        assert d <= EOS
        if d == EOS:
            break
        if d == UNK:
            print('Warning: Has a UNK character')
            res.append('口')  # Or any special character not in the target dict
        res.append(dict95[d])
    return ''.join(res)


def modify_annotation(ann, num_sample, start_img_id=0, start_ann_id=0):
    ann['text'] = digit2text(ann.pop('rec'))
    # Get hide egmentation points
    polygon_pts = bezier2polygon(ann['bezier_pts'], num_sample=num_sample)
    ann['segmentation'] = np.asarray(sort_points(polygon_pts)).reshape(
        1, -1).tolist()
    ann['image_id'] += start_img_id
    ann['id'] += start_ann_id
    return ann


def modify_image_info(image_info, path_prefix, start_img_id=0):
    image_info['file_name'] = osp.join(path_prefix, image_info['file_name'])
    image_info['id'] += start_img_id
    return image_info


def parse_args():
    parser = argparse.ArgumentParser(
        description='Convert CurvedSynText150k to COCO format')
    parser.add_argument('root_path', help='CurvedSynText150k  root path')
    parser.add_argument('-o', '--out-dir', help='Output path')
    parser.add_argument(
        '-n',
        '--num-sample',
        type=int,
        default=4,
        help='Number of sample points at each Bezier curve.')
    parser.add_argument(
        '--nproc', default=1, type=int, help='Number of processes')
    args = parser.parse_args()
    return args


def convert_annotations(data,
                        path_prefix,
                        num_sample,
                        nproc,
                        start_img_id=0,
                        start_ann_id=0):
    modify_image_info_with_params = partial(
        modify_image_info, path_prefix=path_prefix, start_img_id=start_img_id)
    modify_annotation_with_params = partial(
        modify_annotation,
        num_sample=num_sample,
        start_img_id=start_img_id,
        start_ann_id=start_ann_id)
    if nproc > 1:
        data['annotations'] = mmengine.track_parallel_progress(
            modify_annotation_with_params, data['annotations'], nproc=nproc)
        data['images'] = mmengine.track_parallel_progress(
            modify_image_info_with_params, data['images'], nproc=nproc)
    else:
        data['annotations'] = mmengine.track_progress(
            modify_annotation_with_params, data['annotations'])
        data['images'] = mmengine.track_progress(
            modify_image_info_with_params,
            data['images'],
        )
    data['categories'] = [{'id': 1, 'name': 'text'}]
    return data


def main():
    args = parse_args()
    root_path = args.root_path
    out_dir = args.out_dir if args.out_dir else root_path
    mmengine.mkdir_or_exist(out_dir)

    anns = mmengine.load(osp.join(root_path, 'train1.json'))
    data1 = convert_annotations(anns, 'syntext_word_eng', args.num_sample,
                                args.nproc)

    # Get the maximum image id from data1
    start_img_id = max(data1['images'], key=lambda x: x['id'])['id'] + 1
    start_ann_id = max(data1['annotations'], key=lambda x: x['id'])['id'] + 1
    anns = mmengine.load(osp.join(root_path, 'train2.json'))
    data2 = convert_annotations(
        anns,
        'emcs_imgs',
        args.num_sample,
        args.nproc,
        start_img_id=start_img_id,
        start_ann_id=start_ann_id)

    data1['images'] += data2['images']
    data1['annotations'] += data2['annotations']
    mmengine.dump(data1, osp.join(out_dir, 'instances_training.json'))


if __name__ == '__main__':
    main()