from pathlib import Path
import fnmatch
import os
import math
import random
import shutil
import numpy as np

# SETTINGS ###################################################################################
# MODE #######################################################################################
# mode 0 - analyze images, mode
# mode 1 - copy picked images to new iter folder
# mode 2 - mode 1 + use 'analyze_merge_dataset_output.txt'
mode = 1

# FOLDERS ####################################################################################
# path of images folder (for gen_01+ use merge_images)
path_images = Path("Image_database/gen_04/images")
# path of parameters folder (for gen_01+ use merge_parameters)
path_parameters = Path("Image_database/gen_04/parameters")
# path for new iteration folder
path_new_iter_images = Path("")

# TYPE SELECTION ###########################################################
# class (airplane, bird, drone, helicopter, insect)
class_name = [20, 20, 20, 20, 20]
# complexity of image (50% low; 50% high) # parameter [%]
complexity = [50, 50]
# area of object (40% low, 40% medium, 20% high) [%]
area = [40, 40, 20]
# weather in image (50% sunny, 50% not_sunny)
weather = [50, 50]
# number of images with multiple objects (just for print, only information)
multiple_object_min = 10


# number of multiple object has minimum value
# percentage representation of all parameters is valid across all sub-groups created
# by parameters (example: in sub-groups [class: bird, complexity: low, ...] is valid:
# area: percentage representation as specify)
# number of object in single object (min value of images with multiple object)
# has recommended value in whole database
##############################################################################################
# number of images used in reverse search in next iteration
number_of_images_to_pick = 16000
##############################################################################################


# get all files of defined type from folder
def listdir(dir_name, pattern="*"):
    return fnmatch.filter(os.listdir(dir_name), pattern)


def print_files_sorted_by_type(sorted_array, number_of_types_array, number_of_types_ideal_array):
    num_missing = 0
    for j in range(len(number_of_types_array)):
        # get parameters from parameter file name
        par_class = math.floor(j/12)
        par_area = math.floor((j % 12)/4)
        par_complexity = math.floor(((j % 12) % 4)/2)
        par_weather = j % 2

        if par_class == 0:
            str_class = "airplane"
        elif par_class == 1:
            str_class = "bird"
        elif par_class == 2:
            str_class = "drone"
        elif par_class == 3:
            str_class = "helicopter"
        else:
            str_class = "insect"

        if par_area == 0:
            str_area = "small"
        elif par_area == 1:
            str_area = "medium"
        else:
            str_area = "large"

        if par_complexity == 0:
            str_complexity = "small"
        else:
            str_complexity = "big"

        if par_weather == 0:
            str_weather = "sunny"
        else:
            str_weather = "not_sunny"
        '''
        # print all ##########################################################################
        print("class: "+str_class+", area: "+str_area+", complexity: "+str_complexity+
              ", weather: "+str_weather)
        print("files of exact type: "+str(sorted_array[j]))
        print("number of files of given parameters: "+str(number_of_types[j]) +
              ", ideal number of files: "+str(number_of_types_ideal_array[j]))
        print("")
        ######################################################################################
        '''
        # print missing ######################################################################
        if number_of_types_ideal_array[j] > number_of_types_array[j]:
            print("class: " + str_class + ", area: " + str_area + ", complexity: " +
                  str_complexity + ", weather: " + str_weather)
            print("files of exact file: " + str(sorted_array[j]))
            print("number of files of given parameters: " + str(number_of_types[j]) +
                  ", ideal number of files: " + str(number_of_types_ideal_array[j]))
            print("")
            num_missing += number_of_types_ideal_array[j] - number_of_types_array[j]
        ######################################################################################

    print("Number of all files: "+str(sum(number_of_types)))
    print("Percentage of missing images: "+str(num_missing/number_of_images_to_pick*100)+"%")


# GET DATA ###################################################################################
# find all parameters and images in path folder
images_array = listdir(str(path_images), "*.jpg")
parameters_array = listdir(str(path_parameters), "*.txt")

# names of parameters and images must be the same
valid = 1
parameters = []
if len(images_array) == len(parameters_array):
    for i in range(len(parameters_array)):
        name1 = parameters_array[i][:-4]
        parameters.append(name1[len(name1)-5:])
        # print(parameters[-1])
        name2 = images_array[i][:-4]
        if name1[:len(name1)-6] != name2:
            valid = 0
            break
else:
    valid = 0

if valid == 1:
    print("Name arrays of images and parameters are the same => valid.")
else:
    print("Name arrays of images and parameters are not the same => invalid.")
##############################################################################################

if valid:
    # ANALYSIS ###############################################################################
    number_of_multiple_object_in_images = 0
    number_of_types = [0] * 60
    number_of_types_ideal = [0] * 60
    names_of_types = []
    names_of_types_images = []

    # calculate optimal number of image types
    for i in range(60):
        # creating 2D array
        names_of_types.append([])
        names_of_types_images.append([])
        # ideal numbers images of certain type
        p_class = math.floor(i / 12)
        p_area = math.floor((i % 12) / 4)
        p_complexity = math.floor(((i % 12) % 4) / 2)
        p_weather = i % 2
        # calculating
        number_of_types_ideal[i] = round(class_name[p_class]/100 * area[p_area]/100 * complexity[p_complexity]/100
                                         * weather[p_weather]/100 * number_of_images_to_pick)

    for i in range(len(parameters_array)):
        string = parameters[i]
        # get parameters
        p_class = int(string[0])
        p_complexity = int(string[1])
        p_area = int(string[2])
        p_number = int(string[3])
        p_weather = int(string[4])
        type_number = p_weather + 2*p_complexity + 4*p_area + 12*p_class
        # increment of multiple
        number_of_multiple_object_in_images += p_number
        # increment of type counter
        number_of_types[type_number] += 1
        # append name of file to type array
        names_of_types[type_number].append(parameters_array[i])
        names_of_types_images[type_number].append(images_array[i])
    # PRINT ANALYSIS #########################################################################
    if mode == 0:
        print_files_sorted_by_type(names_of_types, number_of_types, number_of_types_ideal)
        print("percentage of multiple object in single images: "
              + str(number_of_multiple_object_in_images/len(parameters_array)*100)
              + "%, ideal minimum percentage: " + str(multiple_object_min) + "%")
    ##########################################################################################
    # CREATE GEN+1 ITER FOLDER ###############################################################
    elif mode == 1:
        print("random pick")
        for i in range(60):
            for k in range(min(len(names_of_types[i]), number_of_types_ideal[i])):
                pos_rand = random.randint(0, len(names_of_types_images[i])-1)
                print(names_of_types_images[i][pos_rand])
                # move images to "new_iter_images" folder
                shutil.copy(str(path_images)+'/'+str(names_of_types_images[i][pos_rand]),
                            str(path_new_iter_images))
                del names_of_types_images[i][pos_rand]
    elif mode == 2:
        print("random pick + don't pick types with sufficient contribution in dataset")
        enough = np.loadtxt('analyze_merge_dataset_output.txt', dtype=int)
        for i in range(60):
            if i in enough:
                print("skipped " + str(i))
                continue
            print("processing.. "+str(i))
            for k in range(min(len(names_of_types[i]), number_of_types_ideal[i])):
                pos_rand = random.randint(0, len(names_of_types_images[i]) - 1)
                print(names_of_types_images[i][pos_rand])
                # move images to "new_iter_images" folder
                shutil.copy(str(path_images) + '/' + str(names_of_types_images[i][pos_rand]),
                            str(path_new_iter_images))
                del names_of_types_images[i][pos_rand]
##############################################################################################
