import torch
import torch.nn as nn
import gradio as gr
import numpy as np
from torchvision.transforms import ToPILImage

# === Generator Model Definition ===
class Generator(nn.Module):
    def __init__(self, latent_dim, label_dim, img_shape):
        super().__init__()
        input_dim = latent_dim + label_dim
        self.model = nn.Sequential(
            nn.Linear(input_dim, 128),
            nn.ReLU(True),
            nn.Linear(128, 256),
            nn.BatchNorm1d(256),
            nn.ReLU(True),
            nn.Linear(256, 512),
            nn.BatchNorm1d(512),
            nn.ReLU(True),
            nn.Linear(512, int(torch.prod(torch.tensor(img_shape)))),
            nn.Tanh()
        )
        self.img_shape = img_shape

    def forward(self, noise, labels):
        x = torch.cat((noise, labels), dim=1)
        img = self.model(x)
        return img.view(img.size(0), *self.img_shape)

# === Load Model ===
latent_dim = 100
label_dim = 10
img_shape = (1, 28, 28)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
generator = Generator(latent_dim, label_dim, img_shape).to(device)
generator.load_state_dict(torch.load("models/mnist_generator.pth", map_location=device))
generator.eval()

# === Generate Function ===
def generate_images(digit: str):
    digit = int(digit)
    z = torch.randn(5, latent_dim, device=device)
    labels = torch.eye(label_dim, device=device)[[digit] * 5]
    with torch.no_grad():
        gen_imgs = generator(z, labels).cpu()
    gen_imgs = (gen_imgs + 1) / 2  # Normalize to [0, 1]
    images = [ToPILImage()(img.squeeze(0)) for img in gen_imgs]
    return images

# === Gradio Interface ===
iface = gr.Interface(
    fn=generate_images,
    inputs=gr.Dropdown(choices=[str(i) for i in range(10)], label="Pick a digit"),
    outputs=[gr.Image(type="pil") for _ in range(5)],
    title="MNIST Digit Generator",
    description="Select a digit from 0–9 to generate 5 synthetic handwritten digits using a GAN trained on MNIST."
)

iface.launch()