# Copyright (c) OpenMMLab. All rights reserved.
import json
import os
import re
import time
import warnings

import cv2
import numpy as np
import xmltodict
from xtcocotools.coco import COCO

np.random.seed(0)


def list_all_files(root_dir, ext='.xml'):
    """List all files in the root directory and all its sub directories.

    :param root_dir: root directory
    :param ext: filename extension
    :return: list of files
    """
    files = []
    file_list = os.listdir(root_dir)
    for i in range(0, len(file_list)):
        path = os.path.join(root_dir, file_list[i])
        if os.path.isdir(path):
            files.extend(list_all_files(path))
        if os.path.isfile(path):
            if path.lower().endswith(ext):
                files.append(path)
    return files


def get_anno_info():
    keypoints_info = [
        'L_Eye',
        'R_Eye',
        'L_EarBase',
        'R_EarBase',
        'Nose',
        'Throat',
        'TailBase',
        'Withers',
        'L_F_Elbow',
        'R_F_Elbow',
        'L_B_Elbow',
        'R_B_Elbow',
        'L_F_Knee',
        'R_F_Knee',
        'L_B_Knee',
        'R_B_Knee',
        'L_F_Paw',
        'R_F_Paw',
        'L_B_Paw',
        'R_B_Paw',
    ]
    skeleton_info = [[1, 2], [1, 3], [2, 4], [1, 5], [2, 5], [5, 6], [6, 8],
                     [7, 8], [6, 9], [9, 13], [13, 17], [6, 10], [10, 14],
                     [14, 18], [7, 11], [11, 15], [15, 19], [7, 12], [12, 16],
                     [16, 20]]
    category_info = [{
        'supercategory': 'animal',
        'id': 1,
        'name': 'animal',
        'keypoints': keypoints_info,
        'skeleton': skeleton_info
    }]

    return keypoints_info, skeleton_info, category_info


def xml2coco_trainval(file_list, img_root, save_path, start_ann_id=0):
    """Save annotations in coco-format.

    :param file_list: list of data annotation files.
    :param img_root: the root dir to load images.
    :param save_path: the path to save transformed annotation file.
    :param start_ann_id: the starting point to count the annotation id.
    :param val_num: the number of annotated objects for validation.
    """
    images = []
    annotations = []
    img_ids = []
    ann_ids = []

    ann_id = start_ann_id

    name2id = {
        'L_Eye': 0,
        'R_Eye': 1,
        'L_EarBase': 2,
        'R_EarBase': 3,
        'Nose': 4,
        'Throat': 5,
        'TailBase': 6,
        'Withers': 7,
        'L_F_Elbow': 8,
        'R_F_Elbow': 9,
        'L_B_Elbow': 10,
        'R_B_Elbow': 11,
        'L_F_Knee': 12,
        'R_F_Knee': 13,
        'L_B_Knee': 14,
        'R_B_Knee': 15,
        'L_F_Paw': 16,
        'R_F_Paw': 17,
        'L_B_Paw': 18,
        'R_B_Paw': 19
    }
    for file in file_list:
        data_anno = xmltodict.parse(open(file).read())['annotation']

        img_id = int(data_anno['image'].split('_')[0] +
                     data_anno['image'].split('_')[1])

        if img_id not in img_ids:
            image_name = 'VOC2012/JPEGImages/' + data_anno['image'] + '.jpg'
            img = cv2.imread(os.path.join(img_root, image_name))

            image = {}
            image['id'] = img_id
            image['file_name'] = image_name
            image['height'] = img.shape[0]
            image['width'] = img.shape[1]

            images.append(image)
            img_ids.append(img_id)
        else:
            pass

        keypoint_anno = data_anno['keypoints']['keypoint']
        assert len(keypoint_anno) == 20

        keypoints = np.zeros([20, 3], dtype=np.float32)

        for kpt_anno in keypoint_anno:
            keypoint_name = kpt_anno['@name']
            keypoint_id = name2id[keypoint_name]

            visibility = int(kpt_anno['@visible'])

            if visibility == 0:
                continue
            else:
                keypoints[keypoint_id, 0] = float(kpt_anno['@x'])
                keypoints[keypoint_id, 1] = float(kpt_anno['@y'])
                keypoints[keypoint_id, 2] = 2

        anno = {}
        anno['keypoints'] = keypoints.reshape(-1).tolist()
        anno['image_id'] = img_id
        anno['id'] = ann_id
        anno['num_keypoints'] = int(sum(keypoints[:, 2] > 0))

        visible_bounds = data_anno['visible_bounds']
        anno['bbox'] = [
            float(visible_bounds['@xmin']),
            float(visible_bounds['@ymin']),
            float(visible_bounds['@width']),
            float(visible_bounds['@height'])
        ]
        anno['iscrowd'] = 0
        anno['area'] = float(anno['bbox'][2] * anno['bbox'][3])
        anno['category_id'] = 1

        annotations.append(anno)
        ann_ids.append(ann_id)
        ann_id += 1

    cocotype = {}

    cocotype['info'] = {}
    cocotype['info'][
        'description'] = 'AnimalPose dataset Generated by MMPose Team'
    cocotype['info']['version'] = '1.0'
    cocotype['info']['year'] = time.strftime('%Y', time.localtime())
    cocotype['info']['date_created'] = time.strftime('%Y/%m/%d',
                                                     time.localtime())

    cocotype['images'] = images
    cocotype['annotations'] = annotations

    keypoints_info, skeleton_info, category_info = get_anno_info()

    cocotype['categories'] = category_info

    os.makedirs(os.path.dirname(save_path), exist_ok=True)
    json.dump(cocotype, open(save_path, 'w'), indent=4)
    print('number of images:', len(img_ids))
    print('number of annotations:', len(ann_ids))
    print(f'done {save_path}')


