import os.path
import albumentations as A
import cv2
from pathlib import Path
import datetime

class ImageAnnotation:

    def read_annotations(self) -> [[]]:

        with open(self.dataset_annotation_path, 'r') as f:
            rows = f.readlines()
            arr = []
            for row in rows:
                str_coordinates = row.split(' ')
                class_num = str_coordinates.pop(0)
                coordinates = [float(n) for n in str_coordinates]
                coordinates.append(class_num)
                arr.append(coordinates)

            return arr

    def __init__(self, image_path, current_dataset_dir, labels_dir, images_dir):
        self.image_path = image_path
        self.image_name = Path(image_path).stem
        self.dataset_image_path = os.path.join(current_dataset_dir, images_dir, self.image_path + '.jpg')
        self.image = cv2.imread(self.dataset_image_path)

        self.annotation_path = os.path.join(labels_dir, self.image_path + '.txt')
        self.dataset_annotation_path = os.path.join(current_dataset_dir, labels_dir, self.image_path + '.txt')
        self.annotations = self.read_annotations()


def image_processing(img_ann, current_dataset_dir,labels_dir,images_dir: ImageAnnotation) -> [ImageAnnotation]:
    category_ids = []
    bboxes = ImageAnnotation(img_ann,current_dataset_dir,labels_dir,images_dir).read_annotations()
    for i in range(len(bboxes)):
        category_ids.append(bboxes[i][4])
        bboxes[i].pop(4)

    transform = A.Compose([
        A.HorizontalFlip(p=0.5),
        A.RandomBrightnessContrast(p=0.2),
    ], bbox_params=A.BboxParams(format='yolo', label_fields=['category_ids']))

    bboxes = bboxes
    imag = ImageAnnotation(img_ann,current_dataset_dir,labels_dir,images_dir).image

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

    return transformed_image, transformed_bboxes, transformed_category_ids


def write_results(img_ann, current_dataset_dir, labels_dir, images_dir: ImageAnnotation):

    file_start_save = 'Zombobase-' + str(datetime.date.today())
    for i in range(5):

        transformed_image, transformed_bboxes, transformed_category_ids, = image_processing(img_ann, current_dataset_dir, labels_dir, images_dir)
        cv2.imwrite(os.path.join(current_dataset_dir, images_dir,str(i)+ImageAnnotation(img_ann,current_dataset_dir,labels_dir,images_dir).image_path + '.jpg'), transformed_image)
        with open(os.path.join(current_dataset_dir, labels_dir, str(i)+ImageAnnotation(img_ann, current_dataset_dir, labels_dir, images_dir).image_path + '.txt'), 'w') as f:
            print(os.path.join(current_dataset_dir, labels_dir, str(i)+ImageAnnotation(img_ann, current_dataset_dir, labels_dir, images_dir).image_path + '.txt'))
            for bbox, category_id in zip(transformed_bboxes, transformed_category_ids):
                print(bbox)
                x_center, y_center, width, height = bbox
                cla = category_id
                f.write(f"{cla} {x_center} {y_center} {width} {height}\n")
#
#
def process_image(current_dataset_dir,images_dir):
    file_annotation = []
    file_annotation_finished =[]
    for foldername, subfolders, filenames in os.walk(os.path.join(current_dataset_dir,images_dir)):
        file_annotation.append(filenames)
        print(file_annotation)

    for i in range(len(file_annotation[0])):
        ff = file_annotation[0][i].split('.')
        ff.pop(-1)
        ff = '.'.join(ff)
        file_annotation_finished.append(ff)

    return file_annotation_finished

def main(current_dataset_dir, labels_dir, images_dir):
    process_image(current_dataset_dir,images_dir)
    for i in process_image(current_dataset_dir,images_dir):
        write_results(i, current_dataset_dir, labels_dir, images_dir)

#main(os.path.join('Zombobase-' + str(datetime.date.today()), 'datasets', 'zombobase-current','test'),'labels', 'images')