import gradio as gr
import PIL
from PIL import Image, ImageOps
from collections import Counter
import numpy as np

size_options ={"1200 x 1500": (1200, 1500),
    "1200 x 1200": (1200, 1200),
    "1200 x 628": (1200, 628),
    "Meta - Square (Feed)": (1200, 1200),
    "Meta - Story": (1080, 1920),
    "Tiktok Pangle - 1280 x 720": (1280, 720),
    "Tiktok Pangle - 1200 x 628": (1200, 628),
    "Tiktok Pangle - 640 x 640": (640, 640),
    "Tiktok Pangle - 300 x 250": (300, 250),
    "Tiktok Pangle - 720 x 1280": (720, 1280),
    "Tiktok Pangle - 300 x 50": (300, 50),
    "Tiktok Pangle - 320 x 50": (320, 50),
    "Tiktok Pangle - 320 x 480": (320, 480),
    "Tiktok Pangle - 250 x 250": (250, 250),
    "Tiktok Pangle - 600 x 770": (600, 770),
    "Tiktok Pangle - 720 x 1067": (720, 1067),
    "Tiktok Pangle - 480 x 320": (480, 320)}

def get_dominant_color(image):
    image = image.resize((50, 50))  # Resize to speed up
    pixels = np.array(image)
    pixels = pixels.reshape((-1, 3))
    counter = Counter(map(tuple, pixels))
    dominant_color = counter.most_common(1)[0][0]
    return dominant_color

def change_size(image, selected_size, size_align):
    target_size = size_options[selected_size]
    
    # Maintain aspect ratio without stretching
    image.thumbnail(target_size, PIL.Image.Resampling.LANCZOS)
    
    dominant_color = get_dominant_color(image)
    
    new_image = Image.new("RGB", target_size, dominant_color)
    
    # Calculate offset based on alignment choice
    x_offset = (target_size[0] - image.width) // 2 if size_align == "Center" else 0 if size_align == "Left" else target_size[0] - image.width
    y_offset = (target_size[1] - image.height) // 2 # this to keep the image on center
    
    new_image.paste(image, (x_offset, y_offset))
    
    return new_image

def gradio_interface():
    with gr.Blocks() as demo:
        gr.Markdown("# Advanced Image Resizer with Background and Margins")
        with gr.Row():
            image_input = gr.Image(type="pil", label="Upload Image")
            size_input = gr.Dropdown(list(size_options.keys()), label="Size Options")
        with gr.Row():
            size_alignment = gr.Radio(["Left", "Center", "Right"], value="Center", label="Size Alignment")
        image_output = gr.Image(type="pil", label="Resized Image with Background")
        submit_button = gr.Button("Resize Image")

        submit_button.click(fn=change_size, 
                            inputs=[image_input, size_input, size_alignment], 
                            outputs=image_output)

    return demo

demo = gradio_interface()
demo.launch(share=True)