from matplotlib import pyplot as plt
import cv2
import json
from pathlib import Path



def compute_categories_id(data_name, world):
    #Category = 'banana1'
    #Category = 'pear2'
    #Category = "orange2"
    # Opening JSON file
    f = open(f'{data_name}/Meta/{world}.json')
    
    # returns JSON object as 
    # a dictionary
    data = json.load(f)
    
    # Iterating through the json
    # list

    catergories_label_to_id={}
    catergories_id_to_label={}
    catergories_instance_array_cat_to_id={}
    catergories_instance_array_id_to_cat={}

    for k in data['categories']:
        catergories_label_to_id[k['label']]=k['id']
        catergories_id_to_label[k['id']]=k['label']
        catergories_instance_array_cat_to_id[k['label']]=[]

    for k in data['objects']:
        #print(k)
        #catergories_instance_array[catergories_id_to_label[i['category_id']]]
        catergories_instance_array_id_to_cat[k['id']] = catergories_id_to_label[k['category_id']]
        catergories_instance_array_cat_to_id[catergories_id_to_label[k['category_id']]].append(k['id'])
        # if i['category_id'] == id_category :
        #     print("Hello fruits instance")
        #     id_instances.append(i['id'])
        #     print(i['id']) 

    # print("catergories_instance_array_cat_to_id : ", catergories_instance_array_cat_to_id)
    # print("catergories_instance_array_id_to_cat : ", catergories_instance_array_id_to_cat)


    # Closing file
    f.close()


    return catergories_instance_array_id_to_cat, catergories_instance_array_cat_to_id



def compute_id_good_occ(data_name, count, catergories_instance_array_id_to_cat, catergories_instance_array_cat_to_id, Occ_wanted):

    f2 = open(f'{data_name}/Occlusion/{count}.json')

    data2 = json.load(f2)


    catergories_occ_array = {}

    for cat in catergories_instance_array_cat_to_id :
        #print(cat)

        catergories_occ_array[cat] = []

    for i in data2:
        #print('i : ',i)
        #print(i['id'])
        #print(id_instances)
        if i['occlusion_value'] > 0.5 :
            catergories_occ_array[catergories_instance_array_id_to_cat[i['id']]].append(i['id'])



        # if i['id'] in id_instances :
        #     print("Hello banana instance occ")
        #     if i['occlusion_value'] > 0.5 :
        #         id_instances_good.append(i['id'])
        #         print(i['id'])
        #         print(i['occlusion_value'])

    print(catergories_occ_array)




    # Closing file
    f2.close()

    return catergories_occ_array




def instance(im, id):
    #im = im * 255
    im[im == id] = 255
    im[im < 255] = 0
    im = im / 255
    #plt.imshow(im, cmap='gray')
    #plt.show()
    return im


def generate_instance_mask(data_name, Nb_camera, Nb_world,list_categories, occ_target):

    for i in range(1, Nb_world + 1): # worlds
        
        catergories_instance_array_id_to_cat, catergories_instance_array_cat_to_id = compute_categories_id(data_name, i)
        


        for j in range(1, Nb_camera+1): # cameras
            k = ((i-1)*Nb_camera) + j

            catergories_occ_array = compute_id_good_occ(data_name, k, catergories_instance_array_id_to_cat, catergories_instance_array_cat_to_id, occ_target)


            for categories in list_categories:

                if len(catergories_occ_array[categories]) == 1 :


                    path = Path(f"{data_name}/Instance_Segmentation/{k}.png") 
                    if path.is_file():
                        id = catergories_occ_array[categories][0]
                        print("iddd : ",id)
                        img = cv2.imread(f"{data_name}/Instance_Segmentation/{k}.png", cv2.IMREAD_UNCHANGED) # plt.imread(path)

                        #print("img[817][308] : ", img[817][308])
                        print("img[308][817] : ", img[308][817])


                        instance_img = instance(img, id)
                        print("instance_img[308][817] : ", instance_img[308][817])


                        cv2.imwrite(f"{data_name}/Generated/Instance_Mask/{categories}/{k}.png", 255*instance_img)