import numpy as np
import matplotlib.pyplot as plt
import time

from tensorflow import keras
from tensorflow.keras import layers
app = FastAPI()


# Model / data parameters
num_classes = 9
input_shape = (28, 28, 3)
batch_size = 1000
epochs = 1

# Define baseline model
def baseline_model():

  # Create model
  model = keras.Sequential(
    [
        keras.Input(shape=input_shape),
        layers.Conv2D(64, kernel_size=(3, 3), activation="relu"),
        layers.Conv2D(128, kernel_size=(3, 3), activation="relu"),
        layers.Flatten(),
        layers.Dropout(0.5),
        layers.Dense(num_classes, activation="softmax"),
    ]
  )
  model.summary()

  # Compile model
  model.compile(loss="categorical_crossentropy", optimizer="adam", metrics=["accuracy"])

  return model

# Load Data
path = './pathmnist.npz'
with np.load(path) as data:
  x_train = data['train_images']
  y_train = data['train_labels']
  x_test = data['test_images']
  y_test = data['test_labels']
  x_val = data['val_images']
  y_val = data['val_labels']

# Show DataSet Images
for image in x_train:
  plt.imshow(image)
  plt.show()
  break

# Normalize images to the [0, 1] range
x_train = x_train.astype("float32") / 255
x_test = x_test.astype("float32") / 255
x_val = x_val.astype("float32") / 255

print("x_train shape:", x_train.shape)
print(x_train.shape[0], "train samples")
print(x_test.shape[0], "test samples")
print(x_val.shape[0], "test samples")

# Convert class vectors to binary class matrices
y_train = keras.utils.to_categorical(y_train, num_classes)
y_test = keras.utils.to_categorical(y_test, num_classes)
y_val = keras.utils.to_categorical(y_val, num_classes)

model = baseline_model()

# Fit model
#history = model.fit(x_train, y_train, batch_size=batch_size, epochs=epochs, validation_split=0.1)
inicio = time.time()
history = model.fit(x_train, y_train, batch_size=batch_size, epochs=epochs, validation_data=(x_val, y_val))
fin = time.time()
print(fin-inicio)

@app.get("/generate")
def generate(x: np.array):
    print(model.predict(x))