import os, time
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
from matplotlib.colors import ListedColormap, LinearSegmentedColormap

import tensorflow as tf
from tensorflow import keras

###############################################################################
# 04-keras-tf_playground-xor.py
# @title: Introduction aux réseaux de neurones - Portage de TensorFlow Playground vers Keras - Points en XOR
# @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)
# - model.add : ajout d'une couche
# - keras.layers.Flatten : couche de formatage de mise à plat
# - keras.layers.Dense : couche de neurones
# - keras.backend.clear_session() : reset de la session
# - model.compile : compilation du modèle
# - model.fit : entrainement du modèle
# - model.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(figsize=(15, 5))
fig.suptitle("Réseaux de neurones avec Keras - Classificateur : XOR")
model_ax = fig.add_subplot(131) # Modèle
apts_ax = fig.add_subplot(132) # Courbes d'apprentissage
donnees_ax = fig.add_subplot(133) # Observations : x1,x2 et cibles : y

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

# Observations d'apprentissage
m = 1000 # Nombre d'observations
bg = 1 # Quantité du bruit gaussien # FIXME : pas en place
x1 = np.random.uniform(-5, 5, m) # Liste des observations x1
x2 = np.random.uniform(-5, 5, m) # Liste des observations x2
y = np.empty(m)  # Liste des observations cible (XOR)
marge = 0.15
for i in range (m):
    if x1[i]>= 0 and x1[i] < marge: x1[i] += marge
    if x1[i]<0 and x1[i] > -marge: x1[i] -= marge
    if x2[i]>= 0 and x2[i] < marge: x2[i] += marge
    if x2[i]<0 and x2[i] > -marge: x2[i] -= marge
    if x1[i]*x2[i]<0:
        y[i]=0
    else:
        y[i]=1

# Split en observations d'entrainement et de validation
test_size=0.1 # Ratio du lot de test
m_train = int(np.round(m*(1-test_size)))
x1_train, x2_train, y_train  = x1[:m_train], x2[:m_train], y[:m_train] # Jeu d'entrainement
x1_valid, x2_valid, y_valid  = x1[m_train:], x2[m_train:], y[m_train:] # Jeu de validation
X_train  = np.c_[x1_train, x2_train]
X_valid  = np.c_[x1_valid, x2_valid]

# Plots
donnees_ax.plot(x1_train[y_train==1], x2_train[y_train==1], "o", markerfacecolor="tab:blue", markeredgecolor='white', markeredgewidth=0.75)
donnees_ax.plot(x1_train[y_train==0], x2_train[y_train==0], "o" , markerfacecolor="tab:orange", markeredgecolor='white', markeredgewidth=0.75)
donnees_ax.plot(x1_valid[y_valid==1], x2_valid[y_valid==1], "o", markerfacecolor='tab:blue', markeredgecolor='black')
donnees_ax.plot(x1_valid[y_valid==0], x2_valid[y_valid==0], "o", markerfacecolor='tab:orange', markeredgecolor='black')

# Nouvelles observations
m_new = 100 # Résolution par axes
x1_new=np.linspace(-6, 6, m_new).reshape(-1, 1)
x2_new=np.linspace(-6, 6, m_new).reshape(-1, 1)
x1_new_mg, x2_new_mg = np.meshgrid(x1_new, x2_new)
X_new = np.c_[x1_new_mg.ravel(), x2_new_mg.ravel()]

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

n = 200 # Nombre d'itérations (valeur par défaut : 200 , 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)
# perte='mean_absolute_error'

keras.backend.clear_session()
model = keras.models.Sequential() # Modèle de réseau de neurones
model.add(keras.layers.Dense(4, input_dim=2, activation="relu")) # Couche 1 : 4 nodes
model.add(keras.layers.Dense(4, activation="relu")) # Couche 2 : 4 nodes
model.add(keras.layers.Dense(1, activation="sigmoid")) # Couche de sortie : 1 node par classe

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
###############################################################################

y_predict=model.predict(X_new) # Prédiction
y_predict_map = y_predict.reshape(x1_new_mg.shape)

###############################################################################
# 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.05, 1.05))
apts_ax.set_xlabel("Époque")
apts_ax.legend()

# Plot des données
donnees_ax.set_title("Données")
new_colors = ["tab:orange", "white", "tab:blue"]
new_cmap = LinearSegmentedColormap.from_list("mycmap", new_colors) # FIXME : faire un dégradé
cc = donnees_ax.contourf(x1_new_mg, x2_new_mg, y_predict_map, cmap=new_cmap)
donnees_ax.set_xticks([-5,0,5])
donnees_ax.set_yticks([-5,0,5])
donnees_ax.set_xlabel(r'$x_1$')
donnees_ax.set_ylabel(r'$x_2$', rotation=0)
donnees_ax.set(xlim=(-5.25, 5.25), ylim=(-5.25, 5.25))
fig.colorbar(cc, ax=donnees_ax)
plt.show()

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