import streamlit as st
import torch
import torch.nn as nn
import numpy as np
from PIL import Image

# --- Configuration ---
# These must match the training script parameters
LATENT_DIM = 100
N_CLASSES = 10
IMG_SIZE = 28
CHANNELS = 1

# Use CPU for inference as the deployment environment may not have a GPU
DEVICE = torch.device('cpu')

# --- Model Architecture ---
# The model class must be defined exactly as it was during training
# so that we can load the saved weights (state_dict).

class Generator(nn.Module):
    def __init__(self):
        super(Generator, self).__init__()
        self.label_embedding = nn.Embedding(N_CLASSES, N_CLASSES)

        self.model = nn.Sequential(
            nn.Linear(LATENT_DIM + N_CLASSES, 256),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear(256, 512),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear(512, 1024),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear(1024, IMG_SIZE * IMG_SIZE * CHANNELS),
            nn.Tanh()
        )

    def forward(self, z, labels):
        label_emb = self.label_embedding(labels)
        gen_input = torch.cat((z, label_emb), -1)
        img = self.model(gen_input)
        img = img.view(img.size(0), CHANNELS, IMG_SIZE, IMG_SIZE)
        return img

# --- Helper Function to Load the Model ---
# Use st.cache_resource to load the model only once
@st.cache_resource
def load_model(model_path):
    """Loads the pre-trained generator model."""
    model = Generator().to(DEVICE)
    # Load the state dictionary. map_location ensures it loads on CPU.
    model.load_state_dict(torch.load(model_path, map_location=DEVICE))
    model.eval()  # Set the model to evaluation mode
    return model

# --- Image Generation Function ---
def generate_images(model, digit_to_generate, num_images=5):
    """Generates a specified number of images for a given digit."""
    with torch.no_grad():
        # Create random noise vectors (latent space)
        z = torch.randn(num_images, LATENT_DIM, device=DEVICE)
        
        # Create labels for the desired digit
        labels = torch.LongTensor([digit_to_generate] * num_images).to(DEVICE)
        
        # Generate images
        generated_imgs_tensor = model(z, labels)
        
        # Post-process images for display
        # 1. Move to CPU and convert to numpy
        # 2. Denormalize from [-1, 1] to [0, 1]
        # 3. Reshape from (N, C, H, W) to (N, H, W)
        generated_imgs_np = generated_imgs_tensor.cpu().numpy()
        generated_imgs_np = 0.5 * generated_imgs_np + 0.5  # Denormalize
        generated_imgs_np = generated_imgs_np.squeeze() # Remove channel dim
        
        return generated_imgs_np


# --- Streamlit Web App UI ---

st.set_page_config(page_title="Digit Generator", layout="wide")

st.title("✍️ Handwritten Digit Generator")
st.write(
    "This web app uses a Conditional Generative Adversarial Network (cGAN) "
    "trained on the MNIST dataset to generate new images of handwritten digits. "
    "Select a digit from the sidebar and click 'Generate'!"
)

# --- Sidebar Controls ---
st.sidebar.header("Controls")
digit_to_generate = st.sidebar.selectbox(
    "Select a digit (0-9):",
    options=list(range(10))
)

generate_button = st.sidebar.button("Generate Images", type="primary")

# --- Main Page Display ---
if generate_button:
    # Load the generator model
    try:
        generator = load_model("src/cgan_generator.pth")
        
        st.subheader(f"Generating 5 images for the digit: {digit_to_generate}")
        
        with st.spinner("🧠 Model is thinking..."):
            # Generate the images
            images_to_display = generate_images(generator, digit_to_generate, num_images=5)

            # Create 5 columns to display images side-by-side
            cols = st.columns(5)
            for i, image_array in enumerate(images_to_display):
                with cols[i]:
                    st.image(
                        image_array,
                        caption=f"Generated Image {i+1}",
                        width=150, # Control the display size
                        use_column_width='auto'
                    )
        
        st.success("Done!")

    except FileNotFoundError:
        st.error(
            "Model file 'cgan_generator.pth' not found. "
            "Please make sure the trained model file is in the same directory as this script."
        )
    except Exception as e:
        st.error(f"An error occurred: {e}")
else:
    st.info("Select a digit and click the 'Generate Images' button in the sidebar to start.")