import math
import os
import datetime
import cv2

file_start_save = 'Zombobase-'+str(datetime.date.today())
file_txt = []
file_jpg = []
labels_dir = 'labels'
images_dir = 'images'
current_dataset_dir = os.path.join('datasets', 'zombobase-current')


def save(annotations, image, i, f_n):
    cv2.imwrite(os.path.join(file_start_save, current_dataset_dir, f_n, images_dir, file_jpg[i]), image)
    with open(os.path.join(file_start_save, current_dataset_dir, f_n, labels_dir, file_txt[i]), 'w') as f:
        for iii in range(len(annotations)):

            f.write(annotations[iii])
    annotations = []


def piercing_photo_file():
    for foldername, subfolders, filenames in os.walk(current_dataset_dir):
        for subfolder in subfolders:
            cu = os.path.join('datasets', 'zombobase-current',subfolder)
            for foldername, subfolders, filenames in os.walk(cu):
                for filename in filenames:
                    f = filename.split('.')
                    if f[-1] == 'txt':
                        file_txt.append(filename)
                    elif f[-1] == 'jpg':
                        file_jpg.append(filename)


def file_validation(annotations, file_txt, file_jpg):
    try:
        os.makedirs(os.path.join(file_start_save, current_dataset_dir,'validation'))
        os.makedirs(os.path.join(file_start_save, current_dataset_dir, 'validation', images_dir))
        os.makedirs(os.path.join(file_start_save, current_dataset_dir, 'validation', labels_dir))
    except FileExistsError:
        pass
    for i in range(len(file_txt)):
        image = cv2.imread(os.path.join(current_dataset_dir, images_dir, file_jpg[i]))
        with open(os.path.join(current_dataset_dir, labels_dir, file_txt[i]), 'r') as file:
            lines = file.readlines()
            for line in lines:
                annotations.append(line)
        save(annotations, image, i, 'validation')
        annotations = []


def sort_file():
    try:
        os.makedirs(os.path.join(file_start_save, current_dataset_dir))
    except FileExistsError:
        pass
    annotations = []
    folder_name = ['test', 'train']
    percent_file = [0.20, 0.10]
    for f_n, p_f in zip(folder_name, percent_file):
        try:
            os.makedirs(os.path.join(file_start_save, current_dataset_dir, f_n))
            os.makedirs(os.path.join(file_start_save, current_dataset_dir, f_n, images_dir))
            os.makedirs(os.path.join(file_start_save, current_dataset_dir, f_n, labels_dir))
        except FileExistsError:
            pass
        for i in range(math.ceil(len(file_txt) * p_f)):
            image = cv2.imread(os.path.join(current_dataset_dir, images_dir, file_jpg[i]))
            with open(os.path.join(current_dataset_dir, labels_dir, file_txt[i]), 'r') as file:
                lines = file.readlines()


                for line in lines:
                    annotations.append(line)
                save(annotations, image, i, f_n)

                annotations = []
                print(annotations)
                file_txt.pop(i)
                file_jpg.pop(i)
    file_validation(annotations, file_txt, file_jpg)


def main():
    piercing_photo_file()
    sort_file()