#!/usr/bin/env python
'''Train a simple deep CNN on the CIFAR10 small images dataset.

Distributed training.

run:
    salloc -N2  -p hsw_p40 --comment=docker  --qos=short --time=04:00:00

    srun -l --ntasks-per-node=2 python \
    examples/cifar/cifar10_cnn_distrib_slurm.py --epochs=5 --network=ib.cluster

    srun -l --ntasks-per-node=2 \
      python examples/cifar/cifar10_cnn_distrib_slurm.py --epochs=5 --rdma
    # rdma default is verbs

    # On psgcluster --network=ib.cluster is required for --rdma=gdr option
    # couldn't get --rdma=gdr option to work.

'''
from __future__ import print_function

import sys

from argparse import SUPPRESS

# from time import sleep

import tensorflow as tf

from keras.utils import to_categorical
from keras.datasets import cifar10
from keras.preprocessing.image import ImageDataGenerator
from keras import backend as KB

from keras.callbacks import ModelCheckpoint

from keras.optimizers import RMSprop

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

from keras_exp.distrib.cluster_mgrs.tfcmgr import (
    TFClusterManagerFacade, JobType)  # , DevType
from keras_exp.distrib.cluster_mgrs.tfclusterdefs import ProtocolType


from keras_exp.distrib.cluster_parsers.slurm import SlurmClusterParser
from cifar_common import make_model

from parser_common import (parser_def_mgpu, remove_options)

# from functools import partial

_DEVPROF = False


def parser_(desc):
    '''CLI Parser for options to run Tensorflow over SLURM cluster.'''
    parser = parser_def_mgpu(desc)
    remove_options(parser, ['--mgpu', '--nccl', '--enqueue'])

    checkptfile = 'cifar10_cnn_distrib.weights.best.hdf5'
    parser.add_argument(
        '--checkpt', action='store', nargs='?',
        const=checkptfile, default=SUPPRESS,
        help='S|Save (overwrites) and load the model weights if available.'
        '\nOptionally specify a file/filepath if the default name is '
        'undesired.\n(default: {})'.format(checkptfile))

    parser.add_argument('--aug', action='store_true', default=False,
                        help='S|Perform data augmentation on cifar10 set.\n')

    parser.add_argument('--logdevp', action='store_true', default=False,
                        help='S|Log device placement in Tensorflow.\n')

    args = parser.parse_args()

    return args


