from huggingface_hub import from_pretrained_fastai
import gradio as gr
from fastai.vision.all import *
from tensorflow.keras.layers import BatchNormalization
from tensorflow.keras.layers import Conv2D
from tensorflow.keras.layers import Conv2DTranspose
from tensorflow.keras.layers import LeakyReLU
from tensorflow.keras.layers import Activation
from tensorflow.keras.layers import Flatten
from tensorflow.keras.layers import Dense
from tensorflow.keras.layers import Reshape
from tensorflow.keras.layers import Input
from tensorflow.keras.models import Model
from tensorflow.keras.optimizers.legacy import Adam
from tensorflow.keras.datasets import mnist
from tensorflow.keras import backend as K
from sklearn.model_selection import train_test_split
from google.colab.patches import cv2_imshow
import matplotlib.pyplot as plt
import random
import pickle
import cv2
import numpy as np

class ConvAutoencoder:
	@staticmethod
	def build(width, height, depth, filters=(32, 64), latentDim=16):

		# initialize the input shape to be "channels last" along with the channels
    # dimension itself
		inputShape = (height, width, depth)
		chanDim = -1

		# define the input to the encoder
		inputs = Input(shape=inputShape)
		x = inputs

		for f in filters:
			# apply a CONV => RELU => BN operation
			x = Conv2D(f, (3, 3), strides=2, padding="same")(x)
			x = LeakyReLU(alpha=0.2)(x)
			x = BatchNormalization(axis=chanDim)(x)

		# flatten the network and then construct our latent vector
		volumeSize = K.int_shape(x)
		x = Flatten()(x)
		latent = Dense(latentDim)(x)

		# build the encoder model
		encoder = Model(inputs, latent, name="encoder")

    # start building the decoder model which will accept the output of the
    # encoder as its inputs
		latentInputs = Input(shape=(latentDim,))
		x = Dense(np.prod(volumeSize[1:]))(latentInputs)
		x = Reshape((volumeSize[1], volumeSize[2], volumeSize[3]))(x)

		# loop over our number of filters again, but this time in reverse order
		for f in filters[::-1]:
			# apply a CONV_TRANSPOSE => RELU => BN operation
			x = Conv2DTranspose(f, (3, 3), strides=2, padding="same")(x)
			x = LeakyReLU(alpha=0.2)(x)
			x = BatchNormalization(axis=chanDim)(x)

		# apply a single CONV_TRANSPOSE layer used to recover the original depth of
    # the image
		x = Conv2DTranspose(depth, (3, 3), padding="same")(x)
		outputs = Activation("sigmoid")(x)

		# build the decoder model
		decoder = Model(latentInputs, outputs, name="decoder")

		# our autoencoder is the encoder + decoder
		autoencoder = Model(inputs, decoder(encoder(inputs)), name="autoencoder")

		return (encoder, decoder, autoencoder)


def build_unsupervised_dataset(data, labels, validLabel=1, anomalyLabel=3,
                               contam=0.01, seed=42):

	# grab all indexes of the supplied class label that are *truly* that particular
  # label, then grab the indexes of the image labels that will serve as "anomalies"
	validIdxs = np.where(labels == validLabel)[0]
	anomalyIdxs = np.where(labels == anomalyLabel)[0]

	random.shuffle(validIdxs)
	random.shuffle(anomalyIdxs)

	# compute the total number of anomaly data points to select
	i = int(len(validIdxs) * contam)
	anomalyIdxs = anomalyIdxs[:i]

	# use NumPy array indexing to extract both the valid images and "anomlay" images
	validImages = data[validIdxs]
	anomalyImages = data[anomalyIdxs]

	# stack the valid images and anomaly images together to form a single data
  # matrix and then shuffle the rows
	images = np.vstack([validImages, anomalyImages])
	np.random.seed(seed)
	np.random.shuffle(images)

	return images


EPOCHS = 20
INIT_LR = 1e-3
BS = 32


((trainX, trainY), (testX, testY)) = mnist.load_data()


images = build_unsupervised_dataset(trainX, trainY, validLabel=1, anomalyLabel=3,
                                    contam=0.01)


(encoder, decoder, autoencoder) = ConvAutoencoder.build(28, 28, 1)
opt = Adam(learning_rate=INIT_LR, decay=INIT_LR / EPOCHS)
autoencoder.compile(loss="mse", optimizer=opt)


H = autoencoder.fit(trainX, trainX, validation_data=(testX, testX), epochs=EPOCHS,
                    batch_size=BS)


# Definimos una función que se encarga de llevar a cabo las predicciones
def predict(img):
    img = PILImage.create(img)
    pred = autoencoder.predict(img)
    mse = np.mean((img - pred) ** 2)
    thresh = 0.02767541486024857
    if mse >= thresh:
      label = 'anomaly'
    else:
      label = 'not anomaly'
    return label
    
# Creamos la interfaz y la lanzamos. 
gr.Interface(fn=predict, inputs=gr.inputs.Image(shape=(128, 128)), outputs=label, examples=['uno.png','tres.png']).launch(share=False)