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]
    new_width, new_height = target_size
    
    # Resize the image to maintain aspect ratio and fit within the half-height
    image.thumbnail((new_width, new_height // 2), PIL.Image.Resampling.LANCZOS)

    # Obtain the dominant color of the image
    dominant_color = get_dominant_color(image)
    
    # Create a new image with the dominant color as the background
    new_image = Image.new("RGB", target_size, dominant_color)

    # Size alignment for position: left, right, center, bottom, top
    if size_align == "Left":
        x_offset = 0
        y_offset = 0  # Place at the top
    elif size_align == "Right":
        x_offset = target_size[0] - image.width
        y_offset = 0  # Place at the top
    elif size_align == "Top":
        x_offset = (target_size[0] - image.width) // 2
        y_offset = 0  # Place at the top
    elif size_align == "Bottom":
        x_offset = (target_size[0] - image.width) // 2
        y_offset = (target_size[1] - image.height)  # Place at the bottom
    else:  # Center
        x_offset = (target_size[0] - image.width) // 2
        y_offset = 0  # Place at the top

    # Paste the resized image onto the new background
    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", "Top", "Bottom"], 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)