def main(argv=None):
    '''Train a simple deep CNN on the CIFAR10 small images dataset on a SLURM
    cluster.'''
    main.__doc__ = __doc__

    argv = sys.argv if argv is None else sys.argv.extend(argv)
    desc = main.__doc__  # .format(os.path.basename(__file__))
    # CLI parser
    args = parser_(desc)
    # print('RDMA: {}'.format(args.rdma))
    # rdma = getattr(args, 'rdma', None)
    rdma = args.rdma
    network = args.network
    # print('NETWORK: {}'.format(network))

    checkpt = getattr(args, 'checkpt', None)
    checkpt_flag = False if checkpt is None else True
    filepath = checkpt
    # print('CHECKPT:', checkpt)

    batch_size = 32
    num_classes = 10
    epochs = args.epochs
    data_augmentation = args.aug

    logdevp = args.logdevp

    # ---------------------------------------------- Distributed setup on SLURM
    # Specifying network necessary for protocol='grpc+gdr'. GDR doesn't find
    # IB addresses automatically like 'grpc+verbs'.
    # The 'ib.cluster' is specific to NVIDIA psgcluster.
    # network = 'ib.cluster' if rdma == 'gdr' else None
    # network = 'ib.cluster'
    # On fast network even without RDMA speed up significant. RDMA still helps.
    scpar = SlurmClusterParser(network=network)
    cmgr_facade = TFClusterManagerFacade(scpar)

    logdevp_flag = True if _DEVPROF or logdevp else False
    gpu_options = tf.GPUOptions(allow_growth=True)
    config = tf.ConfigProto(log_device_placement=logdevp_flag,  # True,
                            allow_soft_placement=True,
                            gpu_options=gpu_options)

    print('\n\tCLUSTER_SPEC_DICT: {}\n'.format(cmgr_facade.clusterspec_dict))

    # TF 1.2.x RDMA: specify protocol='grpc+verbs' in server below.
    protocol = ProtocolType.get_server_protocol_str(rdma)
    # print('PROTOCOL: {}'.format(protocol))
    server = cmgr_facade.get_server(
        config,
        protocol=protocol)
    tfsess = cmgr_facade.get_session(server)
    KB.set_session(tfsess)

    #: :type cluster_spec: tf.train.ClusterSpec
    # cluster_spec = cmgr_facade.get_cluster_spec()
    job_type = cmgr_facade.myjobtype
    # task_id = cmgr_facade.mytask_id

    is_chief = cmgr_facade.is_chief

    if job_type == JobType.ps:
        # JOIN PARAMETER SERVERS
        # server.join()
        cmgr_facade.join(server)

    # Once the server is started everything but the chief worker can join
    # the server and wait to process/service graph computations. Chief pushes
    # the compute graph. COMPARE TO: cifar10_cnn_distrib_v2_slurm
    if not is_chief:
        # JOIN WORKERS EXCEPT FOR CHIEF
        cmgr_facade.join(server)

    # sleep(2)  # Have the chief wait just in case. Occasionally get errors.

    # The ngpus per host needs to be done with MPI or somehow sync'd. Currently
    # assuming all hosts have the same number of GPUs.
    gdev_list = get_available_gpus()
    ngpus = len(gdev_list)

    # List of all devices. The devices might be associated to the same worker.
    wgdev_list = cmgr_facade.get_allworkers_devlist(ngpus)
    print('\n\tWGDEV_LIST: {}\n'
          .format([dev.to_string() for dev in wgdev_list]))  # DEBUG
    # If 2 workers ea. w/ 4 devices then nworker_devices_total == 2 * 4 = 8
    # If 4 workers ea. w/ 1 devices then nworker_devices_total == 4 * 1 = 4
    nworker_devices_total = len(wgdev_list)
    batch_size = batch_size * nworker_devices_total

    psdev_list = cmgr_facade.get_allps_devlist()
    print('\n\tPSDEV_LIST: {}\n'
          .format([dev.to_string() for dev in psdev_list]))  # DEBUG

    # ------------------------------------ Data loading and basic preprocessing
    # The data, shuffled and split between train and test sets:
    (x_train, y_train), (x_test, y_test) = cifar10.load_data()
    # Convert class vectors to binary class matrices.
    y_train = to_categorical(y_train, num_classes)
    y_test = to_categorical(y_test, num_classes)

    x_train = x_train.astype('float32')
    x_test = x_test.astype('float32')
    x_train /= 255
    x_test /= 255

    nsamples = x_train.shape[0]
    steps_per_epoch = nsamples // batch_size

    print(x_train.shape[0], 'train samples')
    print(x_test.shape[0], 'test samples')

    # --------------------------------------------- Setup model and parallelize
    def _load_fn(unused_op):
        return 1

    cspec = cmgr_facade.get_cluster_spec()
    num_ps = cmgr_facade.num_ps
    ps_strategy = \
        tf.contrib.training.GreedyLoadBalancingStrategy(num_ps, _load_fn)

    # ps_device = tf.DeviceSpec(job=JobType.ps, device_type=DevType.cpu,
    #                           device_index=0).to_string()

    rdsetter = tf.train.replica_device_setter(
        cluster=cspec,
        ps_strategy=ps_strategy,
        # ps_device=ps_device,  # '/job:ps/cpu:0'  # seems to work
        # ps_device='/gpu:0'  # for gdr maybe
    )
    with tf.device(rdsetter):
        model_init = make_model(
            x_train.shape[1:], num_classes,
            filepath if checkpt_flag else None
        )

    callbacks = None
    if checkpt_flag:
        checkpoint = ModelCheckpoint(filepath, monitor='val_acc', verbose=1,
                                     save_best_only=True, mode='max')
        callbacks = [checkpoint]

    # Data-Parallelize the model via function or class.
    model = make_parallel(model_init, wgdev_list)  # , ps_device='/gpu:0'
    print_mgpu_modelsummary(model)

    # ------------------------------------------------------------ Run training
    lr = 0.0001 * nworker_devices_total
    opt = RMSprop(lr=lr, decay=1e-6)

    # Let's train the model using RMSprop
    model.compile(loss='categorical_crossentropy',
                  optimizer=opt,
                  metrics=['accuracy'])

    if not data_augmentation:
        print('Not using data augmentation.')
        model.fit(x_train, y_train,
                  batch_size=batch_size,
                  epochs=epochs,
                  validation_data=(x_test, y_test),
                  shuffle=True,
                  callbacks=callbacks)

    else:
        print('Using real-time data augmentation.')
        # This will do preprocessing and realtime data augmentation:
        datagen = ImageDataGenerator(
            featurewise_center=False,  # set input mean to 0 over the dataset
            samplewise_center=False,  # set each sample mean to 0
            # divide inputs by std of the dataset
            featurewise_std_normalization=False,
            samplewise_std_normalization=False,  # divide each input by its std
            zca_whitening=False,  # apply ZCA whitening
            # randomly rotate images in the range (degrees, 0 to 180)
            rotation_range=0,
            # randomly shift images horizontally (fraction of total width)
            width_shift_range=0.1,
            # randomly shift images vertically (fraction of total height)
            height_shift_range=0.1,
            horizontal_flip=True,  # randomly flip images
            vertical_flip=False)  # randomly flip images

        # Compute quantities required for feature-wise normalization
        # (std, mean, and principal components if ZCA whitening is applied).
        datagen.fit(x_train)

        # Fit the model on the batches generated by datagen.flow().
        model.fit_generator(datagen.flow(x_train, y_train,
                                         batch_size=batch_size),
                            steps_per_epoch=steps_per_epoch,
                            epochs=epochs,
                            validation_data=(x_test, y_test),
                            callbacks=callbacks)

    # Run Validation
    if is_chief:
        model_init.compile(loss='categorical_crossentropy',
                           optimizer=opt,
                           metrics=['accuracy'])
        metrics = model_init.evaluate(
            x=x_test, y=y_test,
            batch_size=batch_size)
        print('\nCIFAR VALIDATION LOSS, ACC: {}, {}'.format(*metrics))

    # ------------------------------------------------------------- STOP SERVER
    cmgr_facade.stop_chief(server)


if __name__ == '__main__':
    main()
