import math
import datetime
import cv2
import albumentations as A
import numpy as np
import os
def file_magnification(folder_path):
    file_txt = []
    file_jpg = []
    for foldername, subfolders, filenames in os.walk(folder_path):
        for filename in filenames:
            f = filename.split('.')

            if f[-1] == 'txt':
                file_txt.append(filename)
            elif f[-1] == 'jpg':
                file_jpg.append(filename)

    for k in range(len(file_jpg)):
        image = cv2.imread(f'{folder_path}\\{file_jpg[k]}')
        annotations = []
        with open(f'{folder_path}\\{file_txt[k]}', 'r') as file:
            lines = file.readlines()
            for line in lines:
                annotations.append(line)


        main_fillet_yolo_conversion = []
        fillet_yolo_bboxes = []
        fillet_yolo_class = []
        print(annotations)
        for ii in range(len(annotations)):
            a = annotations[ii].split(' ')
            for i in range(len(a)):
                try:
                    main_fillet_yolo_conversion.append(int(a[i]))
                except ValueError:
                    main_fillet_yolo_conversion.append(float(a[i]))
                print(main_fillet_yolo_conversion)
            fillet_yolo_class.append(main_fillet_yolo_conversion[0])
            del main_fillet_yolo_conversion[0]
            fillet_yolo_bboxes.append(main_fillet_yolo_conversion)
            main_fillet_yolo_conversion = []

        for o in range(10):
            bboxes = fillet_yolo_bboxes
            category_ids = fillet_yolo_class

            transform = A.Compose([
                A.HorizontalFlip(p=0.5),
                A.RandomBrightnessContrast(p=0.2),
                A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.2, rotate_limit=15, p=0.5),
            ], bbox_params=A.BboxParams(format='yolo', label_fields=['category_ids']))

            transformed = transform(image=image, bboxes=bboxes, category_ids=category_ids)
            transformed_image = transformed['image']
            transformed_bboxes = transformed['bboxes']
            transformed_category_ids = transformed['category_ids']

            cv2.imwrite(f'{folder_path}\\{o}_{file_jpg[k]}', transformed_image)
            with open(f'{folder_path}\\{o}_{file_txt[k]}', 'w') as f:
                for bbox, category_id in zip(transformed_bboxes, transformed_category_ids):
                    x_center, y_center, width, height = bbox
                    cla = category_id
                    f.write(f"{cla} {x_center} {y_center} {width} {height}\n")
file_txt_1 = []
file_jpg_1 = []
file_start = 'Zombobase-'+str(datetime.date.today())
folder_path = ('train')
for foldername, subfolders, filenames in os.walk(folder_path):
    for subfolder in subfolders:
        folder_path = (f'train\\{subfolder}')
        for foldername, subfolders, filenames in os.walk(folder_path):
            for filename in filenames:

                f = filename.split('.')

                if f[-1] == 'txt':
                    file_txt_1.append(filename)
                elif f[-1] == 'jpg':
                    file_jpg_1.append(filename)

annotations = []
os.makedirs(file_start)
file = ['test', 'train']
percent_fille = [0.20,0.10]
for fi, p_f in zip(file, percent_fille):
    os.makedirs(f'{file_start}\\{fi}')

    for i in range(math.ceil(len(file_txt_1)* p_f)):
        image = cv2.imread(f'D:\\train\\images\\{file_jpg_1[i]}')
        with open(f'D:\\train\\labels\\{file_txt_1[i]}', 'r') as file:
            lines = file.readlines()
            for line in lines:
                annotations.append(line)
        cv2.imwrite(f'{file_start}\\{fi}\\{file_jpg_1[i]}', image)
        with open(f'{file_start}\\{fi}\\{file_txt_1[i]}', 'w') as f:
            for iii in range(len(annotations)):
                f.write(annotations[iii])

        annotations = []
        del file_txt_1[i]
        del file_jpg_1[i]

os.makedirs(f'{file_start}\\validation')
for a, j in zip(file_txt_1, file_jpg_1):
    image = cv2.imread(f'D:\\train\\images\\{j}')
    with open(f'D:\\train\\labels\\{a}', 'r') as file:
        lines = file.readlines()
        for line in lines:
            annotations.append(line)
    cv2.imwrite(f'{file_start}\\Validation\\{j}', image)
    with open(f'{file_start}\\Validation\\{a}', 'w') as f:
        for iii in range(len(annotations)):
            f.write(annotations[iii])
    annotations = []
file = ['test', 'train','validation']
for i in file:
    file_magnification(f'{file_start}\\{i}')