'''
MNIST dataset with TensorFlow TFRecords. Refer to:
    https://github.com/fchollet/keras/blob/master/examples/mnist_tfrecord.py

This is an implementation for multi-GPU systems.

Call this example:
    # best achieved performance with 2 GPUs
    time CUDA_VISIBLE_DEVICES=0,1 python mnist_tfrecord_mgpu.py
    # or no timeing info:
    CUDA_VISIBLE_DEVICES=0,1 python mnist_tfrecord_mgpu.py

    # vary CUDA_VISIBLE_DEVICES for more than 2 GPUs. Performance starts to
    # degrade.
    time CUDA_VISIBLE_DEVICES=0,1,3,4 python mnist_tfrecord_mgpu.py

    # compare 2 GPUs to one GPU
    time CUDA_VISIBLE_DEVICES=0 python mnist_tfrecord_mgpu.py

    time CUDA_VISIBLE_DEVICES=0,1 python mnist_tfrecord_mgpu.py
    # The overall walltime for this whole example might be faster with 1 GPU
    # due to startup time being longer for multigpu case, but the training
    # portion is faster with 2 GPUs. Running on P100s 50 epochs:
    #     nGPUs | training (sec). | walltime (sec.)
    #         1 | ~ 2             | ~ 16
    #         2 | ~ 1             | ~ 12
    #
    # Degrades with > 2 GPUs because the mnist model is not significant enough
    # to stress the computing of the GPUs so the startup/comm. Overhead is
    # greater than speedup achieved due to data-parallelism.

Using TFRecord queues with Keras can be a significant performance booster. When
the model and batch sizes are signifcantly large, using multigpu with TFRecord
queue can give additional performance boost to Keras.

'''
from __future__ import print_function
import time

import numpy as np

import tensorflow as tf
from tensorflow.contrib.learn.python.learn.datasets import mnist

from keras import backend as KB
from keras.models import Model
from keras.layers import (
    Dense, Dropout, Flatten, Input, Conv2D, MaxPooling2D)
from keras.optimizers import RMSprop

from keras.models import Sequential
import keras.layers as KL

from keras.callbacks import ModelCheckpoint

from keras.utils import to_categorical

from keras_exp.multigpu import (
    get_available_gpus, make_parallel, print_mgpu_modelsummary)


if KB.backend() != 'tensorflow':
    raise RuntimeError('This example can only run with the '
                       'TensorFlow backend, '
                       'because it requires TFRecords, which '
                       'are not supported on other platforms.')


def cnn_layers_list(nclasses):
    '''List of instantitated cnn layers to use for the mnist data.'''
    ll = []
    ll.append(Conv2D(32, (3, 3), activation='relu', padding='valid'))
    ll.append(Conv2D(64, (3, 3), activation='relu'))
    ll.append(MaxPooling2D(pool_size=(2, 2)))
    ll.append(Dropout(0.25))
    ll.append(Flatten())
    ll.append(Dense(128, activation='relu'))
    ll.append(Dropout(0.5))
    ll.append(Dense(nclasses, activation='softmax', name='x_train_out'))

    return ll


def cnn_layers(x_train_input, nclasses):
    '''Create a model via functional API using cnn layers.'''
    ll = cnn_layers_list(nclasses)
    xi = x_train_input
    for il in ll:
        xi = il(xi)

    return xi


def make_model(x_train_input, nclasses):
    '''Non-functional model definition.'''
    model = Sequential()
    model.add(KL.InputLayer(input_tensor=x_train_input))
    ll = cnn_layers_list(nclasses)
    for il in ll:
        model.add(il)

    return model


