import gradio as gr
from PIL import Image, ImageDraw, ImageOps
import io
import os
import subprocess
from math import cos, radians, sin

def preprocess_image(image):
    # Ensure the image is square by padding it
    size = max(image.size)
    new_image = Image.new("RGBA", (size, size), (255, 255, 255, 0))
    new_image.paste(image, ((size - image.width) // 2, (size - image.height) // 2))
    return new_image

def create_gif(img1, img2, transition_type):
    frames = []
    duration = 100  # Duration for each frame in milliseconds
    total_frames = 18  # Total number of frames

    try:
        # Preprocess images to make them square and same size
        img1 = preprocess_image(img1)
        img2 = preprocess_image(img2)

        # Set size for the GIF
        size = (256, 256)
        img1 = img1.resize(size, Image.LANCZOS)
        img2 = img2.resize(size, Image.LANCZOS)

        if transition_type == "default":
            # Calculate step size for consistent speed
            full_width = size[0]
            step = full_width // (total_frames // 2)  # Divide by 2 as we have 2 parts to the animation

            # Generate frames from left to right
            for i in range(0, full_width, step):
                frame = Image.new('RGBA', size)
                frame.paste(img1, (0, 0))
                frame.paste(img2.crop((i, 0, full_width, size[1])), (i, 0), mask=img2.crop((i, 0, full_width, size[1])))

                draw = ImageDraw.Draw(frame)
                draw.line((i, 0, i, size[1]), fill=(0, 255, 0), width=2)

                frame = frame.convert('P', palette=Image.ADAPTIVE)
                frames.append(frame)

            # Generate frames from right to left
            for i in range(full_width, step, -step):
                frame = Image.new('RGBA', size)
                frame.paste(img1, (0, 0))
                frame.paste(img2.crop((i, 0, full_width, size[1])), (i, 0), mask=img2.crop((i, 0, full_width, size[1])))

                draw = ImageDraw.Draw(frame)
                draw.line((i, 0, i, size[1]), fill=(0, 255, 0), width=2)

                frame = frame.convert('P', palette=Image.ADAPTIVE)
                frames.append(frame)

        elif transition_type == "rotate":
            # Create a mask that is 2x the size of the image
            mask_size = (size[0] * 2, size[1] * 2)
            mask = Image.new('L', mask_size, 0)
            draw = ImageDraw.Draw(mask)
            draw.rectangle([size[0], 0, mask_size[0], mask_size[1]], fill=255)

            center_x, center_y = size[0] // 2, size[1] // 2

            for angle in range(0, 360, 360 // total_frames):
                rotated_mask = mask.rotate(angle, center=(mask_size[0] // 2, mask_size[1] // 2), expand=False)
                cropped_mask = rotated_mask.crop((size[0] // 2, size[1] // 2, size[0] // 2 + size[0], size[1] // 2 + size[1]))
                frame = Image.composite(img1, img2, cropped_mask)

                draw = ImageDraw.Draw(frame)
                reverse_angle = -angle + 90
                end_x1 = center_x + int(size[0] * 1.5 * cos(radians(reverse_angle)))
                end_y1 = center_y + int(size[1] * 1.5 * sin(radians(reverse_angle)))
                end_x2 = center_x - int(size[0] * 1.5 * cos(radians(reverse_angle)))
                end_y2 = center_y - int(size[1] * 1.5 * sin(radians(reverse_angle)))
                draw.line([center_x, center_y, end_x1, end_y1], fill=(0, 255, 0), width=3)
                draw.line([center_x, center_y, end_x2, end_y2], fill=(0, 255, 0), width=3)

                frame = frame.convert('P', palette=Image.ADAPTIVE)
                frames.append(frame)

        # Save as GIF
        output = io.BytesIO()
        frames[0].save(output, format='GIF', save_all=True, append_images=frames[1:], duration=duration, loop=0, optimize=True)
        output.seek(0)

        return output.getvalue()
    except Exception as e:
        raise ValueError(f"Error creating GIF: {e}")

def create_gif_gradio(image1, image2, transition_type):
    gif_data = create_gif(image1, image2, transition_type)
    
    # Save the GIF to a temporary file
    temp_output_path = "output.gif"
    with open(temp_output_path, "wb") as f:
        f.write(gif_data)
    
    return temp_output_path

# Gradio interface
iface = gr.Interface(
    fn=create_gif_gradio,
    inputs=[
        gr.Image(type="pil", label="Image 1"),
        gr.Image(type="pil", label="Image 2"),
        gr.Radio(["default", "rotate"], label="Transition Type")
    ],
    outputs=gr.Image(type="filepath", label="Generated GIF"),
    title="GIF Generator",
    description="Upload two images and select a transition type to generate an animated GIF."
)

# Launch the interface
iface.launch()