# Réseau inspiré de http://yann.lecun.com/exdb/publis/pdf/lecun-98.pdf


from resnet18 import ResNet18
import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
# Pour les utilisateurs de MacOS  (pour utiliser plt & keras en même temps)
import os
#os.environ['KMP_DUPLICATE_LIB_OK']='True'


def displayConvFilers(model, layer_name, num_filter=4, filter_size=(3,3), num_channel=0, fig_size=(2,2)):
    
    layer_dict = dict([(layer.name, layer) for layer in mnist_model.layers])
    
    weight, biais = layer_dict[layer_name].get_weights()
    print(weight.shape) 
    plt.figure(figsize=fig_size)
    for i in range(num_filter):
        plt.subplot(fig_size[0],fig_size[1],i+1)
        plt.xticks([])
        plt.yticks([])
        plt.grid(False)
        vis = np.reshape(weight[:,:,num_channel,i],filter_size)
        plt.imshow(vis, cmap=plt.cm.binary)
    plt.show()


def snake(x):
    return x + tf.sin(x)**2


## Chargement et normalisation des données
resnet18 = tf.keras.datasets.cifar10
(train_images, train_labels), (test_images, test_labels) = resnet18.load_data()
train_images = train_images / 255.0
test_images = test_images / 255.0

# POUR LES CNN : On rajoute une dimension pour spécifier qu'il s'agit d'imgages en NdG
train_images = train_images.reshape(50000,32,32,3)
test_images = test_images.reshape(10000,32,32,3)


# One hot encoding
train_labels = tf.keras.utils.to_categorical(train_labels)
test_labels = tf.keras.utils.to_categorical(test_labels)

filter_size_conv1 = (3,3)

model = ResNet18(10)
model.build(input_shape = (None,32,32,3))


# expliquer le nombre de paramètre de ce réseau
print(model.summary())


sgd = tf.keras.optimizers.Adam()

model.compile(sgd, loss='categorical_crossentropy', metrics=['accuracy'])

# On visualise avant l'entrainement
'''
displayConvFilers(mnist_model, 'conv2d', 
                    num_filter=6, 
                    filter_size=filter_size_conv1, 
                    num_channel=0, 
                    fig_size=(2,3)
                )
'''


model.fit(train_images,
         train_labels,
         batch_size=64,
         epochs=4
         )

## Evaluation du modèle 
test_loss, test_acc = model.evaluate(test_images, test_labels)
print('Test accuracy:', test_acc)


'''
displayConvFilers(mnist_model,
                    'conv2d', 
                    num_filter=6, 
                    filter_size=filter_size_conv1, 
                    num_channel=0, 
                    fig_size=(2,3)
                )
'''

'''
displayConvFilers(mnist_model, 'conv2d_1', 
                    num_filter=16, 
                    filter_size=(5,5), 
                    num_channel=1, 
                    fig_size=(4,4)
                )
'''