'''
'''
import logging
import threading
from pathlib import Path

from PIL import Image

from .coco_converter import CocoBaseStrategy

from ....common import id2rgb, write_to_json, tqdm_converter

logger = logging.getLogger("superannotate-python-sdk")


class CocoPanopticConverterStrategy(CocoBaseStrategy):
    def __init__(self, args):
        super().__init__(args)

    def _sa_to_coco_single(self, id_, annotation_json, id_generator):

        image_commons = self._prepare_single_image_commons(
            id_, annotation_json['metadata']
        )
        res = self.conversion_algorithm(
            image_commons, annotation_json['instances'], id_generator
        )
        return res

    def sa_to_output_format(self):
        out_json = self._create_skeleton()
        out_json['categories'] = self._create_categories(
            self.export_root / 'classes_mapper.json'
        )

        images = []
        annotations = []
        id_generator = self._make_id_generator()
        annot_id_generator = self._make_id_generator()

        jsons = self.make_anno_json_generator()
        images_converted = []
        images_not_converted = []
        total_num = self.get_num_total_images()
        finish_event = threading.Event()
        tqdm_thread = threading.Thread(
            target=tqdm_converter,
            args=(
                total_num, images_converted, images_not_converted, finish_event
            ),
            daemon=True
        )
        logger.info('Converting to COCO JSON format')
        tqdm_thread.start()
        for json_ in jsons:
            idx = next(annot_id_generator)
            res = self._sa_to_coco_single(idx, json_, id_generator)

            panoptic_mask = json_['metadata']['panoptic_mask']

            Image.fromarray(id2rgb(res[2])).save(panoptic_mask)

            annotation = {
                'image_id': res[0]['id'],
                'file_name': Path(panoptic_mask).name,
                'segments_info': res[1]
            }
            annotations.append(annotation)

            images.append(res[0])
            images_converted.append(json_)

        out_json['annotations'] = annotations
        out_json['images'] = images
        write_to_json(
            self.output_dir / '{}.json'.format(self.dataset_name), out_json
        )
        finish_event.set()
        tqdm_thread.join()


class CocoObjectDetectionStrategy(CocoBaseStrategy):
    def __init__(self, args):
        super().__init__(args)

    def _sa_to_coco_single(self, id_, annotation_json, id_generator):

        image_commons = self._prepare_single_image_commons(
            id_, annotation_json['metadata']
        )

        def make_annotation(
            category_id, image_id, bbox, segmentation, area, anno_id
        ):
            if self.task == 'object_detection':
                segmentation = [
                    [
                        bbox[0], bbox[1], bbox[0], bbox[1] + bbox[3],
                        bbox[0] + bbox[2], bbox[1] + bbox[3], bbox[0] + bbox[2],
                        bbox[1]
                    ]
                ]
            annotation = {
                'id': anno_id,  # making sure ids are unique
                'image_id': image_id,
                'segmentation': segmentation,
                'iscrowd': 0,
                'bbox': bbox,
                'area': area,
                'category_id': category_id
            }

            return annotation

        res = self.conversion_algorithm(
            make_annotation, image_commons, annotation_json['instances'],
            id_generator
        )
        return res

    def sa_to_output_format(self):
        out_json = self._create_skeleton()
        out_json['categories'] = self._create_categories(
            self.export_root / 'classes_mapper.json'
        )

        images = []
        annotations = []
        id_generator = self._make_id_generator()
        annot_id_generator = self._make_id_generator()
        jsons = self.make_anno_json_generator()

        images_converted = []
        images_not_converted = []
        finish_event = threading.Event()
        total_num = self.get_num_total_images()
        tqdm_thread = threading.Thread(
            target=tqdm_converter,
            args=(
                total_num, images_converted, images_not_converted, finish_event
            ),
            daemon=True
        )
        logger.info('Converting to COCO JSON format')
        tqdm_thread.start()
        for json_ in jsons:
            idx = next(annot_id_generator)
            try:
                images_converted.append(json_)
                res = self._sa_to_coco_single(idx, json_, id_generator)
            except Exception as e:
                images_not_converted.append(json_)
                raise

            images.append(res[0])
            if len(res[1]) < 1:
                self.increase_converted_count()
            for ann in res[1]:
                annotations.append(ann)
        out_json['annotations'] = annotations
        out_json['images'] = images

        write_to_json(
            self.output_dir / '{}.json'.format(self.dataset_name), out_json
        )
        finish_event.set()
        tqdm_thread.join()


class CocoKeypointDetectionStrategy(CocoBaseStrategy):
    def __init__(self, args):
        super().__init__(args)

    def sa_to_output_format(self):
        out_json = self._create_skeleton()
        jsons = self.make_anno_json_generator()

        total_num = self.get_num_total_images()
        id_generator = self._make_id_generator()
        id_generator_anno = self._make_id_generator()
        id_generator_img = self._make_id_generator()

        res = self.conversion_algorithm(
            jsons, id_generator, id_generator_anno, id_generator_img,
            self._make_image_info, total_num
        )

        out_json['categories'] = res[0]
        out_json['annotations'] = res[1]
        out_json['images'] = res[2]
        write_to_json(
            self.output_dir / '{}.json'.format(self.dataset_name), out_json
        )
        self.set_num_converted(len(out_json['images']))
