import gradio as gr
from PIL import Image, ImageEnhance
import numpy as np

# Function to adjust contrast of the input image
def adjust_contrast(img, contrast_level):
    if img is None:
        return None
    enhancer = ImageEnhance.Contrast(img)
    return enhancer.enhance(contrast_level)

# Function to convert image to grayscale
def convert_to_grayscale(img):
    if img is None:
        return None
    return img.convert('L').convert('RGB')  # Grayscale and back to RGB for compatibility

# Function to apply Multiply blending with texture (simulating print effect)
def apply_texture(img, texture):
    if img is None or texture is None:
        return None

    # Resize texture to match image size
    texture = texture.resize(img.size, Image.LANCZOS)

    # Convert images to NumPy arrays
    gray_arr = np.array(img, dtype=np.float32) / 255.0  # Normalize grayscale image
    texture_arr = np.array(texture, dtype=np.float32) / 255.0  # Normalize texture

    # Apply Multiply blending: Multiply each pixel and normalize back to 0-255
    blended = gray_arr * texture_arr

    # Convert array back to image
    return Image.fromarray((blended * 255).astype(np.uint8))

# Gradio GUI
def main():
    with gr.Blocks() as demo:
        gr.Markdown("## 🏡 Simulated Printing on Wood Texture (Multiply Blending)")

        with gr.Row():
            input_img = gr.Image(label="Input Image", type="pil")
            texture_img = gr.Image(label="Upload Wood Texture", type="pil")

        contrast = gr.Slider(0.0, 1.0, value=0.4, label="Contrast Adjustment")

        with gr.Row():
            contrast_img = gr.Image(label="Contrast Adjusted Image", interactive=False)
            grayscale_img = gr.Image(label="Grayscale Image", interactive=False)
            output_img = gr.Image(label="Final Printed Effect", interactive=False)

        # Step 1: Adjust contrast
        def process_contrast(img, contrast_level):
            contrast_img = adjust_contrast(img, contrast_level)
            return contrast_img

        # Step 2: Convert contrast-adjusted image to grayscale
        def process_grayscale(img, contrast_level):
            contrast_img = adjust_contrast(img, contrast_level)
            grayscale_img = convert_to_grayscale(contrast_img)
            return contrast_img, grayscale_img

        # Step 3: Apply Multiply blending with texture
        def process_texture(img, texture, contrast_level):
            contrast_img = adjust_contrast(img, contrast_level)  # Adjust contrast first
            grayscale_img = convert_to_grayscale(contrast_img)  # Convert to grayscale
            result = apply_texture(grayscale_img, texture)  # Multiply blend with texture
            return contrast_img, grayscale_img, result

        # Update contrast-adjusted image
        input_img.change(process_contrast, inputs=[input_img, contrast], outputs=contrast_img)

        # Update grayscale image
        input_img.change(process_grayscale, inputs=[input_img, contrast], outputs=[contrast_img, grayscale_img])
        contrast.change(process_grayscale, inputs=[input_img, contrast], outputs=[contrast_img, grayscale_img])

        # Update final processed image (Multiply blend with texture)
        input_img.change(process_texture, inputs=[input_img, texture_img, contrast], outputs=[contrast_img, grayscale_img, output_img])
        texture_img.change(process_texture, inputs=[input_img, texture_img, contrast], outputs=[contrast_img, grayscale_img, output_img])
        contrast.change(process_texture, inputs=[input_img, texture_img, contrast], outputs=[contrast_img, grayscale_img, output_img])

    demo.launch()

if __name__ == "__main__":
    main()