##### -------------------------------------------------- Librerías -------------------------------------------------- #####

import os
import sys
import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt

sys.path.append("../tools")  # Herramientas propias de MantaFlow.
import uniio  # Lectura de ficheros .uni

from tensorflow.keras.layers import Input, Dropout, Conv2D, Conv2DTranspose, BatchNormalization, Flatten, Activation, Reshape
from tensorflow.keras.layers import LeakyReLU
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.models import Model, load_model
from tensorflow.keras.callbacks import ModelCheckpoint

import h5py

##### -------------------------------------------------- Selección GPU ---------------------------------------------- #####

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"

os.environ["CUDA_VISIBLE_DEVICES"] = "0"

##### -------------------------------------------------- Hiperparámetros -------------------------------------------- #####

num_sims = 2000  # Índice máximo escenas. 
num_scenes = num_sims - 1000  # Número de escenas.
frames = 200  # Frames por escena.

ae_epochs = 5  # Epochs para entrenamiento normal.
ae_epochs_list = 50  # Epochs para cada batch size de la lista.
pre_epochs = 1  # Epochs de preentrenamiento.

ae_batch_multiple = False  # Probar distintos batch sizes, comparar diferencias.
pre_training = True  # Realizar preentrenamiento.

ae_batch_list = [1024, 512, 256, 128, 64, 32]  # Posibles batch sizes de prueba.
ae_batch_size = 64  # Batch size oficial.
pre_batch_size = 128  # Bacth size para preentrenamiento.

##### -------------------------------------------------- Carga de datos --------------------------------------------- #####

np.random.seed(13) 

base_path = "../data"

print("Cargamos {} escenas, con {} frames cada una.".format(num_scenes, frames))
print("Trabajamos con un total de {} frames.".format((num_scenes) * frames))

# ------------------------------------------------------------------------------------------------------------------- #

base_path = "../data"

densities = []

for sim in range(1000, num_sims):
    
    if os.path.exists("%s/simSimple_%04d" % (base_path, sim)):  # Comprueba la existencia de las carpetas (cada una 200 frames de datos).
        
        for i in range(0, frames):
            
            filename = "%s/simSimple_%04d/density_%04d.uni"  # Nombre de cada frame (densidad).
            uni_path = filename % (base_path, sim, i)  # 200 frames por sim, rellena parametros de la ruta.
            header, content = uniio.readUni(uni_path)  # Devuelve un array Numpy [Z, Y, X, C].
            
            h = header["dimX"]
            w = header["dimY"]
            
            arr = content[:, ::-1, :, :]  # Cambia el orden del eje Y.
            arr = np.reshape(arr, [w, h, 1])  # Deshecha el eje Z.
            
            densities.append(arr)

# ------------------------------------------------------------------------------------------------------------------- #

load_num = len(densities)

if load_num < 2 * frames:
    
    print("Error - Usa al menos dos simulaciones completas")
    
    exit(True)

densities = np.reshape(densities, (len(densities), 64, 64, 1))  # Reconvierte la lista en array de Numpy.

print("Forma del array: {}".format(densities.shape))
print("Dimensiones del array: {}".format(densities.ndim))
print("Número de pixels en total: {}".format(densities.size))

# ------------------------------------------------------------------------------------------------------------------- #

vali_set_size = max(200, int(load_num * 0.1))  # Al menos una sim completa o el 10% de los datos.

vali_data = densities[load_num - vali_set_size : load_num, :]  # "load_num" datos del final de "densities".
train_data = densities[0 : load_num - vali_set_size, :]  # El resto de datos de "densities".

print("Separamos en {} frames de entrenamiento y {} frames de validación.".format(train_data.shape[0], vali_data.shape[0]))

# ------------------------------------------------------------------------------------------------------------------- #

train_data = np.reshape(train_data, (len(train_data), 64, 64, 1))  # Reconvertimos a arrays de Numpy.
vali_data = np.reshape(vali_data, (len(vali_data), 64, 64, 1))

