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

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

transform = transforms.Compose([
    transforms.Resize((32, 32)),
    transforms.ToTensor()
])

resize_output = transforms.Resize((512, 512))

def interpolate_vectors(v1, v2, num_steps):
    return [v1 * (1 - alpha) + v2 * alpha for alpha in np.linspace(0, 1, num_steps)]

def to_pil(img_tensor):
    img = img_tensor.squeeze(0).permute(1, 2, 0).cpu().numpy()
    img = (img * 255).clip(0, 255).astype(np.uint8)
    return Image.fromarray(img)

def interpolate_images_gif(img1, img2, num_interpolations=10, duration=100):
    img1 = Image.fromarray(img1).convert('RGB')
    img2 = Image.fromarray(img2).convert('RGB')
    img1_tensor = transform(img1).unsqueeze(0).to(device)
    img2_tensor = transform(img2).unsqueeze(0).to(device)

    with torch.no_grad():
        mu1, _ = model.encode(img1_tensor)
        mu2, _ = model.encode(img2_tensor)
        interpolated = interpolate_vectors(mu1, mu2, num_interpolations)
        decoded_images = []
        for z in interpolated:
            out = model.decode(z)
            img = to_pil(out)
            img_resized = resize_output(img)
            decoded_images.append(img_resized)

    tmp_file = tempfile.NamedTemporaryFile(suffix=".gif", delete=False)
    decoded_images[0].save(
        tmp_file.name,
        save_all=True,
        append_images=decoded_images[1:],
        duration=duration,
        loop=0
    )
    return tmp_file.name

def get_interface():
    with gr.Blocks() as iface:
        gr.Markdown("## Latent Space Interpolation (GIF Output)")
        with gr.Row():
            img1 = gr.Image(label="First Image", type="numpy")
            img2 = gr.Image(label="Second Image", type="numpy")
        slider_steps = gr.Slider(5, 30, value=10, step=1, label="Number of Interpolations")
        slider_duration = gr.Slider(50, 500, value=100, step=10, label="Duration per Frame (ms)")
        output_gif = gr.Image(label="Interpolation GIF")
        run_button = gr.Button("Interpolate")

        run_button.click(
            fn=interpolate_images_gif,
            inputs=[img1, img2, slider_steps, slider_duration],
            outputs=output_gif
        )

        examples = [
            ["example_images/image1.jpg", "example_images/image2.jpg", 10, 100],
            ["example_images/image3.jpg", "example_images/image4.jpg", 15, 150],
            ["example_images/image5.jpg", "example_images/image6.jpg", 20, 200],
            ["example_images/image7.jpg", "example_images/image8.jpg", 25, 250],
        ]

        gr.Examples(
            examples=examples,
            inputs=[img1, img2, slider_steps, slider_duration],
            outputs=output_gif,
            fn=interpolate_images_gif,
            cache_examples=False
        )
    return iface