from __future__ import print_function
from keras.datasets import mnist
from keras.datasets import cifar10
from keras.utils.np_utils import to_categorical
import numpy as np
from keras import backend as K
from evolution import Evolution
from genome_handler import GenomeHandler
import tensorflow as tf
#import mlflow.keras
#import mlflow
#import mlflow.tensorflow
#mlflow.tensorflow.autolog() 
#mlflow.keras.autolog()

print("Num GPUs Available: ", len(tf.config.list_physical_devices('GPU')))

K.set_image_data_format("channels_last")

#(x_train, y_train), (x_test, y_test) = mnist.load_data()
(x_train, y_train), (x_test, y_test) = cifar10.load_data()

x_train = x_train.reshape(x_train.shape[0], x_train.shape[1], x_train.shape[2],x_train.shape[3]).astype('float32') / 255
x_test = x_test.reshape(x_test.shape[0], x_test.shape[1], x_test.shape[2], x_test.shape[3]).astype('float32') / 255
# nCLasses
y_train = to_categorical(y_train)
y_test = to_categorical(y_test)
#y_train.shape
dataset = ((x_train, y_train), (x_test, y_test))

genome_handler = GenomeHandler(max_conv_layers=4, 
                               max_dense_layers=2, # includes final dense layer
                               max_filters=512,
                               max_dense_nodes=1024,
                               input_shape=x_train.shape[1:],
                               n_classes=10)

evo = Evolution(genome_handler, data_path="log/evo_cifar10_gen40_pop10_e20.csv")
model = evo.run(dataset=dataset,
                  num_generations=40,
                  pop_size=10,
                  epochs=20,metric='acc')
                  #epochs=10,metric='loss')
print(model.summary())
[((494, 534), 'keras.backend.set_image_data_format', 'K.set_image_data_format', (['"""channels_last"""'], {}), "('channels_last')\n", (517, 534), True, 'from keras import backend as K\n'), ((633, 652), 'keras.datasets.cifar10.load_data', 'cifar10.load_data', ([], {}), '()\n', (650, 652), False, 'from keras.datasets import cifar10\n'), ((912, 935), 'keras.utils.np_utils.to_categorical', 'to_categorical', (['y_train'], {}), '(y_train)\n', (926, 935), False, 'from keras.utils.np_utils import to_categorical\n'), ((945, 967), 'keras.utils.np_utils.to_categorical', 'to_categorical', (['y_test'], {}), '(y_test)\n', (959, 967), False, 'from keras.utils.np_utils import to_categorical\n'), ((1050, 1190), 'genome_handler.GenomeHandler', 'GenomeHandler', ([], {'max_conv_layers': '(4)', 'max_dense_layers': '(2)', 'max_filters': '(512)', 'max_dense_nodes': '(1024)', 'input_shape': 'x_train.shape[1:]', 'n_classes': '(10)'}), '(max_conv_layers=4, max_dense_layers=2, max_filters=512,\n    max_dense_nodes=1024, input_shape=x_train.shape[1:], n_classes=10)\n', (1063, 1190), False, 'from genome_handler import GenomeHandler\n'), ((1379, 1453), 'evolution.Evolution', 'Evolution', (['genome_handler'], {'data_path': '"""log/evo_cifar10_gen40_pop10_e20.csv"""'}), "(genome_handler, data_path='log/evo_cifar10_gen40_pop10_e20.csv')\n", (1388, 1453), False, 'from evolution import Evolution\n'), ((452, 490), 'tensorflow.config.list_physical_devices', 'tf.config.list_physical_devices', (['"""GPU"""'], {}), "('GPU')\n", (483, 490), True, 'import tensorflow as tf\n')]