import gradio as gr
import torch
import os
from torchvision import transforms
from PIL import Image
from src.lightning_module import StyleTransferModule

MODEL_URL = "https://huggingface.co/Michal-Raszkowski/adain-style-transfer/resolve/main/style-transfer-best-v2.ckpt?download=true"
CHECKPOINT_PATH = "model.ckpt"

def download_model_if_missing():
    if not os.path.exists(CHECKPOINT_PATH):
        print("Downloading model checkpoint...")
        torch.hub.download_url_to_file(MODEL_URL, CHECKPOINT_PATH)

def load_model():
    download_model_if_missing()
    # Loading to CPU by default; change to "cuda" if GPU is available
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = StyleTransferModule.load_from_checkpoint(CHECKPOINT_PATH, map_location=device)
    model.eval()
    return model, device

# Initialize model and get target device
model, device = load_model()

def stylize(content_image, style_image, alpha):
    if content_image is None or style_image is None:
        return None
    
    transform = transforms.Compose([
        transforms.Resize((512, 512)),
        transforms.ToTensor()
    ])
    
    # Transform and push tensors to the correct device (CPU/GPU)
    c = transform(content_image).unsqueeze(0).to(device)
    s = transform(style_image).unsqueeze(0).to(device)
    
    with torch.no_grad():
        generated_tensor, _ = model(c, s, alpha=alpha)
    
    # Bring back to CPU, clamp values, remove batch dimension, and convert to PIL
    generated_tensor = torch.clamp(generated_tensor, 0, 1).cpu().squeeze(0)
    result_image = transforms.ToPILImage()(generated_tensor)
    
    return result_image

# Build the Gradio Interface
with gr.Blocks(title="Style Transfer Demo", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# Neural Style Transfer")
    gr.Markdown("Upload a content image and a style image to combine them using AdaIN.")
    
    with gr.Row():
        with gr.Column():
            input_content = gr.Image(label="Content Image", type="pil", height=300)
            input_style = gr.Image(label="Style Image", type="pil", height=300)
            slider = gr.Slider(minimum=0.0, maximum=1.0, value=1.0, step=0.1, label="Style Strength")
            btn = gr.Button("Generate", variant="primary")
        
        with gr.Column():
            output = gr.Image(label="Output Image", type="pil")
    
    # Set up the click event listener
    btn.click(
        fn=stylize, 
        inputs=[input_content, input_style, slider], 
        outputs=output
    )
    
    # Optional: Uncomment if you want to include default examples
    # gr.Examples(
    #     examples=[["examples/c.jpg", "examples/s.jpg", 1.0]], 
    #     inputs=[input_content, input_style, slider]
    # )

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