import os
import scipy.io as scio
import numpy as np
import cv2
import shutil
import glob
import json


def _mkdir(_dir):
    if not os.path.exists(_dir):
        os.mkdir(_dir)


if __name__ == '__main__':
    _categories = {0: 'background', 1: 'Lower-arm', 2: 'Head', 3: 'Upper-leg', 4: 'Torso', 5: 'Lower-leg',
                   6: 'Upper-arm'}

    key = 'test'
    datasets = {'train': ['train', 'Training'], 'val': ['val', 'Validation'], 'test': ['test', 'Testing']}
    key_list = datasets[key]

    ids_file = 'pascal-person-part/tools/{}_id.txt'.format(key_list[0])
    image_root = 'pascal-person-part/{}/Images/'.format(key_list[1])
    human_dir = 'pascal-person-part/{}/Human_ids/'.format(key_list[1])
    cate_dir = 'pascal-person-part/{}/Category_ids/'.format(key_list[1])
    save_json_file = 'pascal-person-part/annotations/PASCAL-Person-Part_parsing_{}.json'.format(key_list[0])

    ids_list = []
    f = open(ids_file, 'r')
    for l in f:
        img_id = l.strip()
        ids_list.append(img_id)
    f.close()

    # generate json
    categories = [{'id': 0, 'name': 'Background'}, {'id': 1, 'name': 'Lower-arm'}, {'id': 2, 'name': 'Head'},
                  {'id': 3, 'name': 'Upper-leg'}, {'id': 4, 'name': 'Torso'}, {'id': 5, 'name': 'Lower-leg'},
                  {'id': 6, 'name': 'Upper-arm'}, {'id': 7, 'name': 'Human'}]
    images = []
    annotations = []
    ins_id = 1
    for idx, img_id in enumerate(ids_list):
        img = cv2.imread(image_root + '{}.jpg'.format(img_id))
        height, width, _ = img.shape
        image = {'height': height, 'width': width, 'id': idx + 1, 'file_name': '{}.jpg'.format(img_id)}
        images.append(image)

        # if key == 'test':
        #     print(idx, img_id)
        #     continue

        semseg = cv2.imread(cate_dir + '{}.png'.format(img_id), 0)
        h, w = semseg.shape
        assert height == h and width == w

        mask = cv2.imread(human_dir + '{}.png'.format(img_id), 0)
        human_ids = np.unique(mask)
        # bg_id_index = np.where(human_ids == 0)[0]
        # human_ids = np.delete(human_ids, bg_id_index)
        for h_id in human_ids:
            human_mask = np.where(mask == h_id, 1, 0).astype(np.uint8)
            contours, _ = cv2.findContours(human_mask.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            segmentation = []
            for j in contours:
                if len(j) < 20:
                    continue
                _segs = []
                for p in j:
                    _segs.extend([int(p[0][1]), int(p[0][0])])
                    # img[p[0][1]][p[0][0]] = [0, 0, 255]
                segmentation.append(_segs)
            rect = cv2.boundingRect(human_mask.copy())
            x, y, w, h = rect

            if h_id == 0:
                print(idx, img_id, ins_id, h_id, 0)
                anno = {'segmentation': segmentation, 'iscrowd': 0, 'area': w * h, 'image_id': idx + 1,
                        'bbox': [x, y, w, h], 'category_id': 0, 'id': ins_id, 'parsing_id': int(h_id),
                        'ispart': 0, 'isfg': 0}
                annotations.append(anno)
                ins_id += 1
            else:
                print(idx, img_id, ins_id, h_id, 7)
                anno = {'segmentation': segmentation, 'iscrowd': 0, 'area': w * h, 'image_id': idx + 1,
                        'bbox': [x, y, w, h], 'category_id': 7, 'id': ins_id, 'parsing_id': int(h_id),
                        'ispart': 0, 'isfg': 1}
                annotations.append(anno)
                ins_id += 1

                human_semseg = semseg * human_mask
                part_ids = np.unique(human_semseg)
                bg_id_index = np.where(part_ids == 0)[0]
                part_ids = np.delete(part_ids, bg_id_index)
                for p_id in part_ids:
                    part_mask = np.where(human_semseg == p_id, 1, 0).astype(np.uint8)
                    contours, _ = cv2.findContours(part_mask.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
                    segmentation = []
                    for j in contours:
                        if len(j) < 10:
                            continue
                        _segs = []
                        for p in j:
                            _segs.extend([int(p[0][1]), int(p[0][0])])
                            # img[p[0][1]][p[0][0]] = [0, 0, 255]
                        segmentation.append(_segs)
                    rect = cv2.boundingRect(part_mask.copy())
                    x, y, w, h = rect

                    print(idx, img_id, ins_id, h_id, p_id)
                    anno = {'segmentation': segmentation, 'iscrowd': 0, 'area': w * h, 'image_id': idx + 1,
                            'bbox': [x, y, w, h], 'category_id': int(p_id), 'id': ins_id, 'parsing_id': int(h_id),
                            'ispart': 1, 'isfg': 1}
                    annotations.append(anno)
                    ins_id += 1

    json.dump({'images': images, 'categories': categories, 'annotations': annotations}, open(save_json_file, 'w'))