# -*- coding: utf-8 -*-
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved

import time
import functools
import multiprocessing as mp
import numpy as np
import os
from lvis import LVIS
from pycocotools import mask as maskUtils


def annToRLE(ann, img_size):
    h, w = img_size
    segm = ann['segmentation']
    if type(segm) == list:
        # polygon -- a single object might consist of multiple parts
        # we merge all parts into one mask rle code
        rles = maskUtils.frPyObjects(segm, h, w)
        rle = maskUtils.merge(rles)
    elif type(segm['counts']) == list:
        # uncompressed RLE
        rle = maskUtils.frPyObjects(segm, h, w)
    else:
        # rle
        rle = ann['segmentation']
    return rle


def annToMask(ann, img_size):
    rle = annToRLE(ann, img_size)
    m = maskUtils.decode(rle)
    return m


def _process_instance_to_semantic(anns, output_semantic, img):
    img_size = (img["height"], img["width"])
    output = np.zeros(img_size, dtype=np.uint8)
    for ann in anns:
        mask = annToMask(ann, img_size)
        output[mask == 1] = ann["category_id"] // 5
    # save as compressed npz
    np.savez_compressed(output_semantic, mask=output)
    # Image.fromarray(output).save(output_semantic)


def create_lvis_semantic_from_instance(instance_json, sem_seg_root):
    """
    Create semantic segmentation annotations from panoptic segmentation
    annotations, to be used by PanopticFPN.

    It maps all thing categories to contiguous ids starting from 1, and maps all unlabeled pixels to class 0

    Args:
        instance_json (str): path to the instance json file, in COCO's format.
        sem_seg_root (str): a directory to output semantic annotation files
    """
    os.makedirs(sem_seg_root, exist_ok=True)

    lvis_detection = LVIS(instance_json)

    def iter_annotations():
        for img_id in lvis_detection.get_img_ids():
            anns_ids = lvis_detection.get_ann_ids([img_id])
            anns = lvis_detection.load_anns(anns_ids)
            img = lvis_detection.load_imgs([img_id])[0]
            file_name = os.path.splitext(img["file_name"])[0]
            output = os.path.join(sem_seg_root, file_name + '.npz')
            yield anns, output, img

    # # single process
    # print("Start writing to {} ...".format(sem_seg_root))
    # start = time.time()
    # for anno, oup, img in iter_annotations():
    #     _process_instance_to_semantic(
    #         anno, oup, img)
    # print("Finished. time: {:.2f}s".format(time.time() - start))
    # return

    pool = mp.Pool(processes=max(mp.cpu_count() // 2, 4))

    print("Start writing to {} ...".format(sem_seg_root))
    start = time.time()
    pool.starmap(
        functools.partial(
            _process_instance_to_semantic),
        iter_annotations(),
        chunksize=100,
    )
    print("Finished. time: {:.2f}s".format(time.time() - start))


if __name__ == "__main__":
    dataset_dir = os.path.join(os.path.dirname(__file__), "lvis")
    for s in ["train"]:
        create_lvis_semantic_from_instance(
            os.path.join(dataset_dir, "lvis_v0.5_{}.json".format(s)),
            os.path.join(dataset_dir, "thing_{}".format(s)),
        )