import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import gradio as gr
import torch
import torchvision.transforms as transforms
from skimage.filters import sobel
from skimage.restoration import denoise_tv_chambolle
from scipy.interpolate import Rbf
from scipy.ndimage import map_coordinates


# Function to estimate a normal map from the cloth texture
def estimate_normal_map(image):
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    sobel_x = sobel(gray)
    sobel_y = sobel(gray)
    normal_map = np.stack([sobel_x, sobel_y, np.ones_like(sobel_x)], axis=-1)
    normal_map = normal_map / np.linalg.norm(normal_map, axis=-1, keepdims=True)
    return (normal_map * 255).astype(np.uint8)


# Function to apply Thin Plate Spline (TPS) warping
def apply_tps_warping(text_image, normal_map):
    h, w = text_image.shape[:2]
    x, y = np.meshgrid(np.arange(w), np.arange(h))
    
    # Generate control points from the normal map
    control_x = x + (normal_map[:, :, 0] - 128) * 0.5
    control_y = y + (normal_map[:, :, 1] - 128) * 0.5

    # Interpolate using Radial Basis Function (RBF)
    rbf_x = Rbf(x.flatten(), y.flatten(), control_x.flatten(), function='thin_plate')
    rbf_y = Rbf(x.flatten(), y.flatten(), control_y.flatten(), function='thin_plate')

    warped_x = rbf_x(x, y).astype(np.float32)
    warped_y = rbf_y(x, y).astype(np.float32)

    # Apply warping
    warped_text = cv2.remap(text_image, warped_x, warped_y, interpolation=cv2.INTER_LINEAR)
    
    return warped_text


# Function to blend text onto cloth using Poisson editing
def blend_text_cloth(cloth, text, x=50, y=50):
    cloth_bgr = np.array(cloth)
    text_bgr = np.array(text)

    normal_map = estimate_normal_map(cloth_bgr)

    # Resize text to fit on cloth
    text_resized = cv2.resize(text_bgr, (cloth_bgr.shape[1] // 2, cloth_bgr.shape[0] // 5))
    
    # Convert to grayscale and create a mask
    text_gray = cv2.cvtColor(text_resized, cv2.COLOR_BGR2GRAY)
    _, mask = cv2.threshold(text_gray, 1, 255, cv2.THRESH_BINARY)

    # Warp text using normal map
    warped_text = apply_tps_warping(text_resized, normal_map)

    # Blend the text using Poisson editing
    center = (x + text_resized.shape[1] // 2, y + text_resized.shape[0] // 2)
    blended = cv2.seamlessClone(warped_text, cloth_bgr, mask, center, cv2.MIXED_CLONE)

    return Image.fromarray(blended)


# Gradio function
def process_image(cloth_image, text, font_size=32, font_color=(255, 0, 0), x=50, y=50):
    # Convert font color input
    font_color = tuple(map(int, font_color.strip("()").split(",")))

    # Create a blank image with text
    text_img = Image.new('RGB', (400, 200), (0, 0, 0, 0))
    draw = ImageDraw.Draw(text_img)

    try:
        font = ImageFont.truetype("arial.ttf", font_size)
    except:
        font = ImageFont.load_default()

    draw.text((50, 50), text, font=font, fill=font_color)
    text_img = np.array(text_img)

    # Blend text onto cloth
    result = blend_text_cloth(cloth_image, text_img, x, y)

    return result


# Gradio Interface
interface = gr.Interface(
    fn=process_image,
    inputs=[
        gr.Image(type="pil", label="Upload Cloth Image"),
        gr.Textbox(label="Text to Blend", value="Sample Text"),
        gr.Slider(10, 100, step=2, label="Font Size", value=32),
        gr.Textbox(label="Font Color (RGB)", value="(255, 0, 0)"),
        gr.Slider(0, 1000, step=10, label="X Coordinate", value=50),
        gr.Slider(0, 1000, step=10, label="Y Coordinate", value=50),
    ],
    outputs=gr.Image(type="pil", label="Blended Output"),
    title="Advanced Text-Cloth Blending",
    description="Upload a cloth image and blend text naturally using advanced warping & blending techniques.",
)

# Launch the app
if __name__ == "__main__":
    interface.launch(server_name="0.0.0.0", server_port=7860)