import cv2
import json
import os
import numpy as np

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']

        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])
        for annot in self.annot_data['annotations']:
            if img_id == annot['image_id']:
                if annot['category_id'] in list(self.labels_dict.values()):
                    annots.append([annot['category_id'],annot['bbox'][0],annot['bbox'][1],annot['bbox'][2],annot['bbox'][3]])
        
        return np.array(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__':
    # get_annots()
    annots_obj = AnnotsGTGetter()
    gt_dict = annots_obj.get_gt_annots()
    print (gt_dict)
    # annots_obj.get_imgs()