import torch
import gradio as gr
from torchvision import transforms
from PIL import Image
import numpy as np
from model import model

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

resize_input = transforms.Resize((32, 32))
to_tensor = transforms.ToTensor()

def reconstruct_image(image):
    image = Image.fromarray(image).convert('RGB')
    image_32 = resize_input(image)
    image_tensor = to_tensor(image_32).unsqueeze(0).to(device)
    with torch.no_grad():
        mu, _ = model.encode(image_tensor)
        recon = model.decode(mu)
    recon_np = recon.squeeze(0).permute(1, 2, 0).cpu().numpy()
    recon_img = Image.fromarray((recon_np * 255).astype(np.uint8)).resize((512, 512))
    orig_resized = image_32.resize((512, 512))
    return orig_resized, recon_img

def get_interface():
    with gr.Blocks() as iface:
        gr.Markdown("## Encoding & Reconstruction")
        with gr.Row():
            input_image = gr.Image(label="Input (Downsampled to 32x32)", type="numpy")
            output_image = gr.Image(label="Reconstructed", type="pil")
        run_button = gr.Button("Run Reconstruction")

        run_button.click(fn=reconstruct_image, inputs=input_image, outputs=[input_image, output_image])

        examples = [
            ["example_images/image1.jpg"],
            ["example_images/image2.jpg"],
            ["example_images/image3.jpg"],
            ["example_images/image10.jpg"],
            ["example_images/image4.jpg"],
            ["example_images/image5.jpg"],
            ["example_images/image6.jpg"],
            ["example_images/image7.jpg"],
            ["example_images/image8.jpg"],
        ]

        gr.Examples(
            examples=examples,
            inputs=[input_image],
        )
    return iface