print("Forma del set de entrenamiento: {}".format(train_data.shape))
print("Forma del set de validación: {}".format(vali_data.shape))

##### -------------------------------------------------- Funciones -------------------------------------------------- #####

def training_plot(network_train, epochs, batch_size, dropout, identification, loss, metric):
    
    plot_epochs = range(epochs)
    plot_loss = network_train.history["loss"]
    plot_val_loss = network_train.history["val_loss"]
    plot_metric = network_train.history[metric]
    plot_val_metric = network_train.history["val_" + metric]

    plt.figure(figsize = (15, 5))

    ax = plt.subplot(1, 2, 1)
    plt.plot(plot_epochs, plot_loss, label = loss.upper())
    plt.plot(plot_epochs, plot_val_loss, label = "Validation " + loss.upper())
    plt.legend()
    plt.xlabel("Epoch")
    plt.ylabel(loss.upper())

    ax = plt.subplot(1, 2, 2)
    plt.plot(plot_epochs, plot_metric, label = metric.upper())
    plt.plot(plot_epochs, plot_val_metric, label = "Validation " + metric.upper())
    plt.legend()
    plt.xlabel("Epoch")
    plt.ylabel(metric.upper())

    if dropout > 0.0:

        plt.savefig("../plots/model_ae_" + identification + "_DO-" + str(dropout * 100) + "_BS-" + str(batch_size) + ".png")

    else:

        plt.savefig("../plots/model_ae_" + identification + "_BS-" + str(batch_size) + ".png")

##### -------------------------------------------------- Autoencoder 2D --------------------------------------------- #####

feature_multiplier = 8  # Controla la cantidad de filtros de convolución utilizados por el autoencoder, y la dimension del espacio latente de la red. 
surface_kernel_size = 4  # Matriz 4x4
kernel_size = 2  # Matriz 2x2
dropout = 0.0  # Porcentaje de nodos que apagar mediante dropout.
init_func = "glorot_normal"  # Función de inicialización de los pesos de la red neuronal.

input_shape = (train_data.shape[1], 
               train_data.shape[2], 
               train_data.shape[3])

# ------------------------------------------------------------------------------------------------------------------- #

layer_conv = []

### Conv 1 ###

# Input #

conv1_input_shape = input_shape

conv1_input = Input(shape = conv1_input_shape)

x = conv1_input

# Layer 0 #

x = Conv2D(filters = feature_multiplier * 1, 
            kernel_size = surface_kernel_size,
            strides = 1,
            padding = "same",
            kernel_initializer = init_func)(x)
x = LeakyReLU(alpha = 0.2)(x)
x = BatchNormalization()(x)

# Layer 1 #

x = Conv2D(filters = feature_multiplier * 1, 
            kernel_size = surface_kernel_size,
            strides = 1,
            kernel_initializer = init_func,
            padding = "same")(x)
x = LeakyReLU(alpha = 0.2)(x)
x = BatchNormalization()(x)

# Layer 2 #

x = Conv2D(filters = feature_multiplier * 1, 
            kernel_size = surface_kernel_size,
            strides = 2,
            kernel_initializer = init_func, 
            padding = "same")(x)
x = LeakyReLU(alpha = 0.2)(x)
x = BatchNormalization()(x)
x = Dropout(dropout)(x) if dropout > 0.0 else x

# Output #

conv1_output = x

convolution_1 = Model(conv1_input, conv1_output)
layer_conv.append(convolution_1)

conv1_output_shape = (convolution_1.output_shape[1],
                      convolution_1.output_shape[2],
                      convolution_1.output_shape[3])

### Conv 2 ###

# Input #

conv2_input_shape = conv1_output_shape

conv2_input = Input(shape = conv2_input_shape)

x = conv2_input

# Layer 0 #

x = Conv2D(filters = feature_multiplier * 2, 
            kernel_size = kernel_size,
            strides = 1,
            kernel_initializer = init_func, 
            padding = "same")(x)
x = LeakyReLU(alpha = 0.2)(x)
x = BatchNormalization()(x)

