#!/usr/bin/env python
#
# file: $ISIP_EXP/tuh_dpath/exp_0086/scripts/train.py
#
# revision history:
#  20191206 (TE): first version
#
# usage:
#  python train.py mdir data
#
# arguments:
#  mdir: the directory where the output model is stored
#  data: the input data list
#
# This script trains a simple MLP model
#------------------------------------------------------------------------------

# import tensorflow modules
#
import tensorflow as tf
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.utils import to_categorical
from tensorflow.keras.losses import categorical_crossentropy

# import the model and all of its variables/functions
#
from model import *

# import modules
#
import random
import sys
import os

#-----------------------------------------------------------------------------
#
# global variables are listed here
#
#-----------------------------------------------------------------------------

# general global values
#
NUM_ARGS = 2
NUM_EPOCHS = 100
BATCH_SIZE = 100
LEARNING_RATE = "lr"
BETA_1 = "beta_1"
BETA_2 = "beta_2"
EPS = "epsilon"
LOSS_FX = 'categorical_crossentropy'
WEIGHT_DECAY = "decay"
MDL_NAME="model.json"
WGT_NAME="model.h5"

# for reproducibility, we seed the RNGs
#
set_seed(SEED)

# ensure single threaded application to prevent
# non-deterministic operations
#
tf.config.threading.set_inter_op_parallelism_threads(1)
tf.config.threading.set_intra_op_parallelism_threads(1)

#------------------------------------------------------------------------------
#
# the main program starts here
#
#------------------------------------------------------------------------------

# function: main
#
# arguments: none
#
# return: none
#
# This method is the main function.
#
def main(argv):

    # ensure we have the correct amount of arguments
    #
    if(len(argv) != NUM_ARGS):
        print("usage: python nedc_train_mdl.py [MDL_PATH] [TRAIN_SET]")
        exit(-1)

    # define local variables
    #
    mdl_path = argv[0]
    fname = argv[1]
    num_feats = DEF_NUM_FEATS
    if("DL_NUM_FEATS" in os.environ):
        num_feats = int(os.environ["DL_NUM_FEATS"])

    # get the output directory name
    #
    odir = os.path.dirname(mdl_path)

    # if the odir doesn't exits, we make it
    #
    if not os.path.exists(odir):
        os.makedirs(odir)

    # get a file pointer
    #
    try:
        train_fp = open(fname, "r")
    except (IOError) as e:
        print("[%s]: %s" % (fname, e.strerror))
        exit(-1)

    # get array of the data
    # data: [[0, 1, ... 26], [27, 28, ...] ...]
    # labels: [0, 0, 1, ...]
    #
    train_data, train_labels = get_data(train_fp, num_feats)

    # shuffle the lists together as one
    #
    list_shuffle = list(zip(train_data, train_labels))
    random.shuffle(list_shuffle)
    train_data, train_labels = zip(*list_shuffle)

    # close the file
    #
    train_fp.close()

    # get the number of epochs to train on
    #
    epochs = NUM_EPOCHS

    # get the batch size
    #
    batch_size = BATCH_SIZE

    # set the adam optimizer parameters
    #
    opt_params = { LEARNING_RATE: 0.005,
                   BETA_1: 0.9,
                   BETA_2: 0.999,
                   EPS: 1e-08,
                   WEIGHT_DECAY: .00001 }

    # create an optimizer
    #
    adam_opt = Adam(**opt_params)

    # instantiate a model
    #
    model = Model(num_feats, NUM_NODES, NUM_CLASSES, SEED)

    # compile the model
    #
    model.compile(optimizer=adam_opt, loss=LOSS_FX, metrics=['accuracy'])

    # fit the training data to the model
    #
    model.fit(x=np.array(train_data), y=to_categorical(np.array(train_labels)),
              batch_size=batch_size, epochs=epochs, verbose=2, shuffle=False)

    # convert the model to a JSON
    #
    model_json = model.to_json()

    # write the JSON to the mdl dir
    #
    with open(os.path.join(mdl_path, MDL_NAME), "w") as json_file:
        json_file.write(model_json)

    # save the weights in the same directory
    #
    model.save_weights(os.path.join(mdl_path, WGT_NAME))

    # exit gracefully
    #
    return True
#
# end of function


# begin gracefully
#
if __name__ == '__main__':
    main(sys.argv[1:])
#
# end of file