import gradio as gr
from PIL import Image, ImageOps
import numpy as np
import cv2

# Global variables to maintain state
tshirt_img = None
design_img = None
vertices = []
polygon_drawn = False
design_pos = (250, 250)

def load_tshirt(image):
    global tshirt_img
    tshirt_img = image.resize((600, 600), Image.Resampling.LANCZOS)
    return tshirt_img

def load_design(image):
    global design_img, design_pos
    design_img = image.resize((100, 100), Image.Resampling.LANCZOS)
    design_pos = (250, 250)
    return tshirt_img, design_img, design_pos

def add_polygon_vertex(x, y):
    global vertices, polygon_drawn
    if tshirt_img and not polygon_drawn:
        vertices.append((x, y))
        if len(vertices) > 2 and np.linalg.norm(np.array(vertices[-1]) - np.array(vertices[0])) < 10:
            vertices[-1] = vertices[0]  # Close the polygon
            polygon_drawn = True
    return tshirt_img, design_img, design_pos

def fit_design():
    global design_img, tshirt_img, vertices, polygon_drawn
    if design_img and tshirt_img and polygon_drawn:
        src_pts = np.array([[0, 0], [design_img.width, 0], [design_img.width, design_img.height], [0, design_img.height]], dtype="float32")
        dst_pts = np.array(vertices, dtype="float32")

        M = cv2.getPerspectiveTransform(src_pts, dst_pts)
        design_np = np.array(design_img)
        transformed_design = cv2.warpPerspective(design_np, M, (tshirt_img.width, tshirt_img.height))
        transformed_design_img = Image.fromarray(transformed_design)

        design_img = transformed_design_img
        return tshirt_img, design_img, design_pos
    return tshirt_img, design_img, design_pos

def rotate_design(angle):
    global design_img
    if design_img:
        angle = int(angle)
        rotated_design = design_img.rotate(angle, expand=True)
        design_img = rotated_design
        return tshirt_img, design_img, design_pos
    return tshirt_img, design_img, design_pos

def display_image(x=0, y=0, angle=0, image_type=None, image=None):
    if image_type == "tshirt":
        return load_tshirt(image)
    elif image_type == "design":
        return load_design(image)
    elif image_type == "polygon":
        return add_polygon_vertex(x, y)
    elif image_type == "fit":
        return fit_design()
    elif image_type == "rotate":
        return rotate_design(angle)
    return tshirt_img, design_img, design_pos

with gr.Blocks() as demo:
    tshirt_image = gr.Image(label="T-Shirt Image")
    design_image = gr.Image(label="Design Image")
    rotate_slider = gr.Slider(minimum=0, maximum=360, step=1, label="Rotate Design")
    x_coord = gr.Number(label="X Coordinate")
    y_coord = gr.Number(label="Y Coordinate")
    
    tshirt_output = gr.Image(label="Output T-Shirt with Design")
    design_output = gr.Image(label="Output Design on T-Shirt")
    
    # Upload options for T-shirt and design images
    tshirt_image.upload(lambda image: display_image(image_type="tshirt", image=image), inputs=tshirt_image, outputs=tshirt_output)
    design_image.upload(lambda image: display_image(image_type="design", image=image), inputs=design_image, outputs=[tshirt_output, design_output])

    # Controls for adding polygon vertices, fitting the design, and rotating
    gr.Button("Add Polygon Vertex").click(lambda x, y: display_image(x=x, y=y, image_type="polygon"), inputs=[x_coord, y_coord], outputs=[tshirt_output, design_output])
    gr.Button("Fit Design").click(lambda: display_image(image_type="fit"), None, [tshirt_output, design_output])
    rotate_slider.change(lambda angle: display_image(angle=angle, image_type="rotate"), inputs=rotate_slider, outputs=[tshirt_output, design_output])

demo.launch(share=True)