# Layer 1 #

x = Conv2D(filters = feature_multiplier * 2, 
            kernel_size = kernel_size,
            strides = 2,
            kernel_initializer = init_func, 
            padding = "same")(x)
x = LeakyReLU(alpha = 0.2)(x)
x = BatchNormalization()(x)
x = Dropout(dropout)(x) if dropout > 0.0 else x

# Output #

conv2_output = x

convolution_2 = Model(conv2_input, conv2_output)
layer_conv.append(convolution_2)

conv2_output_shape = (convolution_2.output_shape[1],
                      convolution_2.output_shape[2],
                      convolution_2.output_shape[3])

### Conv 3 ###

# Input #

conv3_input_shape = conv2_output_shape

conv3_input = Input(shape = conv3_input_shape)

x = conv3_input

# Layer 0 #

x = Conv2D(filters = feature_multiplier * 4, 
            kernel_size = kernel_size,
            strides = 1,
            kernel_initializer = init_func, 
            padding = "same")(x)
x = LeakyReLU(alpha = 0.2)(x)
x = BatchNormalization()(x)

# Layer 1 #

x = Conv2D(filters = feature_multiplier * 4, 
            kernel_size = kernel_size,
            strides = 2,
            kernel_initializer = init_func, 
            padding = "same")(x)
x = LeakyReLU(alpha = 0.2)(x)
x = BatchNormalization()(x)
x = Dropout(dropout)(x) if dropout > 0.0 else x

# Output #

conv3_output = x

convolution_3 = Model(conv3_input, conv3_output)
layer_conv.append(convolution_3)

conv3_output_shape = (convolution_3.output_shape[1],
                      convolution_3.output_shape[2],
                      convolution_3.output_shape[3])

### Conv 4 ###

# Input #

conv4_input_shape = conv3_output_shape

conv4_input = Input(shape = conv4_input_shape)

x = conv4_input

# Layer 0 #

x = Conv2D(filters = feature_multiplier * 8, 
            kernel_size = kernel_size,
            strides = 1,
            kernel_initializer = init_func, 
            padding = "same")(x)
x = LeakyReLU(alpha = 0.2)(x)
x = BatchNormalization()(x)

# Layer 1 #

x = Conv2D(filters = feature_multiplier * 8, 
            kernel_size = kernel_size,
            strides = 2,
            kernel_initializer = init_func, 
            padding = "same")(x)
x = LeakyReLU(alpha = 0.2)(x)
x = BatchNormalization()(x)
x = Dropout(dropout)(x) if dropout > 0.0 else x

# Output #

conv4_output = x

convolution_4 = Model(conv4_input, conv4_output)
layer_conv.append(convolution_4)

conv4_output_shape = (convolution_4.output_shape[1],
                      convolution_4.output_shape[2],
                      convolution_4.output_shape[3])

### Conv 5 ###

# Input #

conv5_input_shape = conv4_output_shape

conv5_input = Input(shape = conv5_input_shape)

x = conv5_input

# Layer 0 #

x = Conv2D(filters = feature_multiplier * 16, 
            kernel_size = kernel_size,
            strides = 2,
            kernel_initializer = init_func, 
            padding = "same")(x)
x = LeakyReLU(alpha = 0.2)(x)
x = BatchNormalization()(x)
x = Dropout(dropout)(x) if dropout > 0.0 else x

# Output #

conv5_output = x

convolution_5 = Model(conv5_input, conv5_output)
layer_conv.append(convolution_5)

conv5_output_shape = (convolution_5.output_shape[1],
                      convolution_5.output_shape[2],
                      convolution_5.output_shape[3])

### Conv 6 ###

# Input #

conv6_input_shape = conv5_output_shape

conv6_input = Input(shape = conv6_input_shape)

x = conv6_input

# Layer 0 #

x = Conv2D(filters = feature_multiplier * 32, 
            kernel_size = kernel_size,
            strides = 2,
            kernel_initializer = init_func, 
            padding = "same")(x)
