import sys
import os
import tensorflow as tf
import math

# VGG16 es un modelo preentrenado que contiene multitud de imagenes.
# Este modelo preentrado es uno de los muchos que tiene Keras
vgg = tf.keras.applications.vgg16.VGG16()

# Resumen de las capas que tiene la red neuronal VGG16
vgg.summary()

# Creación de la red convolucional
cnn = tf.keras.Sequential()

# Se añaden las capas que se ha visto antes a nuestra propia red
for capa in vgg.layers:
    cnn.add(capa)

# Resumen con las capas de nuestra propia red. Debe coincidir con el resumen
# realizado en la red VGG16
cnn.summary()

# Con esto se elimina la capa de predicción. El pop elimina la ultima capa
# del modelo. La ultima capa era la de "predictions" que tiene 1000 parámetros, tipos de imagenes,
# que pued clasificar. Como queremos usarlo para el ejemplo de gatos, gorilas y perros, se quita dicha
# capa para añadir una propia.
cnn.pop()

# Las capas que provienen del VGG16 no queremos entrenarlas porque ya lo han sido
# en el pasado. Por lo tanto, indicamos que no es necesario entrenarlas.
for layer in cnn.layers:
    layer.trainable = False

# Añadimos una capa de predicción de softmax que tendrá 3 neuronas. Ya que se usará el ejemplo
# de gatos, gorilas y perros.
cnn.add(tf.keras.layers.Dense(3,activation='softmax'))

# Esta función encapsula lo que hemos hecho en el código anterior paso a paso.
def modelo():
    vgg=tf.keras.applications.vgg16.VGG16()
    cnn=tf.keras.Sequential()
    for capa in vgg.layers:
        cnn.add(capa)
    cnn.pop()
    for layer in cnn.layers:
        layer.trainable = False
    cnn.add(tf.keras.layers.Dense(3,activation='softmax'))

    return cnn

# Se limpia todas las variables, estado, etc. de sesión que hemos abierto de keras
tf.keras.backend.clear_session()

# Trozo de código del clasificador de imagenes pero aquí no se indicarán las capas convolucionales
data_entrenamiento = "K:/github/python_test/tensorflow/clasificador_imagenes/data/entrenamiento/"
data_validacion = "K:/github/python_test/tensorflow/clasificador_imagenes/data/validacion/"

# Parametros de la red neuronal
# Número de veces que se iteran los datos en el entrenamiento
epocas = 20
# Tamaño de las imagenes. Este es el tamaño que espera la red VGG16, en la documentación
# oficial lo indica.
longitud, altura = 224, 224
# Numero de imagenes que se procesan en cada paso
batch_size = 32
# Al final de cada época se harán 300 pasos con los datos de validación
# para ir viendo como de bien va aprendiendo
validation_steps = 32
# Número de convoluciones o de capas(profundidad) que tendra la iamagen
# Primera convolucion serán 32
filtrosConv1 = 32
# Segunda convolucion serán 64
filtrosConv2 = 64
# Anchura y altura que va a procesar en cada convolucion.
tamano_filtro1 = (3, 3)
tamano_filtro2 = (2, 2)
# Tamaño que se va usar en el maxpooling
tamano_pool = (2, 2)
# Número de clases que hay en el set de datos: gatos, perros y gorilas
clases = 3
# Es el learning rate. Es decir, cuanto de grandes van a ser los ajustes en la red
# para ajustarse para buscar una solución óptima.
lr = 0.0004

# Como se van a transformar las imagenes para poderlas pasar al procesao de entrenamiento
entrenamiento_datagen = tf.keras.preprocessing.image.ImageDataGenerator(
    rescale=1. / 255,
    shear_range=0.2,
    zoom_range=0.2,
    horizontal_flip=True)

# Lo mismo pero para los datos de test
test_datagen = tf.keras.preprocessing.image.ImageDataGenerator(
    rescale=1. / 255)    

# Proceso de lectura y transformación de los datos para los datos de entrenamiento
entrenamiento_generador = entrenamiento_datagen.flow_from_directory(
    data_entrenamiento,
    target_size=(altura, longitud),
    batch_size=batch_size,
    class_mode='categorical')

# Lo mismo para los datos de test
validacion_generador = test_datagen.flow_from_directory(
    data_validacion,
    target_size=(altura, longitud),
    batch_size=batch_size,
    class_mode='categorical')    
# Se llama a la función que devolverá el modelo adaptado basandonas en la red BGG16
cnn = modelo()

# Se indica como aprenderá el algoritmo
cnn.compile(loss='categorical_crossentropy',
            optimizer=tf.keras.optimizers.Adam(lr=lr),
            metrics=['accuracy'])
# Calculo de pasos por epoca
pasos = math.ceil( ( len(entrenamiento_generador.filenames) / batch_size ) )    

cnn.fit(
    entrenamiento_generador,
    steps_per_epoch=pasos,
    epochs=epocas,
    validation_data=validacion_generador,
    validation_steps=validation_steps )

# Directorio donde se guardará el modelo
target_dir = 'K:/github/python_test/tensorflow/clasificador_imagenes/modelo'
if not os.path.exists(target_dir):
 os.mkdir(target_dir)

# Grabación del modelo
cnn.save('K:/github/python_test/tensorflow/clasificador_imagenes/modelo/modelo.h5')

# Grabación de los pesos de cada una de las capas
cnn.save_weights('K:/github/python_test/tensorflow/clasificador_imagenes/modelo/pesos.h5')