import cv2
import json
import os
import numpy as np
from PIL import Image, ImageDraw

def convert_seg_coord_to_mask(img_size, coords):
    """Converts the segmentation coords found in COCO dataset to mask

    Args:
        img_size (_type_): _description_
        coords (_type_): _description_

    Returns:
        _type_: _description_
    """

    img = Image.new('L', img_size, 0)
    for polygon in coords:
        ImageDraw.Draw(img).polygon(polygon, outline=1, fill=1)
    mask = np.array(img)    

    return mask

class AnnotsGTGetter:
    
    def __init__(self, cfg_obj):

        self.cfg_obj = cfg_obj
        
        self.img_folder_path = cfg_obj['dataset']['img_folder_path']
        self.json_folder_path = cfg_obj['dataset']['annotations_folder_path']
        self.annot_json_fname = cfg_obj['dataset']['annotations_fname']
        self.labels_dict = cfg_obj['error_analysis']['labels_dict']
        self.task = cfg_obj['error_analysis']['task']

        json_file = open(self.json_folder_path + self.annot_json_fname)
        self.annot_data = json.load(json_file)

        self.img_ids_in_json = [annot['image_id'] for annot in self.annot_data['annotations']]
        self.all_imgs = os.listdir(self.img_folder_path)

        return

    def get_imgs(self):
        """method to get the mutually -inclusive- images between the img_ids in json and those in the folder path

        not needed because all images in folder were accounted for in the json...
        """
        all_img_ids_in_folder = [int(i[:-4]) for i in self.all_imgs]

        all_imgs_found = [i for i in all_img_ids_in_folder if i in self.img_ids_in_json]

        print (len(all_imgs_found))


    def get_annots(self, img_fname = '000000576052.jpg'):
        """retrieve annotation given a filename

        Args:
            img_fname (_type_): image file name
        
        Returns:
            np array: all annotations of an image
        """

        # change img_fname for extraction purpose
        # assumes jpg, png, but not jpeg...
        # TODO - what if jpeg?
        annots = []
        img_id = int(img_fname[:-4])
        img = Image.open(self.img_folder_path + img_fname)
        
        for annot in self.annot_data['annotations']:
            if img_id == annot['image_id']:
                if annot['category_id'] in list(self.labels_dict.values()):
                    if self.task == "det":
                        annots.append([annot['category_id'],annot['bbox'][0],annot['bbox'][1],annot['bbox'][2],annot['bbox'][3]])
                    elif self.task == "seg":
                        # call convert_seg_coord_to_mask to convert segmentations [x1,y1,x2,y2,...,xn,yn] to binary mask
                        mask = convert_seg_coord_to_mask(img.size, annot['segmentation'])
                        annots.append([annot['category_id'], mask])

        if self.task == "det":
            return np.array(annots)
        elif self.task == "seg":
            return annots

    def get_gt_annots(self):
        """goes into the image folder, calls get_annots to extract image annotation

        Returns:
            dict: all annotations 
        """

        # create dictionary of gt annots
        # for img in os.listdir(self.img_folder_path):
        #     self.get_annots(img)
        all_gt_annots = {img: self.get_annots(img) for img in os.listdir(self.img_folder_path)}

        return all_gt_annots

if __name__ == '__main__':
    import yaml
    # get_annots()
    cfg_file = open("cfg/cfg.yml")
    cfg_obj = yaml.load(cfg_file, Loader=yaml.FullLoader)
    annots_obj = AnnotsGTGetter(cfg_obj)
    gt_dict = annots_obj.get_gt_annots()
    # print (gt_dict)
    # annots_obj.get_imgs()
    
    # # to output mask
    # img_annots = gt_dict['000000576031.jpg']

    # import matplotlib.pyplot as plt
    # plt.imshow(img_annots[0][1])
    # plt.show()