import os
from PIL import Image
import natsort
import json
from collections import Counter
import random
from sklearn.model_selection import train_test_split

proj = 'hdvehicle' 
path = "/data/si/darwin/si/construction-equipment/releases/"+proj+ '/annotations/'

def separate_train_test(file_list, ratio):
    
    file_list = natsort.natsorted(file_list)
    
    category_list=[]

    for idx,file_name in enumerate(file_list):
        with open(path+file_name, 'r') as f:
            json_data = json.load(f)

        # categories
        for idy,instance in enumerate(json_data["annotations"]):
            print(instance)
            if instance["class"] in category_list:
                pass
            else:
                category_list.append(instance["class"])

    category_list=natsort.natsorted(category_list)
            
    category_dic_count = dict.fromkeys(category_list, 0)
    category_dic_filename = dict.fromkeys(category_list)
    
    for key in category_dic_filename.keys():
        category_dic_filename[key]=[]

    for idx,file_name in enumerate(file_list):
        with open(path+file_name, 'r') as f:
            json_data = json.load(f)

        class_id_list = []
        # categories
        for idy,instance in enumerate(json_data["annotations"]):
            class_id_list.append(instance["class"])
        count_items = Counter(class_id_list)
        random.shuffle(count_items)

        class_id = count_items.most_common(n=1)[0][0]
        category_dic_count[class_id]+=1
        category_dic_filename[class_id].append(file_name)
        
    train = [] # train + val
    test = [] # test

    for key in category_dic_count.keys():
        if (category_dic_count[key]):
            temp=category_dic_filename[key]

            train_temp, test_temp = train_test_split(temp, test_size=ratio)

            train.extend(train_temp)
            test.extend(test_temp)

    return train, test

def class_count(file_list, ratio):
    
    file_list = natsort.natsorted(file_list)
    
    category_list=[]

    for idx,file_name in enumerate(file_list):
        with open(path+file_name, 'r') as f:
            json_data = json.load(f)

        # categories
        for idy,instance in enumerate(json_data["annotations"]):
            if instance["class"] in category_list:
                pass
            else:
                category_list.append(instance["class"])

    category_list=natsort.natsorted(category_list)
            
    category_dic_count = dict.fromkeys(category_list, 0)
    category_dic_filename = dict.fromkeys(category_list)
    
    for key in category_dic_filename.keys():
        category_dic_filename[key]=[]

    for idx,file_name in enumerate(file_list):
        with open(path+file_name, 'r') as f:
            json_data = json.load(f)

        class_id_list = []
        # categories
        for idy,instance in enumerate(json_data["annotations"]):
            class_id_list.append(instance["class"])
        count_items = Counter(class_id_list)
        random.shuffle(count_items)

        class_id = count_items.most_common(n=1)[0][0]
        category_dic_count[class_id]+=1
        category_dic_filename[class_id].append(file_name)
        
    train = [] # train + val
    test = [] # test

    for key in category_dic_count.keys():
        if (category_dic_count[key]):
            temp=category_dic_filename[key]

            train_temp, test_temp = train_test_split(temp, test_size=ratio)

            train.extend(train_temp)
            test.extend(test_temp)

    return train, test