import streamlit as st
import torch
from torchvision.utils import save_image
import matplotlib.pyplot as plt
from PIL import Image
import io

# Importa il modello del generatore
class EfficientGenerator(torch.nn.Module):
    def __init__(self, latent_dim):
        super(EfficientGenerator, self).__init__()
        self.init_size = 64 // 16
        self.l1 = torch.nn.Sequential(torch.nn.Linear(latent_dim, 128 * self.init_size ** 2))

        self.conv_blocks = torch.nn.Sequential(
            torch.nn.BatchNorm2d(128),
            torch.nn.Upsample(scale_factor=2),
            torch.nn.Conv2d(128, 128, 3, stride=1, padding=1),
            torch.nn.BatchNorm2d(128, 0.8),
            torch.nn.ReLU(inplace=True),
            torch.nn.Upsample(scale_factor=2),
            torch.nn.Conv2d(128, 64, 3, stride=1, padding=1),
            torch.nn.BatchNorm2d(64, 0.8),
            torch.nn.ReLU(inplace=True),
            torch.nn.Upsample(scale_factor=2),
            torch.nn.Conv2d(64, 32, 3, stride=1, padding=1),
            torch.nn.BatchNorm2d(32, 0.8),
            torch.nn.ReLU(inplace=True),
            torch.nn.Upsample(scale_factor=2),
            torch.nn.Conv2d(32, 3, 3, stride=1, padding=1),
            torch.nn.Tanh(),
        )

    def forward(self, z):
        out = self.l1(z)
        out = out.view(out.shape[0], 128, self.init_size, self.init_size)
        img = self.conv_blocks(out)
        return img

# Carica il modello salvato
def load_model(model_path):
    latent_dim = 32
    model = EfficientGenerator(latent_dim)
    checkpoint = torch.load(model_path, map_location=torch.device('cpu'))
    model.load_state_dict(checkpoint['generator_state_dict'])
    model.eval()
    return model

# Funzione per generare un'immagine
def generate_image(model):
    latent_dim = 32
    z = torch.randn(1, latent_dim)
    with torch.no_grad():
        gen_img = model(z).detach().cpu()
    img = gen_img.squeeze(0).permute(1, 2, 0)
    img = (img + 1) / 2  # Denormalizza
    return img

# Interfaccia Streamlit
st.title("Generazione di Immagini GAN")
st.write("Carica un modello GAN addestrato per generare immagini di gatti.")

uploaded_model = st.file_uploader("Carica il modello", type=["pth"])

if uploaded_model is not None:
    # Carica il modello dal file caricato
    with open("temp_model.pth", "wb") as f:
        f.write(uploaded_model.getbuffer())
    
    model = load_model("temp_model.pth")

    st.write("Modello caricato correttamente!")

    if st.button("Genera Immagine"):
        img = generate_image(model)
        fig, ax = plt.subplots()
        ax.imshow(img)
        ax.axis("off")
        st.pyplot(fig)

        # Salva l'immagine generata
        buffer = io.BytesIO()
        save_image(img.permute(2, 0, 1), buffer, format="PNG", normalize=True)
        st.download_button(
            label="Scarica Immagine",
            data=buffer.getvalue(),
            file_name="generated_image.png",
            mime="image/png"
        )