def xml2coco_test(file_list, img_root, save_path, start_ann_id=0):
    """Save annotations in coco-format.

    :param file_list: list of data annotation files.
    :param img_root: the root dir to load images.
    :param save_path: the path to save transformed annotation file.
    :param start_ann_id: the starting point to count the annotation id.
    """
    images = []
    annotations = []
    img_ids = []
    ann_ids = []

    ann_id = start_ann_id

    name2id = {
        'L_eye': 0,
        'R_eye': 1,
        'L_ear': 2,
        'R_ear': 3,
        'Nose': 4,
        'Throat': 5,
        'Tail': 6,
        'withers': 7,
        'L_F_elbow': 8,
        'R_F_elbow': 9,
        'L_B_elbow': 10,
        'R_B_elbow': 11,
        'L_F_knee': 12,
        'R_F_knee': 13,
        'L_B_knee': 14,
        'R_B_knee': 15,
        'L_F_paw': 16,
        'R_F_paw': 17,
        'L_B_paw': 18,
        'R_B_paw': 19
    }

    cat2id = {'cat': 1, 'cow': 2, 'dog': 3, 'horse': 4, 'sheep': 5}

    for file in file_list:
        data_anno = xmltodict.parse(open(file).read())['annotation']

        category_id = cat2id[data_anno['category']]

        img_id = category_id * 1000 + int(
            re.findall(r'\d+', data_anno['image'])[0])

        assert img_id not in img_ids

        # prepare images
        image_name = os.path.join('animalpose_image_part2',
                                  data_anno['category'], data_anno['image'])
        img = cv2.imread(os.path.join(img_root, image_name))

        image = {}
        image['id'] = img_id
        image['file_name'] = image_name
        image['height'] = img.shape[0]
        image['width'] = img.shape[1]

        images.append(image)
        img_ids.append(img_id)

        # prepare annotations
        keypoint_anno = data_anno['keypoints']['keypoint']
        keypoints = np.zeros([20, 3], dtype=np.float32)

        for kpt_anno in keypoint_anno:
            keypoint_name = kpt_anno['@name']
            keypoint_id = name2id[keypoint_name]

            visibility = int(kpt_anno['@visible'])

            if visibility == 0:
                continue
            else:
                keypoints[keypoint_id, 0] = float(kpt_anno['@x'])
                keypoints[keypoint_id, 1] = float(kpt_anno['@y'])
                keypoints[keypoint_id, 2] = 2

        anno = {}
        anno['keypoints'] = keypoints.reshape(-1).tolist()
        anno['image_id'] = img_id
        anno['id'] = ann_id
        anno['num_keypoints'] = int(sum(keypoints[:, 2] > 0))

        visible_bounds = data_anno['visible_bounds']
        anno['bbox'] = [
            float(visible_bounds['@xmin']),
            float(visible_bounds['@xmax']
                  ),  # typo in original xml: should be 'ymin'
            float(visible_bounds['@width']),
            float(visible_bounds['@height'])
        ]
        anno['iscrowd'] = 0
        anno['area'] = float(anno['bbox'][2] * anno['bbox'][3])
        anno['category_id'] = 1

        annotations.append(anno)
        ann_ids.append(ann_id)
        ann_id += 1

    cocotype = {}

    cocotype['info'] = {}
    cocotype['info'][
        'description'] = 'AnimalPose dataset Generated by MMPose Team'
    cocotype['info']['version'] = '1.0'
    cocotype['info']['year'] = time.strftime('%Y', time.localtime())
    cocotype['info']['date_created'] = time.strftime('%Y/%m/%d',
                                                     time.localtime())

    cocotype['images'] = images
    cocotype['annotations'] = annotations

    keypoints_info, skeleton_info, category_info = get_anno_info()

    cocotype['categories'] = category_info

    os.makedirs(os.path.dirname(save_path), exist_ok=True)
    json.dump(cocotype, open(save_path, 'w'), indent=4)
    print('=========================================================')
    print('number of images:', len(img_ids))
    print('number of annotations:', len(ann_ids))
    print(f'done {save_path}')


