import cv2
import numpy as np
import gradio as gr
from PIL import Image, ImageDraw, ImageFont

def generate_text_image(text, style, width, height):
    """Generate an image of the input text with the selected style."""
    font_path = {
        "Bold": "arialbd.ttf",
        "Italic": "ariali.ttf",
        "Graffiti": "graffiti.ttf",  # Provide custom font if needed
        "Calligraphy": "calligraphy.ttf"  # Provide custom font if needed
    }

    img = Image.new("RGBA", (width, height), (255, 255, 255, 0))
    draw = ImageDraw.Draw(img)

    try:
        font = ImageFont.truetype(font_path[style], size=80)
    except:
        font = ImageFont.load_default()

    text_bbox = draw.textbbox((0, 0), text, font=font)
    text_width = text_bbox[2] - text_bbox[0]
    text_height = text_bbox[3] - text_bbox[1]

    position = ((width - text_width) // 2, (height - text_height) // 2)

    draw.text(position, text, font=font, fill=(0, 0, 0, 255))

    return img

def apply_displacement_map(text_img, clothing_img, strength=20):
    """Apply displacement map to blend text onto clothing."""
    gray = cv2.cvtColor(clothing_img, cv2.COLOR_BGR2GRAY)
    grad_x = cv2.Sobel(gray, cv2.CV_32F, 1, 0, ksize=5)
    grad_y = cv2.Sobel(gray, cv2.CV_32F, 0, 1, ksize=5)

    grad_x = cv2.normalize(grad_x, None, 0, 1, cv2.NORM_MINMAX)
    grad_y = cv2.normalize(grad_y, None, 0, 1, cv2.NORM_MINMAX)

    displacement_map = np.zeros_like(clothing_img, dtype=np.float32)
    displacement_map[:, :, 0] = grad_x * strength
    displacement_map[:, :, 1] = grad_y * strength

    text_warped = cv2.remap(text_img, 
                            displacement_map[:, :, 0].astype(np.float32),
                            displacement_map[:, :, 1].astype(np.float32), 
                            interpolation=cv2.INTER_LINEAR)

    return text_warped

def overlay_text_on_clothing(clothing_image, text_input, style, strength=20, alpha=0.7):
    """Blend generated text onto the clothing image."""
    clothing_img = cv2.imread(clothing_image)  # Read image from file path

    # Generate text image dynamically
    text_img_pil = generate_text_image(text_input, style, clothing_img.shape[1], clothing_img.shape[0])
    text_img = cv2.cvtColor(np.array(text_img_pil), cv2.COLOR_RGBA2BGRA)

    alpha_channel = text_img[:, :, 3] / 255.0
    text_img = text_img[:, :, :3]

    text_warped = apply_displacement_map(text_img, clothing_img, strength)

    for c in range(3):
        clothing_img[:, :, c] = (1 - alpha_channel) * clothing_img[:, :, c] + alpha_channel * text_warped[:, :, c]

    # Convert the final blended image back to a PIL Image and return
    blended_img = Image.fromarray(cv2.cvtColor(clothing_img, cv2.COLOR_BGR2RGB))
    return blended_img

interface = gr.Interface(
    fn=overlay_text_on_clothing,
    inputs=[
        gr.Image(type="filepath", label="Upload Clothing Image", interactive=True),
        gr.Textbox(label="Enter Text for Design"),
        gr.Radio(["Bold", "Italic", "Graffiti", "Calligraphy"], label="Select Style", value="Bold"),
        gr.Slider(10, 50, step=5, value=20, label="Displacement Strength"),
        gr.Slider(0.1, 1.0, step=0.1, value=0.7, label="Alpha Blending")
    ],
    outputs=gr.Image(type="pil", label="Final Design"),
    title="AI-Powered Clothing Text Overlay",
    description="Upload a clothing image, enter a text design, and select a style to blend them onto clothing with displacement mapping.",
    allow_flagging="never"
)

if __name__ == "__main__":
    interface.launch()