import os
import fnmatch
from pathlib import Path
from PIL import Image
import random

# SETTINGS ###################################################################################
# 0 - crop object area   1 - crop user defined area (creating background class)
mode = 0
# multiply width and height of bounding_box (create padding) ; = 1 => without padding
crop_padding = 1
##############################################################################################
# path of images folder
path_images = Path(r"")
# path of labels folder
path_labels = Path(r"")
##############################################################################################
# path of new images folder
path_images_new = Path(r"")
##############################################################################################


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


def num_class_to_txt_class(num_class):
    if num_class == 0:
        return "airplane"
    elif num_class == 1:
        return "bird"
    elif num_class == 2:
        return "drone"
    elif num_class == 3:
        return "helicopter"
    elif num_class == 4:
        return "insect"
    else:
        return "none"


def bounding_box_from_yolo_label(tuple_info, w, h, padding=1):
    x_center_norm = float(tuple_info[1])
    y_center_norm = float(tuple_info[2])
    width_norm = float(tuple_info[3])
    height_norm = float(tuple_info[4])

    left_norm = x_center_norm - padding*width_norm/2
    right_norm = x_center_norm + padding*width_norm/2
    top_norm = y_center_norm - padding*height_norm/2
    bottom_norm = y_center_norm + padding*height_norm/2

    left = max(0, int(round(left_norm * w)))
    right = min(w, int(round(right_norm * w)))
    top = max(0, int(round(top_norm * h)))
    bottom = min(h, int(round(bottom_norm * h)))

    return_tuple = (left, top, right, bottom)
    return return_tuple


def create_background_tuple(tuple_array, w, h):
    last_tuple = tuple_array[-1]
    h_last = last_tuple[3]-last_tuple[1]
    w_last = last_tuple[2]-last_tuple[0]
    # print(w_last)
    # print(h_last)

    # generate random bb
    random_top = random.randint(0, h - h_last)
    random_bottom = min(random_top + h_last, h)
    random_left = random.randint(0, w - w_last)
    random_right = min(random_left + w_last, w)

    valid_bb = 0
    counter_try = 0
    # validate if random area is out of objects
    while valid_bb == 0 and counter_try < 15:
        valid_bb = 1
        for k in range(len(tuple_array)):
            # get tuple info
            current_tuple = tuple_array[k]
            c_left = current_tuple[0]
            c_top = current_tuple[1]
            c_right = current_tuple[2]
            c_bottom = current_tuple[3]

            # validate
            if ((random_left >= c_left) and (random_left <= c_right)) or ((random_right >= c_left) and (random_right <= c_right)):
                if ((random_top >= c_top) and (random_top <= c_bottom)) or ((random_bottom >= c_top) and (random_bottom <= c_bottom)):
                    valid_bb = 0
                    print("random bb is not valid")
                    print(counter_try)
                    break

        if valid_bb == 0:
            # generate random bb
            random_top = random.randint(0, h - h_last)
            random_bottom = min(random_top + h_last, h)
            random_left = random.randint(0, w - w_last)
            random_right = min(random_left + w_last, w)
            counter_try += 1

    bg_tuple = (random_left, random_top, random_right, random_bottom)
    return bg_tuple, valid_bb


# find all labels, images and parameters in path folder
images_array = listdir(str(path_images), "*.jpg")
labels_array = listdir(str(path_labels), "*.txt")

valid = 1
if len(labels_array) == len(images_array):
    for i in range(len(labels_array)):
        name1 = labels_array[i][:-4]
        name2 = images_array[i][:-4]
        if name1 != name2:
            valid = 0
            break
else:
    valid = 0

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

if valid == 1 and mode == 0:
    for i in range(len(images_array)):
        pic_full_path = (str(path_images) + '\\' + str(images_array[i]))
        label_full_path = (str(path_labels) + '\\' + str(labels_array[i]))

        pic1 = Image.open(pic_full_path)
        pic = pic1.convert('RGB')

        f = open(label_full_path, "r")
        line = f.readline()
        j = 0
        while line != "":
            # get information from line
            tuple_label = line.replace("\n", "").split(" ")
            # print(tuple_label)

            # extract class name
            class_str = num_class_to_txt_class(int(tuple_label[0]))
            # print(class_str)

            # get bounding box
            width = pic.size[0]
            height = pic.size[1]
            tuple_bb = bounding_box_from_yolo_label(tuple_label, width, height, crop_padding)
            # crop ###############################################################################
            pic_crop = pic.crop(tuple_bb)

            # saving...
            pic_full_new_path = (str(path_images_new) + '\\' + class_str + '_' + str(j) + str(images_array[i]))
            pic_crop.save(pic_full_new_path)

            # read new line
            line = f.readline()
            j += 1

if valid == 1 and mode == 1:
    for i in range(len(images_array)):
        pic_full_path = (str(path_images) + '\\' + str(images_array[i]))
        label_full_path = (str(path_labels) + '\\' + str(labels_array[i]))

        pic1 = Image.open(pic_full_path)
        pic = pic1.convert('RGB')

        tuple_bb_array = []
        f = open(label_full_path, "r")
        line = f.readline()
        while line != "":
            # get information from line
            tuple_label = line.replace("\n", "").split(" ")
            # print(tuple_label)

            # read new line
            line = f.readline()

            # extract class name
            class_str = num_class_to_txt_class(int(tuple_label[0]))
            # print(class_str)

            # get bounding box
            width = pic.size[0]
            height = pic.size[1]
            tuple_bb = bounding_box_from_yolo_label(tuple_label, width, height)
            tuple_bb_array.append(tuple_bb)

            # crop ###############################################################################
            if line == "":
                bg_tuple_bb, success = create_background_tuple(tuple_bb_array, width, height)
                if success == 1:
                    pic_crop = pic.crop(bg_tuple_bb)
                    # saving...
                    pic_full_new_path = (str(path_images_new) + '\\'+'background_' + str(images_array[i]))
                    pic_crop.save(pic_full_new_path)
                else:
                    print("cannot find valid bg")

##################################################################################################
