from PIL import Image, ImageOps
import random
import os
import numpy as np

RANDOM_SEED = 42
random.seed(RANDOM_SEED)

def resize_image(image, target_size):
    original_width, original_height = image.size
    target_width, target_height = target_size

    # Calculate the aspect ratios
    width_ratio = target_width / original_width
    height_ratio = target_height / original_height

    # Use the smaller ratio to ensure the image fits within the target dimensions
    scaling_factor = min(width_ratio, height_ratio)

    # Calculate new dimensions
    new_width = int(original_width * scaling_factor)
    new_height = int(original_height * scaling_factor)

    # Resize the image
    return image.resize((new_width, new_height), Image.LANCZOS)

def rotate_image(image, angle):
    return image.rotate(angle, expand=True, resample=Image.Resampling.BICUBIC)

def tilt_image(image, max_tilt_angle=10):
    width, height = image.size

    # Simulate tilting by applying an affine transformation
    tilt_angle = random.uniform(-max_tilt_angle, max_tilt_angle)
    x_shift = width * np.sin(np.radians(tilt_angle))
    
    # Affine transformation matrix
    if tilt_angle > 0:
        transform_matrix = (1, x_shift / height, -x_shift / 2, 0, 1, 0)
    else:
        transform_matrix = (1, x_shift / height, x_shift / 2, 0, 1, 0)
    
    return image.transform(image.size, Image.AFFINE, transform_matrix, resample=Image.Resampling.BICUBIC)

def generate_images(object_image, output_folder, background_image_path,
                    background_width=960, background_height=720, num_images=50, 
                    max_rotation_angle=45, max_tilt_angle=10):
    
    background_image = Image.open(background_image_path).convert("RGB")
    background_image = background_image.resize((background_width, background_height))

    for i in range(num_images):
        # Create a copy of the background image
        background = background_image.copy()

        # Randomly rotate the object image
        rotation_angle = random.uniform(-max_rotation_angle, max_rotation_angle)
        rotated_image = rotate_image(object_image, rotation_angle)

        # Apply random tilt to the object image
        tilted_image = tilt_image(rotated_image, max_tilt_angle)

        # Random position for the object
        max_x = background_width - tilted_image.width
        max_y = background_height - tilted_image.height
        x = random.randint(0, max_x)
        y = random.randint(0, max_y)

        # Paste the object image onto the background
        background.paste(tilted_image, (x, y), tilted_image)

        # Save the new image
        background.save(f"{output_folder}/variation_{i + 1}.png")

def main():
    background_width = 960
    background_height = 720
    new_object_size = (300, 300)

    os.makedirs("gen_data", exist_ok=True)
    output_folder = "gen_data/"

    # Load the object image (with transparency)
    object_image = Image.open("img/drone.png").convert("RGBA")
    object_image = resize_image(object_image, new_object_size)

    # Load the background image from a second source
    background_image_path = "img/background.jpg"

    # Generate images with object placed on top of the background
    generate_images(object_image, output_folder, background_image_path,
                    background_width, background_height, num_images=50)

if __name__ == "__main__":
    main()