x = LeakyReLU(alpha = 0.2)(x)
x = BatchNormalization()(x)
x = Dropout(dropout)(x) if dropout > 0.0 else x

# Output #

conv6_output = x

convolution_6 = Model(conv6_input, conv6_output)
layer_conv.append(convolution_6)

conv6_output_shape = (convolution_6.output_shape[1],
                      convolution_6.output_shape[2],
                      convolution_6.output_shape[3])

# ------------------------------------------------------------------------------------------------------------------- #

layer_deconv = []

### Deconv 6 ###

# Input #

deconv6_input_shape = conv6_output_shape

deconv6_input = Input(shape = deconv6_input_shape)

x = deconv6_input

# Layer 0 #

x = Conv2DTranspose(filters = feature_multiplier * 16, 
                    kernel_size = kernel_size,
                    strides = 2,
                    kernel_initializer = init_func,
                    padding = "same")(x)
x = LeakyReLU(alpha = 0.2)(x)
x = BatchNormalization()(x)
x = Dropout(dropout)(x) if dropout > 0.0 else x

# Output #

deconv6_output = x

deconvolution_6 = Model(deconv6_input, deconv6_output)
layer_deconv.append(deconvolution_6)

### Deconv 5 ###

# Input # 

deconv5_input_shape = conv5_output_shape

deconv5_input = Input(shape = deconv5_input_shape)

x = deconv5_input

# Layer 0 #

x = Conv2DTranspose(filters = feature_multiplier * 8, 
                    kernel_size = kernel_size,
                    strides = 2,
                    kernel_initializer = init_func,
                    padding = "same")(x)
x = LeakyReLU(alpha = 0.2)(x)
x = BatchNormalization()(x)

# Layer 1 #

x = Conv2DTranspose(filters = feature_multiplier * 8, 
                    kernel_size = kernel_size,
                    strides = 1,
                    kernel_initializer = init_func,
                    padding = "same")(x)
x = LeakyReLU(alpha = 0.2)(x)
x = BatchNormalization()(x)
x = Dropout(dropout)(x) if dropout > 0.0 else x

# Output #

deconv5_output = x

deconvolution_5 = Model(deconv5_input, deconv5_output)
layer_deconv.append(deconvolution_5)

### Deconv 4 ###

# Input #

deconv4_input_shape = conv4_output_shape

deconv4_input = Input(shape = deconv4_input_shape)

x = deconv4_input

# Layer 0 #

x = Conv2DTranspose(filters = feature_multiplier * 4, 
                    kernel_size = kernel_size,
                    strides = 2,
                    kernel_initializer = init_func,
                    padding = "same")(x)
x = LeakyReLU(alpha = 0.2)(x)
x = BatchNormalization()(x)

# Layer 1 #

x = Conv2DTranspose(filters = feature_multiplier * 4, 
                    kernel_size = kernel_size,
                    strides = 1,
                    kernel_initializer = init_func,
                    padding = "same")(x)
x = LeakyReLU(alpha = 0.2)(x)
x = BatchNormalization()(x)
x = Dropout(dropout)(x) if dropout > 0.0  else x

# Output #

deconv4_output = x

deconvolution_4 = Model(deconv4_input, deconv4_output)
layer_deconv.append(deconvolution_4)

### Deconv 3 ###

# Input #

deconv3_input_shape = conv3_output_shape

deconv3_input = Input(shape = deconv3_input_shape)

x = deconv3_input

# Layer 0 #

x = Conv2DTranspose(filters = feature_multiplier * 2, 
                    kernel_size = kernel_size,
                    strides = 2,
                    kernel_initializer = init_func,
                    padding = "same")(x)
x = LeakyReLU(alpha = 0.2)(x)
x = BatchNormalization()(x)

# Layer 1 #

x = Conv2DTranspose(filters = feature_multiplier * 2, 
                    kernel_size = kernel_size,
                    strides = 1,
                    kernel_initializer = init_func,
                    padding = "same")(x)
