import tensorflow as tf
from tensorflow.keras.layers import Dense, Conv2D, MaxPool2D, Dropout, Add, Layer, Flatten, BatchNormalization, Activation
from tensorflow.keras.models import Model

class ResLayer(Layer):
    def __init__(self, filters, name = "Res_Layer"):
        super(ResLayer, self).__init__(name = name)
        self.filters = filters
        self.f1, self.f2, self.f3 = self.filters
        
        # Camino normal
        self.Conv_1 = Conv2D(filters = self.f1, kernel_size = (1, 1), strides = (1, 1))
        self.MaxPool_1 = MaxPool2D(pool_size = (2, 2))
        self.BatchNorm_1 = BatchNormalization()
        self.Activation_1 = Activation("relu")
        
        self.Conv_2 = Conv2D(filters = self.f2, kernel_size  = (1, 1), strides = (1, 1))
        self.BatchNorm_2 = BatchNormalization()
        self.Activation_2 = Activation("relu")
        
        self.Conv_3 = Conv2D(filters = self.f3, kernel_size = (1, 1), strides = (1, 1))
        self.BatchNorm_3 = BatchNormalization()
        
        # Camino corto
        self.Conv_4 = Conv2D(filters = self.f3, kernel_size = (1, 1), strides = (1, 1))
        self.MaxPool_2 = MaxPool2D(pool_size = (2, 2))
        
        self.Add = Add()
        self.Activation_3 = Activation("relu")
        
    def call(self, inputs):
        X_copy = inputs
        
        X = self.Conv_1(inputs)
        X = self.MaxPool_1(X)
        X = self.BatchNorm_1(X)
        X = self.Activation_1(X)
        
        X = self.Conv_2(X)
        X = self.BatchNorm_2(X)
        X = self.Activation_2(X)
        
        X = self.Conv_3(X)
        X = self.BatchNorm_3(X)
        
        X_copy = self.Conv_4(X_copy)
        X_copy = self.MaxPool_2(X_copy)
        
        outputs = self.Add([X, X_copy])
        outputs = self.Activation_3(outputs)
        
        return outputs
        
class ResNet(Model):
    def __init__(self, filters = [[64, 128, 256]], name = "ResNet"):
        super(ResNet, self).__init__(name = name)
        self.filters = filters
        self.nb_layers = tf.shape(self.filters)[0].numpy()
        self.res_layer = [ResLayer(filters)
                         for i, filters in enumerate(self.filters)]
        self.Flatten = Flatten()
        self.Dense_1 = Dense(units = 128, activation = "relu")
        self.dropout_1 = Dropout(rate = 0.2)
        self.Dense_2 = Dense(units = 64, activation = "relu")
        self.dropout_2 = Dropout(rate = 0.1)
        self.Dense_Out = Dense(units = 10, activation = "softmax")
        
    def call(self, inputs):
        outputs = inputs
        for i in range(self.nb_layers):
            outputs = self.res_layer[i](outputs)
        
        outputs = self.Flatten(outputs)
        outputs = self.Dense_1(outputs)
        outputs = self.dropout_1(outputs)
        outputs = self.Dense_2(outputs)
        outputs = self.dropout_2(outputs)
        outputs = self.Dense_Out(outputs)
        
        return outputs
        
model = ResNet()
model.build(input_shape = [None, 28, 28, 1])
model.load_weights("ResNet_Weights.tf")

import gradio as gr

def digit_recognition(img):
    img = img / 255.
    img = tf.expand_dims(img, axis = -1)
    img = tf.convert_to_tensor([img], dtype = tf.float32)
    prediction = model(img)
    prediction = tf.squeeze(prediction)
    return {"Cero": float(prediction[0]),
           "Uno": float(prediction[1]),
           "Dos": float(prediction[2]),
           "Tres": float(prediction[3]),
           "Cuatro": float(prediction[4]),
           "Cinco": float(prediction[5]),
           "Seis": float(prediction[6]),
           "Siete": float(prediction[7]),
           "Ocho": float(prediction[8]),
           "Nueve": float(prediction[9])}

app = gr.Interface(fn = digit_recognition, inputs = "sketchpad", outputs = "label", description = "Dibuja un número", title = "MNIST Digit Recognition")
app.launch(share = True)