def main():
    # user options
    batch_size = 128
    val_in_train = False  # not sure how the validation part works during fit.
    use_model_checkpt = False

    # demo processing
    sess = tf.Session()
    KB.set_session(sess)

    gdev_list = get_available_gpus()
    ngpus = len(gdev_list)
    batch_size = batch_size * ngpus

    data = mnist.load_mnist()
    X_train = data.train.images
    # X_test = data.test.images
    train_samples = X_train.shape[0]  # 60000
    # test_samples = X_test.shape[0]  # 10000
    height_nrows = 28
    width_ncols = 28
    batch_shape = [batch_size, height_nrows, width_ncols, 1]
    epochs = 5
    steps_per_epoch = train_samples // batch_size
    # validations_steps = test_samples / batch_size
    nclasses = 10

    # The capacity variable controls the maximum queue size
    # allowed when prefetching data for training.
    capacity = 10000

    # min_after_dequeue is the minimum number elements in the queue
    # after a dequeue, which ensures sufficient mixing of elements.
    min_after_dequeue = 3000

    # If `enqueue_many` is `False`, `tensors` is assumed to represent a
    # single example.  An input tensor with shape `[x, y, z]` will be output
    # as a tensor with shape `[batch_size, x, y, z]`.
    #
    # If `enqueue_many` is `True`, `tensors` is assumed to represent a
    # batch of examples, where the first dimension is indexed by example,
    # and all members of `tensors` should have the same size in the
    # first dimension.  If an input tensor has shape `[*, x, y, z]`, the
    # output will have shape `[batch_size, x, y, z]`.
    enqueue_many = True

    x_train_batch, y_train_batch = tf.train.shuffle_batch(
        tensors=[data.train.images, data.train.labels.astype(np.int32)],
        batch_size=batch_size,
        capacity=capacity,
        min_after_dequeue=min_after_dequeue,
        enqueue_many=enqueue_many,
        num_threads=8)

    x_train_batch = tf.cast(x_train_batch, tf.float32)
    x_train_batch = tf.reshape(x_train_batch, shape=batch_shape)

    y_train_batch = tf.cast(y_train_batch, tf.int32)
    y_train_batch = tf.one_hot(y_train_batch, nclasses)

    x_train_input = Input(tensor=x_train_batch)

    # x_test_batch, y_test_batch = tf.train.batch(
    #     tensors=[data.test.images, data.test.labels.astype(np.int32)],
    #     batch_size=batch_size,
    #     capacity=capacity,
    #     enqueue_many=enqueue_many,
    #     num_threads=8)

    # I like the non-functional definition of model more.
    # model_init = make_model(x_train_input, nclasses)
    # x_train_out = model_init.output
    # train_model = Model(inputs=[x_train_input], outputs=[x_train_out])

    x_train_out = cnn_layers(x_train_input, nclasses)
    train_model = Model(inputs=[x_train_input], outputs=[x_train_out])
    if ngpus > 1:
        train_model = make_parallel(train_model, gdev_list)

    lr = 2e-3 * ngpus
    train_model.compile(optimizer=RMSprop(lr=lr, decay=1e-5),
                        loss='categorical_crossentropy',
                        metrics=['accuracy'],
                        target_tensors=[y_train_batch])

    if ngpus > 1:
        print_mgpu_modelsummary(train_model)
    else:
        train_model.summary()

    # Callbacks
    if use_model_checkpt:
        mon = 'val_acc' if val_in_train else 'acc'
        checkpoint = ModelCheckpoint(
            'saved_wt.h5', monitor=mon, verbose=0,
            save_best_only=True,
            save_weights_only=True)
        checkpoint = [checkpoint]
    else:
        checkpoint = []

    callbacks = checkpoint
    # Training slower with callback. Multigpu slower with callback during
    # training than 1 GPU. Again, mnist is too trivial of a model and dataset
    # to benchmark or stress GPU compute capabilities. I set up this example
    # to illustrate potential for speedup of multigpu case trying to use mnist
    # as a stressor.
    # It's like comparing a 5 ft race between a person and a truck. A truck is
    # obviously faster than a person but in a 5 ft race the person will likely
    # win due to slower startup for the truck.
    # I will re-implement this with Cifar that should be a better benchmark.

    # Start the queue runners.
    tf.train.start_queue_runners(sess=sess)

    # Fit the model using data from the TFRecord data tensors.
    coord = tf.train.Coordinator()
    threads = tf.train.start_queue_runners(sess, coord)

    start_time = time.time()
    train_model.fit(
        # validation_data=(x_test_batch, y_test_batch)
        # if val_in_train else None, # validation data is not used???
        # validations_steps if val_in_train else None,
        # validation_steps=val_in_train,
        steps_per_epoch=steps_per_epoch,
        epochs=epochs,
        callbacks=callbacks)
    elapsed_time = time.time() - start_time
    print('[{}] finished in {} s'.format('TRAINING', round(elapsed_time, 3)))

    if not checkpoint:  # empty list
        train_model.save_weights('./saved_wt.h5')

    # Clean up the TF session.
    coord.request_stop()
    coord.join(threads)

    KB.clear_session()

    # Second Session. Demonstrate that the model works and is independent of
    # the TFRecord pipeline, and to test loading trained model without tensors.
    x_test = np.reshape(data.validation.images,
                        (data.validation.images.shape[0], 28, 28, 1))
    y_test = data.validation.labels
    x_test_inp = KL.Input(shape=(x_test.shape[1:]))
    test_out = cnn_layers(x_test_inp, nclasses)
    test_model = Model(inputs=x_test_inp, outputs=test_out)

    test_model.load_weights('saved_wt.h5')
    test_model.compile(optimizer='rmsprop',
                       loss='categorical_crossentropy',
                       metrics=['accuracy'])
    test_model.summary()

    loss, acc = test_model.evaluate(x_test, to_categorical(y_test))
    print('\nTest loss: {0}'.format(loss))
    print('\nTest accuracy: {0}'.format(acc))


if __name__ == '__main__':
    main()