x = LeakyReLU(alpha = 0.2)(x)
x = BatchNormalization()(x)
x = Dropout(dropout)(x) if dropout > 0.0 else x

# Output #

deconv3_output = x

deconvolution_3 = Model(deconv3_input, deconv3_output)
layer_deconv.append(deconvolution_3)

### Deconv 2 ###

# Input #

deconv2_input_shape = conv2_output_shape

deconv2_input = Input(shape = deconv2_input_shape)

x = deconv2_input

# Layer 0 #

x = Conv2DTranspose(filters = feature_multiplier * 1, 
                    kernel_size = kernel_size,
                    strides = 2,
                    kernel_initializer = init_func,
                    padding = "same")(x)
x = LeakyReLU(alpha = 0.2)(x)
x = BatchNormalization()(x) 

# Layer 1 #

x = Conv2DTranspose(filters = feature_multiplier * 1, 
                    kernel_size = kernel_size,
                    strides = 1,
                    kernel_initializer = init_func,
                    padding = "same")(x)
x = LeakyReLU(alpha = 0.2)(x)
x = BatchNormalization()(x)
x = Dropout(dropout)(x) if dropout > 0.0 else x

# Output # 

deconv2_output = x

deconvolution_2 = Model(deconv2_input, deconv2_output)
layer_deconv.append(deconvolution_2)

### Deconv 1 ###

# Input #

deconv1_input_shape = conv1_output_shape

deconv1_input = Input(shape = deconv1_input_shape)

x = deconv1_input

# Layer 0 #

x = Conv2DTranspose(input_shape[-1],
                    kernel_size = surface_kernel_size,
                    strides = 2,
                    padding = "same",
                    kernel_initializer = init_func)(x)
x = Activation("linear")(x)

# Output #

deconv1_output = x

deconvolution_1 = Model(deconv1_input, deconv1_output)
layer_deconv.append(deconvolution_1)

layer_deconv.reverse()

##### -------------------------------------------------- Autoencoder Entrenamiento ---------------------------------- #####

adam_learning_rate = 0.00015  # El learning rate de Adam (tamaño step)
adam_epsilon = 1e-8  # Previene problemas de división por 0.
adam_lr_decay = 1e-05  # Learning rate decay

optimizer = Adam(lr = adam_learning_rate, 
                 epsilon = adam_epsilon, 
                 decay = adam_lr_decay)

# ------------------------------------------------------------------------------------------------------------------- #

stages = []

stage_input = Input(shape = input_shape)

x = stage_input

x = layer_conv[0](x)

x = layer_deconv[0](x)

stage_output = x

stage_1 = Model(inputs = stage_input, outputs = stage_output)
stage_1.compile(optimizer = optimizer, loss = "mse", metrics = ["mae"])
stages.append(stage_1)

# ------------------------------------------------------------------------------------------------------------------- #

stage_input = Input(shape = input_shape)

x = stage_input

x = layer_conv[0](x)
x = layer_conv[1](x)

x = layer_deconv[1](x)
x = layer_deconv[0](x)

stage_output = x

stage_2 = Model(inputs = stage_input, outputs = stage_output)
stage_2.compile(optimizer = optimizer, loss = "mse", metrics = ["mae"])
stages.append(stage_2)

# ------------------------------------------------------------------------------------------------------------------- #

stage_input = Input(shape = input_shape)

x = stage_input

x = layer_conv[0](x)
x = layer_conv[1](x)
x = layer_conv[2](x)

x = layer_deconv[2](x)
x = layer_deconv[1](x)
x = layer_deconv[0](x)

stage_output = x

stage_3 = Model(inputs = stage_input, outputs = stage_output)
stage_3.compile(optimizer = optimizer, loss = "mse", metrics = ["mae"])
stages.append(stage_3)

# ------------------------------------------------------------------------------------------------------------------- #

stage_input = Input(shape = input_shape)

x = stage_input

x = layer_conv[0](x)
x = layer_conv[1](x)
x = layer_conv[2](x)
x = layer_conv[3](x)

