8import os
import face_recognition
import pickle
from tensorflow.keras.preprocessing.image import ImageDataGenerator
import cv2

def generateDatasetImages():
    # Define ImageDataGenerator with desired augmentation parameters
    datagen = ImageDataGenerator(
        rotation_range=20,      # Rotate images randomly within the range of -20 to +20 degrees
        width_shift_range=0.1,  # Shift images horizontally (as a fraction of total width)
        height_shift_range=0.1, # Shift images vertically (as a fraction of total height)
        shear_range=0.1,        # Shear intensity (shear angle in counter-clockwise direction in radians)
        zoom_range=0.1,         # Zoom images randomly by a factor of 0.1
        horizontal_flip=True,   # Flip images horizontally
        vertical_flip=True,     # Flip images vertically
    )

    # Directory containing example images
    input_dir = 'static/data/dataset/'

    # Iterate over each image file in the directory
    for filename in os.listdir(input_dir):
        if filename.endswith(".jpg") or filename.endswith(".jpeg") or filename.endswith(".png"):
            output_dir = os.path.join(input_dir, "augmented_images", os.path.splitext(filename)[0])

            if not os.path.exists(output_dir):
                # Load the image using OpenCV (required for face_recognition)
                img_path = os.path.join(input_dir, filename)
                img = cv2.imread(img_path)
                rgb_img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # Convert BGR to RGB
                
                # Detect faces in the image using face_recognition
                face_locations = face_recognition.face_locations(rgb_img)
                
                if len(face_locations) == 0:
                    print(f"No faces detected in '{filename}'. Skipping augmentation.")
                    continue
                
                # Create a directory for augmented images based on the image filename
                os.makedirs(output_dir, exist_ok=True)
                
                # Convert pixel values to [0, 1]
                rgb_img = rgb_img / 255.0
                
                # Generate augmented images
                img = rgb_img.reshape((1,) + rgb_img.shape)

                # Generate augmented images
                num_samples = 100  # Number of augmented samples to generate
                augmented_images = []

                # Generate augmented images
                for batch in datagen.flow(img, batch_size=1):
                    augmented_image = (batch[0] * 255).astype('uint8')  # Convert back to uint8
                    augmented_images.append(augmented_image)
                    num_samples -= 1
                    if num_samples == 0:
                        break
                
                # Save augmented images
                for i, augmented_image in enumerate(augmented_images):
                    aug_filename = f'{os.path.splitext(filename)[0]}_{i+1}.jpg'
                    aug_filepath = os.path.join(output_dir, aug_filename)
                    cv2.imwrite(aug_filepath, cv2.cvtColor(augmented_image, cv2.COLOR_RGB2BGR))

                print(f"{len(augmented_images)} augmented images saved in '{output_dir}' directory.")
    
    known_face_encodings = []
    known_face_names = []
    dataset_path = "static/data/dataset"

    # Load images and augment dataset
    image_files = [file for file in os.listdir(dataset_path) if file.endswith(('.jpg', '.jpeg', '.png'))]

    for file in image_files:
        name = os.path.splitext(os.path.basename(file))[0]
        images = os.listdir(os.path.join(dataset_path, "augmented_images", name))

        file_path = face_recognition.load_image_file(os.path.join(dataset_path, file))

        face_locations = face_recognition.face_locations(file_path, model="hog")
        encoding = face_recognition.face_encodings(file_path, face_locations)

        known_face_encodings.append(encoding)
        known_face_names.append(name)
        
        # for image in images:            
        #     subfile = face_recognition.load_image_file(os.path.join(dataset_path, "augmented_images", name, image))

        #     face_locations = face_recognition.face_locations(subfile, model="hog")
        #     face_encodings = face_recognition.face_encodings(subfile, face_locations)

        #     for encoding in face_encodings:
        #         known_face_names.append(name)
        #         known_face_encodings.append(encoding)

    # Save the trained model for future use
    with open('static/data/dataset/trained_faces.pkl', 'wb') as f:
        pickle.dump((known_face_encodings, known_face_names), f)

generateDatasetImages()