import gradio as gr
import tensorflow as tf
from tensorflow import keras
import matplotlib.pyplot as plt
from math import sqrt, ceil

from huggingface_hub import from_pretrained_keras

import numpy as np


model = from_pretrained_keras("keras-io/conditional-gan")

latent_dim = 128

def generate_latent_points(digit, latent_dim, n_samples, n_classes=10):
	# generate points in the latent space
	random_latent_vectors = tf.random.normal(shape=(n_samples, latent_dim))
	labels = tf.keras.utils.to_categorical([digit for _ in range(n_samples)], n_classes)
	return tf.concat([random_latent_vectors, labels], 1)

def create_digit_samples(digit, n_samples, latent_dim=latent_dim):
    random_vector_labels = generate_latent_points(digit, latent_dim, n_samples)
    examples = cgan_generator.predict(random_vector_labels)
    examples = examples * 255.0
    size = ceil(sqrt(n_samples))
    digit_images = np.zeros((28*size, 28*size))
    n = 0
    for i in range(size):
        for j in range(size):
            if n == n_samples:
                break
            digit_images[i* 28 : (i+1)*28, j*28 : (j+1)*28] = examples[n, :, :, 0]
            n += 1

    return digit_images

description = "This model is based on the example created here: https://keras.io/examples/generative/conditional_gan/"

title = "Conditional GAN for MNIST"

examples = [[1, 10], [3, 5], [5, 15]]


iface = gr.Interface(
    fn = create_digit_samples,
    inputs = ["number", "number"],
    outputs = [gradio.outputs.Image(invert_colors=True, type="numpy", label="Samples for given digit")],
    examples = examples,
    description = description,
    title = title
    )

iface.launch()