import cv2
import numpy as np
import gradio as gr

def find_bounding_box(mask):
    """Find the bounding box around the largest contour in the mask."""
    mask = mask.astype(np.uint8)
    _, binary_mask = cv2.threshold(mask, 1, 255, cv2.THRESH_BINARY)

    contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if contours:
        x, y, w, h = cv2.boundingRect(np.vstack(contours))
        return x, y, w, h
    return None

def overlay_logo_on_image(original_img, mask, logo):
    """Overlay a scaled logo while maintaining aspect ratio and centering it in the bounding box."""
    bbox = find_bounding_box(mask)
    if bbox is None:
        print("No bounding box found. Returning original image.")
        return original_img

    x, y, w, h = bbox
    print(f"Bounding box found at: x={x}, y={y}, w={w}, h={h}")

    # Get original logo dimensions
    logo_h, logo_w = logo.shape[:2]

    # Compute the scaling factor while maintaining aspect ratio
    scale = min(w / logo_w, h / logo_h)
    new_w, new_h = int(logo_w * scale), int(logo_h * scale)

    # Resize the logo while maintaining aspect ratio
    logo_resized = cv2.resize(logo, (new_w, new_h), interpolation=cv2.INTER_AREA)

    # Ensure logo_resized has 4 channels (RGBA)
    if logo_resized.shape[2] == 3:
        logo_resized = cv2.cvtColor(logo_resized, cv2.COLOR_RGB2RGBA)

    # Create a blank transparent canvas (RGBA)
    logo_canvas = np.zeros((h, w, 4), dtype=np.uint8)

    # Compute centering offsets
    x_offset = (w - new_w) // 2
    y_offset = (h - new_h) // 2

    # Place the resized logo at the center of the bounding box
    logo_canvas[y_offset:y_offset + new_h, x_offset:x_offset + new_w] = logo_resized

    # Convert original image to 4-channel if necessary
    if original_img.shape[2] == 3:
        original_img = cv2.cvtColor(original_img, cv2.COLOR_RGB2RGBA)

    # Blend the logo into the image using alpha blending
    alpha = logo_canvas[:, :, 3] / 255.0
    for c in range(3):
        original_img[y:y+h, x:x+w, c] = (
            (1 - alpha) * original_img[y:y+h, x:x+w, c] + alpha * logo_canvas[:, :, c]
        )

    return original_img

def process_images(original_img, mask, logo, progress=gr.Progress()):
    """Process images: overlay logo using uploaded mask."""
    original_img = cv2.cvtColor(np.array(original_img), cv2.COLOR_RGB2BGR)
    mask = np.array(mask)
    print(f"Logo type before conversion: {type(logo)}, dtype: {getattr(logo, 'dtype', 'N/A')}")
    logo = np.array(logo)
    if logo.dtype != np.uint8:
        logo = logo.astype(np.uint8)
    logo = cv2.cvtColor(np.array(logo), cv2.COLOR_RGB2BGR)


    if len(mask.shape) == 3:
        mask = cv2.cvtColor(mask, cv2.COLOR_RGB2GRAY)

    mask = (mask > 0).astype(np.uint8) * 255 

    print(f"Mask shape: {mask.shape}, unique values: {np.unique(mask)}")

    result_img = overlay_logo_on_image(original_img, mask, logo)
    result_img = cv2.cvtColor(result_img, cv2.COLOR_BGRA2RGBA) if result_img.shape[2] == 4 else cv2.cvtColor(result_img, cv2.COLOR_BGR2RGB)

    return result_img

with gr.Blocks() as demo :
  gr.Markdown("## Inpaint Logo on Uploaded Mask")
  with gr.Row() :
    with gr.Column() :
      with gr.Row():
        original_img = gr.Image(label="Upload Original Image", type="numpy")
      with gr.Row():
        mask = gr.Image(label="Upload Mask (Binary)", type="numpy")
        logo = gr.Image(label="Upload Logo", type="numpy")
      btn = gr.Button("Submit")
    with gr.Column() :
      output = gr.Image(label="Output Image", format="png")
  btn.click(fn=process_images, inputs=[original_img, mask, logo], outputs=output)

demo.launch(share=True, debug=True, show_error=True)