import torch
import torchvision.transforms as transforms
from PIL import Image
import cv2
import numpy as np
import gradio as gr

# Load MiDaS model
midas = torch.hub.load("intel-isl/MiDaS", "DPT_Large", trust_repo=True)
midas.eval()

# Preprocessing function
def preprocess_image(image):
    transform = transforms.Compose([
        transforms.Resize(384),
        transforms.CenterCrop(384),
        transforms.ToTensor(),
        transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225],
        ),
    ])
    return transform(image).unsqueeze(0)

# Function to generate the displacement map
def generate_displacement_map(image_a):
    input_batch = preprocess_image(image_a)

    with torch.no_grad():
        depth_map = midas(input_batch)

    depth_map = depth_map.squeeze().cpu().numpy()
    depth_map = cv2.resize(depth_map, (image_a.width, image_a.height))

    depth_map = (depth_map - depth_map.min()) / (depth_map.max() - depth_map.min())
    displacement_map = depth_map * 30
    return displacement_map

# Function to warp and fit Image-B onto Image-A
def fit_and_warp_design(image_a, image_b, top_left_x, top_left_y, bottom_right_x, bottom_right_y):
    displacement_map = generate_displacement_map(image_a)

    # Extract bounding box coordinates
    top_left = (int(top_left_x), int(top_left_y))
    bottom_right = (int(bottom_right_x), int(bottom_right_y))

    # Resize the design to fit within the specified bounding box
    design_width = bottom_right[0] - top_left[0]
    design_height = bottom_right[1] - top_left[1]
    image_b = image_b.convert("RGBA")  # Ensure design image has an alpha channel
    image_b = image_b.resize((design_width, design_height))
    
    # Create a blank canvas with the same size as Image-A
    canvas = Image.new('RGBA', (displacement_map.shape[1], displacement_map.shape[0]), (0, 0, 0, 0))
    canvas.paste(image_b, top_left, image_b)  # Use image_b as the mask for pasting
    canvas_np = np.array(canvas)

    h, w = displacement_map.shape
    y_indices, x_indices = np.indices((h, w), dtype=np.float32)
    x_displacement = (x_indices + displacement_map).astype(np.float32)
    y_displacement = (y_indices + displacement_map).astype(np.float32)

    x_displacement = np.clip(x_displacement, 0, w - 1)
    y_displacement = np.clip(y_displacement, 0, h - 1)

    warped_canvas = cv2.remap(canvas_np, x_displacement, y_displacement, cv2.INTER_LINEAR, borderMode=cv2.BORDER_TRANSPARENT)

    image_a_rgba = image_a.convert("RGBA")
    image_a_np = np.array(image_a_rgba)

    non_transparent_pixels = warped_canvas[..., 3] > 0
    image_a_np[non_transparent_pixels] = warped_canvas[non_transparent_pixels]

    final_image = Image.fromarray(image_a_np)
    return final_image

# Gradio interface function
def process_images(image_a, image_b, top_left_x, top_left_y, bottom_right_x, bottom_right_y):
    result = fit_and_warp_design(image_a, image_b, top_left_x, top_left_y, bottom_right_x, bottom_right_y)
    return result

# Gradio UI components
image_input_a = gr.Image(label="Upload Clothing Image", type="pil")
image_input_b = gr.Image(label="Upload Design Image", type="pil")
top_left_x = gr.Slider(minimum=0, maximum=1000, label="Top-left X Coordinate", value=50)
top_left_y = gr.Slider(minimum=0, maximum=1000, label="Top-left Y Coordinate", value=100)
bottom_right_x = gr.Slider(minimum=0, maximum=1000, label="Bottom-right X Coordinate", value=300)
bottom_right_y = gr.Slider(minimum=0, maximum=1000, label="Bottom-right Y Coordinate", value=400)

# Define the Gradio interface
iface = gr.Interface(
    fn=process_images,
    inputs=[image_input_a, image_input_b, top_left_x, top_left_y, bottom_right_x, bottom_right_y],
    outputs="image",
    title="Clothing Design Fitting with Adjustable Bounding Box",
    description="Upload a clothing image and a design image. Adjust the design's position and size using sliders to fit the design onto the clothing.",
)

# Launch the interface
iface.launch()