import os
import shutil

from PIL import Image, ImageEnhance
import numpy as np
import random

from tqdm import tqdm


# Color Jittering function
def apply_color_jitter(image, brightness=0.5, contrast=0.5, saturation=0.5):
    """
    Apply color jitter to the image to simulate different lighting conditions.
    """
    enhancer = ImageEnhance.Brightness(image)
    image = enhancer.enhance(brightness)

    enhancer = ImageEnhance.Contrast(image)
    image = enhancer.enhance(contrast)

    enhancer = ImageEnhance.Color(image)
    image = enhancer.enhance(saturation)

    return image


# Noise Injection function
def apply_noise_injection(image, mean=0, std=1):
    """
    Add random noise to an image to improve the robustness of the model to noise.
    """
    np_image = np.array(image)
    noise = np.random.normal(mean, std, np_image.shape).astype(np.uint8)
    noisy_image = np_image + noise
    noisy_image = np.clip(noisy_image, 0, 255).astype(np.uint8)
    return Image.fromarray(noisy_image)


# Occlusion function
def apply_occlusion(image, occlusion_size=(50, 50)):
    """
    Add a random occlusion to an image to simulate real-world occlusions.
    """
    np_image = np.array(image)
    height, width, _ = np_image.shape
    top_left = (np.random.randint(0, height - occlusion_size[0]), np.random.randint(0, width - occlusion_size[1]))
    bottom_right = (top_left[0] + occlusion_size[0], top_left[1] + occlusion_size[1])

    np_image[top_left[0]:bottom_right[0], top_left[1]:bottom_right[1], :] = 0
    return Image.fromarray(np_image)


# Example usage
if __name__ == "__main__":
    # Load the image
    imgs_path = "20240315/images"
    labels_path = "20240315/labels"


    save_images_path = "dataset/enhance_yawn/images"
    save_labels_path = "dataset/enhance_yawn/labels"

    os.makedirs(save_images_path,exist_ok=True)

    os.makedirs(save_labels_path,exist_ok=True)


    for file in tqdm(os.listdir(imgs_path)):


        if file.startswith("yawn"):

            # exit()
            img_path = os.path.join(imgs_path, file)

            img = Image.open(img_path)
            label_path = os.path.join(labels_path, file.split(".")[0] + ".txt")

            if os.path.exists(label_path):
                # Apply transformations
                for i in range(5):
                    img_color_jittered = apply_color_jitter(img, brightness= random.uniform(0.8, 1.5), contrast=random.uniform(0.8, 1.5), saturation=random.uniform(0.8, 1.5))
                    new_img_name = file.split(".")[0]+f"{i}_color"+".png"
                    new_img_path = os.path.join(save_images_path,new_img_name)
                    img_color_jittered.save(new_img_path)
                    new_label_path = os.path.join(save_labels_path,new_img_name.split(".")[0] + ".txt")
                    shutil.copyfile(label_path,new_label_path)


                # 噪声
                img_noisy = apply_noise_injection(img)
                new_img_name = file.split(".")[0] + f"noisy" + ".png"
                new_img_path = os.path.join(save_images_path, new_img_name)
                img_noisy.save(new_img_path)

                new_label_path = os.path.join(save_labels_path, new_img_name.split(".")[0] + ".txt")
                shutil.copyfile(label_path, new_label_path)

                # img_occluded = apply_occlusion(img)
                # new_img_name = file.split(".")[0] + f"occluded" + ".png"
                # new_img_path = os.path.join(save_images_path, new_img_name)
                # img_occluded.save(new_img_path)
                # new_label_path = os.path.join(save_labels_path, new_img_name.split(".")[0] + ".txt")
                # shutil.copyfile(label_path, new_label_path)