import gradio as gr
import tensorflow as tf
import numpy as np
import cv2
import os
from main import tensor_to_image,StyleContentModel, run_style_transfer

# Define layers and instantiate the model once globally
content_layers = ['block5_conv2'] 
style_layers = ['block1_conv1', 'block2_conv1', 'block3_conv1', 'block4_conv1', 'block5_conv1']
extractor = StyleContentModel(style_layers, content_layers)

def style_transfer_wrapper(content_img_np, style_img_np):
    """
    A wrapper to handle I/O for the Gradio interface.
    Saves numpy arrays to temp files to use with the main function.
    """
    if content_img_np is None or style_img_np is None:
        return None # Return None if either image is missing

    # Save numpy arrays to temporary files
    content_path = "temp_content.jpg"
    style_path = "temp_style.jpg"
    
    # Gradio provides RGB, but cv2 saves in BGR order
    cv2.imwrite(content_path, cv2.cvtColor(content_img_np, cv2.COLOR_RGB2BGR))
    cv2.imwrite(style_path, cv2.cvtColor(style_img_np, cv2.COLOR_RGB2BGR))

    # Run the main process (using fewer iterations for a faster demo)
    final_tensor = run_style_transfer(content_path, style_path, iterations=500)
    
    # Convert tensor to displayable image
    final_image = tensor_to_image(final_tensor)
    
    # Clean up temporary files
    os.remove(content_path)
    os.remove(style_path)

    return final_image

## 4. GRADIO UI DEFINITION ##

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🎨 Neural Style Transfer")
    gr.Markdown("Combine the content of one image with the artistic style of another. This demo uses a VGG19 model. Processing can take a minute, especially on CPU.")

    with gr.Row():
        content_img = gr.Image(label="Content Image", type="numpy")
        style_img = gr.Image(label="Style Image", type="numpy")
    
    run_button = gr.Button("Generate Image", variant="primary")
    
    output_img = gr.Image(label="Result")
    
    run_button.click(
        fn=style_transfer_wrapper, 
        inputs=[content_img, style_img], 
        outputs=output_img
    )
    
    gr.Markdown("---")
    gr.Markdown("Based on the paper '[A Neural Algorithm of Artistic Style](https://arxiv.org/abs/1508.06576)' by Gatys et al.")

# Launch the Gradio app
demo.launch(share=True,debug=True)