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

import tensorflow as tf
from tensorflow import keras

###############################################################################
# 06-keras-tf_playground-spiral.py
# @title: Introduction aux réseaux de neurones - Portage de TensorFlow Playground vers Keras - Points en spirale
# @project: Mes scripts de ML
# @lang: fr
# @authors: Philippe Roy <philippe.roy@ac-grenoble.fr>
# @copyright: Copyright (C) 2023 Philippe Roy
# @license: GNU GPL
###############################################################################

###
# TensorFlow Playground : 
# - Site internet : https://playground.tensorflow.org
# - Dépôt git : https://github.com/tensorflow/playground
###

###
# Installation : 
# - pip3 install tensorflow
# - pip3 install keras
# - pip3 install pydot
# - pip3 install graphviz
###

###
# Commandes NumPy :
# - np.array : créer un tableau à partir d'une liste de listes
# - np.c_ : concatène les colonnes des tableaux
# - np.linspace : créer un tableau 1D de la valeur de début à la valeur de fin avec n valeurs
# - np.meshgrid : créer un tableau 2D avec l'ensemble des combinaisons allant des deux valeurs de début aux deux valeurs de fin
# - .reshape : reformater la tableau avec le nombre de lignes et le nombre de colonnes
###

###
# Commandes Keras :
# - keras.models.Sequential() : créer un modèle où les couches de neurones sont reliées séquentiellement (modèle simple)
# - keras.layers.Flatten : ajout d'une couche de formatage de mise à plat
# - keras.layers.Dense : ajout d'une couche de neurones
# - keras.backend.clear_session() : reset de la session
# - .compile : compilation du modèle
# - .fit : entrainement du modèle
# - .predict : prédiction du modèle
# - keras.utils.plot_model : créer le diagramme d'un modèle
###

###############################################################################
# Initialisation
###############################################################################

# Init de Tensorflow + Keras

# tf.__version__
# keras.__version__
# tf.config.list_physical_devices('GPU')

# Init du temps
t_debut = time.time()

# Init des plots
fig = plt.figure(layout="constrained", figsize=(15, 5))
fig.suptitle("Réseaux de neurones avec Keras - Classificateur d'images")
subfigs = fig.subfigures(1, 3)
model_ax = subfigs[0].subplots(1, 1)
apts_ax = subfigs[1].subplots(1, 1)
img_ax = subfigs[2].subplots(4, 8)

###############################################################################
# Observations
###############################################################################

# Observations d'apprentissage, de validation et de test
vetement = keras.datasets.fashion_mnist # Jeu de données Fashion MNIST 
(X, y), (X_test, y_test) = vetement.load_data()
X_train, y_train  = X[5000:]/255.0 , y[5000:]
X_valid, y_valid = X[:5000]/255.0 , y[:5000]
classes = ["Tshirt", "Pantalon", "Pull", "Robe", "Manteau", "Sandale", "Chemise", "Basket", "Sac", "Bottine"]

###############################################################################
# Phase d'apprentissage
###############################################################################

n = 30 # Nombre d'itérations (valeur par défaut : 30 , hyperparamètre)
eta = 0.01 # Taux d'appentissage (valeur par défaut dans Keras : 0.01, hyperparamètre)
lot=32 # Taille de lot (valeur par défaut dans Keras: 32 , hyperparamètre)
perte="sparse_categorical_crossentropy" # Type de perte (hyperparamètre)
#perte="mse" # Type de perte (hyperparamètre)

keras.backend.clear_session()
# np.random.seed(42)
# tf.random.set_seed(42)

model = keras.models.Sequential() # Modèle de reseau de neurones
model.add(keras.layers.Flatten(input_shape=[28, 28])) # Couche d'entrée : mise à plat des données d'entrée -> 1 node / pixel soit 784 (28x28)
model.add(keras.layers.Dense(300, activation="relu")) # Couche 1 : 300 nodes
# model.add(keras.layers.Dense(300, activation="relu")) # Couche 2 : 300 nodes -> passage de 100 à 300
# model.add(keras.layers.Dense(300, activation="relu")) # Couche 3 : 300 nodes -> ajout
model.add(keras.layers.Dense(100, activation="relu")) # Couche 4 : 100 nodes -> ajout
model.add(keras.layers.Dense(10, activation="softmax")) # Couche de sortie : 1 node par classe soit 10

# model.compile(loss="sparse_categorical_crossentropy", optimizer="sgd", metrics=["accuracy"])
optimiseur=keras.optimizers.SGD(learning_rate= eta)
model.compile(loss=perte, optimizer=optimiseur, metrics=["accuracy"]) # Compilation du modèle

apts = model.fit(X_train, y_train, epochs=n, batch_size=lot, validation_data=(X_valid, y_valid)) # Entrainement

###############################################################################
# Phase d'inférence
###############################################################################

# X_new=[]
# y_new=[]
# for i in range(8):
#     idx = np.random.randint(X_test.shape[0]) # Index aléatoire
#     X_new.append(X_test[idx:idx+1]/255.0)
#     y_new.append(y_test[idx:idx+1])

idx = np.random.randint(X_test.shape[0]-32) # Index aléatoire
print ("\n")
print ("Test sur les images de "+ str(idx) + " à "+ str(idx+32) + " sur un jeu de 10 000 images.")
X_new = X_test[idx:idx+32]
y_new = np.argmax(model.predict(X_new), axis=-1)
y_new_test= y_test[idx:idx+32]
print ("\n")

###############################################################################
# Résultats
###############################################################################

# Modèle
model_ax.set_title("Modèle")
keras.utils.plot_model(model, "model.png", show_shapes=True)
model_img=plt.imread("model.png")
model_ax.imshow(model_img)
model_ax.set_axis_off()
os.remove("model.png") # Supression du fichier temporaire

# Courbes d'apprentissage
apts_ax.set_title("Courbes d'apprentissage")
apts_ax.plot(apts.epoch, apts.history['loss'], 'b-', label="Perte - entrainement")
apts_ax.plot(apts.epoch, apts.history['val_loss'], 'r-', label="Perte - validation")
apts_ax.plot(apts.epoch, apts.history['accuracy'], 'b:', label="Précision - entrainement")
apts_ax.plot(apts.epoch, apts.history['val_accuracy'], 'r:', label="Précision - validation")
apts_ax.set(ylim=(0, 1))
apts_ax.set_xlabel("Époque")
apts_ax.legend()

# Prédictions
for i in range (8):
    for j in range (4):
        img_ax[j][i].imshow(X_new[i*2+j], cmap="binary", interpolation="nearest")
        img_ax[j][i].set_axis_off()
        if y_new[i*2+j] == y_new_test[i*2+j]:
            img_ax[j][i].set_title(classes[y_new[i*2+j]], fontsize=10)
        else:
            img_ax[j][i].set_title(classes[y_new[i*2+j]], fontsize=10, color="red")

plt.show()

# Performances
print ("Temps total  : "+str(time.time()-t_debut))