import gradio as gr
import cv2
import numpy as np
import albumentations as A
import random

def apply_augmentations(image, flip_h, flip_v, rotate, crop, gray, scale,
                       prob_flip_h, prob_flip_v, prob_rotate, prob_crop, prob_gray, prob_scale,
                       rotation_limit):
    augmentations = []
    
    if flip_h:
        augmentations.append(A.HorizontalFlip(p=float(prob_flip_h)))
    if flip_v:
        augmentations.append(A.VerticalFlip(p=float(prob_flip_v)))
    if rotate:
        augmentations.append(A.Rotate(limit=(-rotation_limit, rotation_limit), p=float(prob_rotate)))
    if crop:
        augmentations.append(A.RandomResizedCrop(
            size=(image.shape[0], image.shape[1]),
            scale=(0.8, 1.0),
            p=float(prob_crop)
        ))
    if gray:
        augmentations.append(A.ToGray(p=float(prob_gray)))
    if scale:
        scale_factor = random.uniform(0.8, 1.2)
        augmentations.append(A.Resize(
            height=int(image.shape[0] * scale_factor),
            width=int(image.shape[1] * scale_factor),
            p=float(prob_scale)
        ))

    transform = A.Compose(augmentations)
    
    # Generate 9 augmented images
    augmented_images = []
    for _ in range(9):
        augmented = transform(image=image)
        augmented_images.append(augmented['image'])
    
    # Find maximum dimensions
    max_height = max(img.shape[0] for img in augmented_images)
    max_width = max(img.shape[1] for img in augmented_images)
    
    # Add padding to all images
    padded_images = []
    for img in augmented_images:
        h, w = img.shape[:2]
        pad_top = (max_height - h) // 2
        pad_bottom = max_height - h - pad_top
        pad_left = (max_width - w) // 2
        pad_right = max_width - w - pad_left
        
        # Handle both RGB and grayscale images
        if len(img.shape) == 3:
            padded = cv2.copyMakeBorder(img, pad_top, pad_bottom, pad_left, pad_right, 
                                      cv2.BORDER_CONSTANT, value=[128, 128, 128])
        else:
            padded = cv2.copyMakeBorder(img, pad_top, pad_bottom, pad_left, pad_right, 
                                      cv2.BORDER_CONSTANT, value=[128])
        padded_images.append(padded)
    
    # Create a 3x3 grid
    rows = []
    for i in range(0, 9, 3):
        row = np.hstack(padded_images[i:i+3])
        rows.append(row)
    grid = np.vstack(rows)
    
    return grid

def main():
    with gr.Blocks() as demo:
        with gr.Row():
            with gr.Column():
                input_image = gr.Image(label="Input Image")
            with gr.Column():
                output_image = gr.Image(label="Output Image (3x3 Grid)")
        
        with gr.Row():
            with gr.Column():
                flip_h = gr.Checkbox(label="Horizontal Flip")
                prob_flip_h = gr.Slider(minimum=0, maximum=1, value=0.5, label="Probability")
            
            with gr.Column():
                flip_v = gr.Checkbox(label="Vertical Flip")
                prob_flip_v = gr.Slider(minimum=0, maximum=1, value=0.5, label="Probability")
            
            with gr.Column():
                rotate = gr.Checkbox(label="Rotate")
                prob_rotate = gr.Slider(minimum=0, maximum=1, value=0.5, label="Probability")
                rotation_limit = gr.Slider(minimum=0, maximum=90, value=10, label="Rotation Limit (±degrees)")
            
            with gr.Column():
                crop = gr.Checkbox(label="Random Crop")
                prob_crop = gr.Slider(minimum=0, maximum=1, value=0.5, label="Probability")

            with gr.Column():
                gray = gr.Checkbox(label="Grayscale")
                prob_gray = gr.Slider(minimum=0, maximum=1, value=0.5, label="Probability")
                
            with gr.Column():
                scale = gr.Checkbox(label="Random Scale (0.8-1.2x)")
                prob_scale = gr.Slider(minimum=0, maximum=1, value=0.5, label="Probability")

        with gr.Row():
            run_button = gr.Button("Apply Augmentations")

        run_button.click(
            fn=apply_augmentations,
            inputs=[
                input_image,
                flip_h, flip_v, rotate, crop, gray, scale,
                prob_flip_h, prob_flip_v, prob_rotate, prob_crop, prob_gray, prob_scale,
                rotation_limit
            ],
            outputs=output_image
        )

    demo.launch()

if __name__ == "__main__":
    main()