x = layer_deconv[3](x)
x = layer_deconv[2](x)
x = layer_deconv[1](x)
x = layer_deconv[0](x)

stage_output = x

stage_4 = Model(inputs = stage_input, outputs = stage_output)
stage_4.compile(optimizer = optimizer, loss = "mse", metrics = ["mae"])
stages.append(stage_4)

# ------------------------------------------------------------------------------------------------------------------- #

stage_input = Input(shape = input_shape)

x = stage_input

x = layer_conv[0](x)
x = layer_conv[1](x)
x = layer_conv[2](x)
x = layer_conv[3](x)
x = layer_conv[4](x)

x = layer_deconv[4](x)
x = layer_deconv[3](x)
x = layer_deconv[2](x)
x = layer_deconv[1](x)
x = layer_deconv[0](x)

stage_output = x

stage_5 = Model(inputs = stage_input, outputs = stage_output)
stage_5.compile(optimizer = optimizer, loss = "mse", metrics = ["mae"])
stages.append(stage_5)

# ------------------------------------------------------------------------------------------------------------------- #

stage_input = Input(shape = input_shape)

x = stage_input

x = layer_conv[0](x)
x = layer_conv[1](x)
x = layer_conv[2](x)
x = layer_conv[3](x)
x = layer_conv[4](x)
x = layer_conv[5](x)

x = layer_deconv[5](x)
x = layer_deconv[4](x)
x = layer_deconv[3](x)
x = layer_deconv[2](x)
x = layer_deconv[1](x)
x = layer_deconv[0](x)

stage_output = x

stage_6 = Model(inputs = stage_input, outputs = stage_output)
stage_6.compile(optimizer = optimizer, loss = "mse", metrics = ["mae"])
stages.append(stage_6)

autoencoder = stage_6
autoencoder_clean_weights = autoencoder.get_weights()

# ------------------------------------------------------------------------------------------------------------------- #
   
epochs = pre_epochs
batch_size = pre_batch_size

if pre_training:
    
    for stage in stages:
        
        autoencoder_pre_train = stage.fit(train_data, train_data,
                                          epochs = epochs,
                                          batch_size = batch_size,
                                          validation_data = (vali_data, vali_data),
                                          shuffle = True)

    autoencoder.save("../modelos/autoencoder_true_pretraining.h5")
    autoencoder_pre_weights = autoencoder.get_weights()

# ------------------------------------------------------------------------------------------------------------------- #

epochs_list = ae_epochs_list
batch_list = ae_batch_list  # Distintas batch size para comparación

epochs = ae_epochs  # Número de vueltas completas al set de entrenamiento.
batch_size = ae_batch_size  # Número de ejemplos antes de calcular el error de la función de coste.

mc = ModelCheckpoint(filepath = "../modelos/autoencoder_true.h5", 
                     monitor = "val_loss", 
                     mode = "min", 
                     save_best_only = True,
                     verbose = 1)

if ae_batch_multiple:

    for batch_size in batch_list:
        
        autoencoder.set_weights(autoencoder_pre_weights)

        autoencoder_train = autoencoder.fit(train_data, train_data, 
                                            epochs = epochs_list,
                                            batch_size = batch_size,
                                            verbose = 1,
                                            validation_data = (vali_data, vali_data),
                                            shuffle = True,
                                            callbacks = [mc])
        
        training_plot(network_train = autoencoder_train, epochs = epochs_list, batch_size = batch_size, dropout = dropout, loss = "mse", metric = "mae", identification = "true_list")

else:
    
    autoencoder.set_weights(autoencoder_pre_weights)
    
    autoencoder_train = autoencoder.fit(train_data, train_data, 
                                        epochs = epochs,
                                        batch_size = batch_size,
                                        verbose = 1,
                                        validation_data = (vali_data, vali_data),
                                        shuffle = True,
                                        callbacks = [mc])
    
    training_plot(network_train = autoencoder_train, epochs = epochs, batch_size = batch_size, dropout = dropout, loss = "mse", metric = "mae", identification = "true_single")