import json
import os

import cv2
import numpy as np


error_list = ["23382.png", "23441.png", "20714.png", "20727.png", "23300.png", "21200.png"]


def mask2box(mask):
    index = np.argwhere(mask == 1)
    rows = index[:, 0]
    clos = index[:, 1]
    y1 = int(np.min(rows))  # y
    x1 = int(np.min(clos))  # x
    y2 = int(np.max(rows))
    x2 = int(np.max(clos))
    return (x1, y1, x2, y2)


def gen_coco(phase):
    result = {
        "info": {"description": "PIC2.0 dataset."},
        "categories": [{"supercategory": "none", "id": 1, "name": "person"}],
    }
    out_json = phase + "_person.json"
    store_segmentation = True

    images_info = []
    labels_info = []
    img_id = 0
    files = tuple(open("pic/list5/" + phase + "_id", "r"))
    files = (_.strip() for _ in files)

    for index, image_name in enumerate(files):
        image_name = image_name + ".png"
        print(index, image_name)
        if image_name in error_list:
            continue
        instance = cv2.imread(os.path.join("instance", phase, image_name), flags=cv2.IMREAD_GRAYSCALE)
        semantic = cv2.imread(os.path.join("semantic", phase, image_name), flags=cv2.IMREAD_GRAYSCALE)
        # print(instance.shape, semantic.shape)
        h = instance.shape[0]
        w = instance.shape[1]
        images_info.append({"file_name": image_name[:-4] + ".jpg", "height": h, "width": w, "id": index})
        instance_ids = np.unique(instance)
        for instance_id in instance_ids:
            if instance_id == 0:
                continue
            instance_part = instance == instance_id
            object_pos = instance_part.nonzero()
            # category_id_ = int(semantic[object_pos[0][0], object_pos[1][0]])
            category_id = int(np.max(semantic[object_pos[0], object_pos[1]]))
            # assert category_id_ == category_id, (category_id_, category_id)
            if category_id != 1:
                continue
            area = int(instance_part.sum())
            x1, y1, x2, y2 = mask2box(instance_part)
            w = x2 - x1 + 1
            h = y2 - y1 + 1
            segmentation = []
            if store_segmentation:
                contours, hierarchy = cv2.findContours(
                    (instance_part * 255).astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE
                )
                for contour in contours:
                    contour = contour.flatten().tolist()
                    if len(contour) > 4:
                        segmentation.append(contour)
                if len(segmentation) == 0:
                    print("error")
                    continue
            labels_info.append(
                {
                    "segmentation": segmentation,  # poly
                    "area": area,  # segmentation area
                    "iscrowd": 0,
                    "image_id": index,
                    "bbox": [x1, y1, w, h],
                    "category_id": category_id,
                    "id": img_id,
                },
            )
            img_id += 1
        # break
    result["images"] = images_info
    result["annotations"] = labels_info
    with open("pic/annotations/" + out_json, "w") as f:
        json.dump(result, f, indent=4)


if __name__ == "__main__":
    if not os.path.exists("pic/annotations/"):
        os.mkdirs("pic/annotations/")
    gen_coco("train")
    gen_coco("val")
    # gen_coco("test")