def split_train_val(work_dir, trainval_file, train_file, val_file,
                    val_ann_num):
    """Split train-val json file into training and validation files.

    :param work_dir: path to load train-val json file, and save split files.
    :param trainval_file: The input json file combining both train and val.
    :param trainval_file: The output json file for training.
    :param trainval_file: The output json file for validation.
    :param val_ann_num: the number of validation annotations.
    """

    coco = COCO(os.path.join(work_dir, trainval_file))

    img_list = list(coco.imgs.keys())
    np.random.shuffle(img_list)

    count = 0

    images_train = []
    images_val = []
    annotations_train = []
    annotations_val = []

    for img_id in img_list:
        ann_ids = coco.getAnnIds(img_id)

        if count + len(ann_ids) <= val_ann_num:
            # for validation
            count += len(ann_ids)
            images_val.append(coco.imgs[img_id])
            for ann_id in ann_ids:
                annotations_val.append(coco.anns[ann_id])

        else:
            images_train.append(coco.imgs[img_id])
            for ann_id in ann_ids:
                annotations_train.append(coco.anns[ann_id])

    if count == val_ann_num:
        print(f'We have found {count} annotations for validation.')
    else:
        warnings.warn(
            f'We only found {count} annotations, instead of {val_ann_num}.')

    cocotype_train = {}
    cocotype_val = {}

    keypoints_info, skeleton_info, category_info = get_anno_info()

    cocotype_train['info'] = {}
    cocotype_train['info'][
        'description'] = 'AnimalPose dataset Generated by MMPose Team'
    cocotype_train['info']['version'] = '1.0'
    cocotype_train['info']['year'] = time.strftime('%Y', time.localtime())
    cocotype_train['info']['date_created'] = time.strftime(
        '%Y/%m/%d', time.localtime())
    cocotype_train['images'] = images_train
    cocotype_train['annotations'] = annotations_train
    cocotype_train['categories'] = category_info

    json.dump(
        cocotype_train,
        open(os.path.join(work_dir, train_file), 'w'),
        indent=4)
    print('=========================================================')
    print('number of images:', len(images_train))
    print('number of annotations:', len(annotations_train))
    print(f'done {train_file}')

    cocotype_val['info'] = {}
    cocotype_val['info'][
        'description'] = 'AnimalPose dataset Generated by MMPose Team'
    cocotype_val['info']['version'] = '1.0'
    cocotype_val['info']['year'] = time.strftime('%Y', time.localtime())
    cocotype_val['info']['date_created'] = time.strftime(
        '%Y/%m/%d', time.localtime())
    cocotype_val['images'] = images_val
    cocotype_val['annotations'] = annotations_val
    cocotype_val['categories'] = category_info

    json.dump(
        cocotype_val, open(os.path.join(work_dir, val_file), 'w'), indent=4)
    print('=========================================================')
    print('number of images:', len(images_val))
    print('number of annotations:', len(annotations_val))
    print(f'done {val_file}')


dataset_dir = 'data/animalpose/'

# We choose the images from PascalVOC for train + val
# In total, train+val: 3608 images, 5117 annotations
xml2coco_trainval(
    list_all_files(os.path.join(dataset_dir, 'PASCAL2011_animal_annotation')),
    dataset_dir,
    os.path.join(dataset_dir, 'annotations', 'animalpose_trainval.json'),
    start_ann_id=1000000)

# train: 2798 images, 4000 annotations
# val: 810 images, 1117 annotations
split_train_val(
    os.path.join(dataset_dir, 'annotations'),
    'animalpose_trainval.json',
    'animalpose_train.json',
    'animalpose_val.json',
    val_ann_num=1117)

# We choose the remaining 1000 images for test
# 1000 images, 1000 annotations
xml2coco_test(
    list_all_files(os.path.join(dataset_dir, 'animalpose_anno2')),
    dataset_dir,
    os.path.join(dataset_dir, 'annotations', 'animalpose_test.json'),
    start_ann_id=0)