#!/usr/bin/env python

import argparse
import logging
import math
import os
import pprint
import shutil
import sys
import traceback
import glob
import time
import re
import models
import ladder_block as ladder
import ze_utils as utils
import subprocess

MAX_TRY_COUNT = 16

logger = logging.getLogger('train_dnn')
logger.setLevel(logging.INFO)
handler = logging.StreamHandler(sys.stdout)
handler.setLevel(logging.INFO)
formatter = logging.Formatter("%(asctime)s [%(pathname)s:%(lineno)s - "
                              "%(funcName)s - %(levelname)s ] %(message)s")
handler.setFormatter(formatter)
logger.addHandler(handler)


def get_args():
    """ Get args from stdin.
    """

    parser = argparse.ArgumentParser(
        description="""Trains a feed forward raw DNN (without transition model)
        using frame-level objectives like cross-entropy and mean-squared-error.
        DNNs include simple DNNs, TDNNs and CNNs.""",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
        conflict_handler='resolve')

    # Parameters for the optimization
    parser.add_argument("--use-gpu", type=str, dest='use_gpu',
                        choices=["yes", "no", "wait"],
                        help="Use GPU for training.",
                        default="yes")

    parser.add_argument("--momentum", type=float,
                        dest='momentum', default=0.0,
                        help="""Momentum used in update computation.
                        Note: we implemented it in such a way that it
                        doesn't increase the effective learning rate.""")

    parser.add_argument("--targets-scp", type=str, required=False,
                        help="Targets for training neural network.")

    parser.add_argument("--tf-model-class", type=str, dest='tf_model_class', required=True,
                        help="Shows the tensorflow class name which should was defined in models.py file.")

    parser.add_argument("--dir", type=str, required=True,
                        help="Directory to store the models and all other files.")

    parser.add_argument("--egs-dir", type=str, dest='egs_dir', required=True,
                        help="Directory of training egs.")

    parser.add_argument("--num-epochs", type=float, dest='num_epochs', default=6.0,
                        help="Number of epochs to train the model.")

    parser.add_argument("--num-targets", type=int, dest='num_targets', required=True,
                        help="Shows the number of output of the neural network.")

    parser.add_argument("--initial-effective-lrate", type=float,
                        dest='initial_effective_lrate', default=0.0003,
                        help="Learning rate used during the initial iteration.")

    parser.add_argument("--final-effective-lrate", type=float,
                        dest='final_effective_lrate', default=0.00003,
                        help="Learning rate used during the final iteration.")

    parser.add_argument("--num-jobs-initial", type=int, dest='num_jobs_initial', default=1,
                        help="Number of neural net jobs to run in "
                             "parallel at the start of training")

    parser.add_argument("--num-jobs-final", type=int, dest='num_jobs_final', default=8,
                        help="Number of neural net jobs to run in "
                             "parallel at the end of training")

    parser.add_argument("--minibatch-size",
                        type=int, dest='minibatch_size', required=True,
                        help="Size of the minibatch used in SGD training.")

    parser.add_argument("--do-final-combination", dest='do_final_combination', type=str,
                        action=utils.StrToBoolAction,
                        choices=["true", "false"], default=False,
                        help="""Set this to false to disable the final
                             'combine' stage (in this case we just use the
                             last-numbered model as the final.mdl).""")

    parser.add_argument("--random-seed", type=int, dest='random_seed', default=0,
                        help="""Sets the random seed for egs shuffling and tensorflow random seed.
                             Warning: This random seed does not control all aspects of this 
                             experiment.  There might be other random seeds used in other stages of the
                             experiment like data preparation (e.g. volume perturbation).""")

    parser.add_argument("--dropout-schedule", type=str,
                        action=utils.NullStrToNoneAction,
                        dest='dropout_schedule', default=None,
                        help="""Use this to specify the dropout
                             schedule.  You specify a piecewise linear
                             function on the domain [0,1], where 0 is the
                             start and 1 is the end of training; the
                             function-argument (x) rises linearly with the
                             amount of data you have seen, not iteration
                             number (this improves invariance to
                             num-jobs-{initial-final}).  E.g. '0,0.2,0'
                             means 0 at the start; 0.2 after seeing half
                             the data; and 0 at the end.  You may specify
                             the x-value of selected points, e.g.
                             '0,0.2@0.25,0' means that the 0.2
                             dropout-proportion is reached a quarter of the
                             way through the data.   The start/end x-values
                             are at x=0/x=1, and other unspecified x-values
                             are interpolated between known x-values.  You
                             may specify different rules for different
                             component-name patterns using 'pattern1=func1
                             pattern2=func2', e.g. 'relu*=0,0.1,0
                             lstm*=0,0.2,0'.  More general should precede
                             less general patterns, as they are applied
                             sequentially.""")
    parser.add_argument("--max-objective-evaluations",
                        type=int, dest='max_objective_evaluations', default=30,
                        help="""The maximum number of objective evaluations in order to figure out the
                         best number of models to combine. It helps to speedup if the number of models provided to the
                         model combination binary is quite large (e.g. several hundred).""")

    parser.add_argument("--preserve-model-interval", dest="preserve_model_interval",
                        type=int, default=10,
                        help="""Determines iterations for which models
                             will be preserved during cleanup.
                             If mod(iter,preserve_model_interval) == 0
                             model will be preserved.""")

    parser.add_argument("--cleanup", type=str, action=utils.StrToBoolAction,
                        choices=["true", "false"], default=True,
                        help="Clean up models after training")

    parser.add_argument("--max-param-change", type=float, dest='max_param_change',
                        default=2.0, help="""The maximum change in parameters
                             allowed per minibatch, measured in Frobenius
                             norm over the entire model""")

    parser.add_argument("--proportional-shrink", type=float, dest='proportional_shrink',
                        default=0.0, help="""If nonzero, this will set a shrinkage (scaling)
                        factor for the parameters, whose value is set as:
                        shrink-value=(1.0 - proportional-shrink * learning-rate), where
                        'learning-rate' is the learning rate being applied
                        on the current iteration, which will vary from
                        initial-effective-lrate*num-jobs-initial to
                        final-effective-lrate*num-jobs-final.
                        Unlike for train_rnn.py, this is applied unconditionally,
                        it does not depend on saturation of nonlinearities.
                        Can be used to roughly approximate l2 regularization.""")

    parser.add_argument("--stage", type=int, default=-4,
                        help="Specifies the stage of the experiment to execution from")

    parser.add_argument("--cmd", type=str, dest="command",
                        action=utils.NullStrToNoneAction,
                        help="""Specifies the script to launch jobs.
                             e.g. queue.pl for launching on SGE cluster
                                    run.pl for launching on local machine
                             """, default="queue.pl")

    parser.add_argument("--max-models-combine", type=int, dest='max_models_combine',
                        default=20, help="""The maximum number of models used in
                                 the final model combination stage.  These
                                 models will themselves be averages of
                                 iteration-number ranges""")

    print(' '.join(sys.argv))

    args = parser.parse_args()

    [args, run_opts] = process_args(args)

    return [args, run_opts]


def process_args(args):
    """ Process the options got from get_args()
    """

    if not os.path.exists(args.dir):
        raise Exception("This scripts expects {0} to exist.".format(args.dir))

    try:
        if args.tf_model_class == "ModelWithoutDropoutLadderTdnn":
            eval('ladder.%s()' % args.tf_model_class)
        else:
            eval('models.%s()' % args.tf_model_class)
    except AttributeError:
        raise Exception("The specified class name {0} does not exist.".format(args.tf_model_class))

    # set the options corresponding to args.use_gpu
    run_opts = utils.RunOpts()
    if args.use_gpu in ["yes", "wait"]:
        run_opts.train_queue_opt = "--gpu 1"
        run_opts.parallel_train_opts = "--use-gpu={}".format(args.use_gpu)
        run_opts.combine_gpu_opt = "--use-gpu={}".format(args.use_gpu)
        run_opts.combine_queue_opt = "--gpu 1"
        run_opts.prior_gpu_opt = "--use-gpu={}".format(args.use_gpu)
        run_opts.prior_queue_opt = "--gpu 1"
    else:
        logger.warning("Without using a GPU this will be very slow. "
                       "nnet3 does not yet support multiple threads.")
        run_opts.train_queue_opt = ""
        run_opts.parallel_train_opts = "--use-gpu=no"
        run_opts.combine_gpu_opt = "--use-gpu=no"
        run_opts.combine_queue_opt = ""
        run_opts.prior_gpu_opt = "--use-gpu=no"
        run_opts.prior_queue_opt = ""

    run_opts.command = args.command
    # run_opts.num_jobs_compute_prior = args.num_jobs_compute_prior

    return [args, run_opts]


def train_new_models(model_dir, _iter, random_seed, num_jobs,
                     num_archives_processed, num_archives,
                     learning_rate, shrinkage_value, dropout_proportion, egs_dir,
                     momentum, max_param_change, minibatch_size,
                     run_opts, feature_dim, archives_minibatch_count, try_count=0, train_opts=""):
    """ Called from train_one_iteration(), this model does one iteration of
    training with 'num_jobs' jobs, and writes models in dirs like
    exp/tdnn_a/model_24.{1,2,3,..<num_jobs>}
    We cannot easily use a single parallel SGE job to do the main training,
    because the computation of which archive and which --frame option
    to use for each job is a little complex, so we spawn each one separately.
    """

    threads = []

    # the GPU timing info is only printed if we use the --verbose=1 flag; this
    # slows down the computation slightly, so don't accumulate it on every
    # iteration.  Don't do it on iteration 0 either, because we use a smaller
    # than normal minibatch size, and people may get confused thinking it's
    # slower for iteration 0 because of the verbose option.
    verbose_opt = ("--verbose=1" if _iter % 20 == 0 and _iter > 0 else "")

    for job in range(1, num_jobs + 1):
        # k is a zero-based index that we will derive the other indexes from.
        k = num_archives_processed + job - 1

        # work out the 1-based archive index.
        archive_index = (k % num_archives) + 1
        minibatch_count = archives_minibatch_count[archive_index]

        if try_count > 0 and utils.is_correct_model_dir('{0}/model_{1}.{2}'.format(model_dir, _iter + 1, job)):
            continue

        egs_rspecifier = \
            '--ranges-file="{egs_dir}/temp/ranges.{archive_index}" ' \
            '--scp-file="{egs_dir}/temp/feats.scp.{archive_index}" ' \
            '--shuffle=True --minibatch-size={minibatch_size}'.format(
                egs_dir=egs_dir, archive_index=archive_index,
                minibatch_size=minibatch_size)

        # check whether tar file exist or not. If it was generated, so lets pass it to the script for speedup
        tar_file = '{egs_dir}/egs.{archive_index}.tar'.format(egs_dir=egs_dir, archive_index=archive_index)
        if os.path.exists(tar_file):
            egs_rspecifier = '--tar-file="{0}" {1}'.format(tar_file, egs_rspecifier)

        _command = '{command} {train_queue_opt} {dir}/log/train.{iter}.{job}.log ' \
                   'local/tf/train_dnn_one_iteration.py ' \
                   '{parallel_train_opts} ' \
                   '{verbose_opt} --print-interval=10 ' \
                   '--momentum={momentum} ' \
                   '--max-param-change={max_param_change} ' \
                   '--l2-regularize-factor={l2_regularize_factor} ' \
                   '--random-seed={random_seed} {train_opts} ' \
                   '--learning-rate={learning_rate} ' \
                   '--scale={shrinkage_value} ' \
                   '--minibatch-count={minibatch_count} ' \
                   '--feature-dim={feature_dim} ' \
                   '--dropout-proportion={dropout_proportion} ' \
                   '{egs_rspecifier} ' \
                   '--input-dir={dir}/model_{iter} ' \
                   '--output-dir={dir}/model_{next_iter}.{job}' \
            .format(command=run_opts.command,
                    train_queue_opt=run_opts.train_queue_opt,
                    dir=model_dir, iter=_iter,
                    next_iter=_iter + 1, random_seed=_iter + random_seed,
                    job=job,
                    parallel_train_opts=run_opts.parallel_train_opts,
                    verbose_opt=verbose_opt,
                    momentum=momentum, max_param_change=max_param_change,
                    l2_regularize_factor=1.0 / num_jobs,
                    train_opts=train_opts,
                    learning_rate=learning_rate,
                    shrinkage_value=shrinkage_value,
                    minibatch_count=minibatch_count,
                    feature_dim=feature_dim,
                    dropout_proportion=dropout_proportion,
                    egs_rspecifier=egs_rspecifier)

        thread = utils.background_command(_command, require_zero_status=False)
        threads.append(thread)

    for thread in threads:
        thread.join()


def train_one_iteration(model_dir, _iter, random_seed, egs_dir,
                        num_jobs, num_archives_processed, num_archives,
                        learning_rate, minibatch_size,
                        momentum, max_param_change, run_opts,
                        feature_dim, archives_minibatch_count,
                        shrinkage_value=1.0, current_dropout=0.0):
    """ Called from train for one iteration of neural network training
    Selected args:
        shrinkage_value: If value is 1.0, no shrinkage is done; otherwise
            parameter values are scaled by this value.
    """

    # check if different iterations use the same random seed
    random_seed_file = '{0}/random_seed'.format(model_dir)
    if os.path.exists(random_seed_file):
        try:
            with open(random_seed_file, 'r') as fid:
                saved_random_seed = int(fid.readline().strip())
        except (IOError, ValueError):
            logger.error("Exception while reading the random seed for training")
            raise
        if random_seed != saved_random_seed:
            logger.warning("The random seed provided to this iteration (random_seed={0}) is "
                           "different from the one saved last time (random_seed={1}). "
                           "Using random_seed={0}.".format(random_seed, saved_random_seed))
    else:
        with open(random_seed_file, 'w') as fid:
            fid.write(str(random_seed))

    # Sets off some background jobs to compute train and validation set objectives
    eval_trained_dnn(model_dir, _iter, egs_dir, run_opts)

    new_model = "{0}/model_{1}/model.meta".format(model_dir, _iter + 1)
    if utils.is_correct_model_dir("{0}/model_{1}".format(model_dir, _iter + 1)):
        logger.info('The output model {0} was exist and so I do not continue this iteration.'.format(new_model))
        return

    do_average = (_iter > 0)

    dropout_proportion = current_dropout

    if do_average or True:  # TODO
        cur_minibatch_size = minibatch_size
        cur_max_param_change = max_param_change
    else:
        # on iteration zero, use a smaller minibatch size (and we will later
        # choose the output of just one of the jobs): the model-averaging isn't
        # always helpful when the model is changing too fast (i.e. it can worsen
        # the objective function), and the smaller minibatch size will help to
        # keep the update stable.
        cur_minibatch_size = minibatch_size / 2
        cur_max_param_change = float(max_param_change) / math.sqrt(2)

    try_count = 0
    training_flag = True
    while try_count < MAX_TRY_COUNT and training_flag:
        train_new_models(model_dir=model_dir, _iter=_iter, random_seed=random_seed, num_jobs=num_jobs,
                         num_archives_processed=num_archives_processed,
                         num_archives=num_archives, learning_rate=learning_rate,
                         shrinkage_value=shrinkage_value, dropout_proportion=dropout_proportion,
                         egs_dir=egs_dir, momentum=momentum, max_param_change=cur_max_param_change,
                         minibatch_size=cur_minibatch_size, feature_dim=feature_dim, try_count=try_count,
                         archives_minibatch_count=archives_minibatch_count, run_opts=run_opts)
        try_count += 1
        training_flag = False
        cnt = 0
        for job in range(1, num_jobs + 1):
            if not utils.is_correct_model_dir("{0}/model_{1}.{2}".format(model_dir, _iter + 1, job)):
                # move logs to prevent rewriting.
                log_file = "{0}/log/train.{1}.{2}.log".format(model_dir, _iter, job)
                if os.path.exists(log_file):
                    os.rename(log_file, "{0}.{1}".format(log_file, try_count))
                training_flag = True
                cnt += 1
        if training_flag:
            if try_count < MAX_TRY_COUNT:
                logger.warn("{0}/{1} of jobs failed. Resubmitting them may solved the problem. "
                            "Start resubmitting them after 30 seconds ...".format(cnt, num_jobs))
                # sleep for 30 seconds before resubmitting
                time.sleep(30)
            else:
                logger.error("{0}/{1} of jobs failed and maximum number of retrying is reached. "
                             "Stop the training ...".format(cnt, num_jobs))

    if training_flag:
        raise Exception("Some training jobs failed more than %d times. "
                        "Please check the log files." % MAX_TRY_COUNT)

    [models_to_average, best_model] = utils.get_successful_models(
        num_jobs, '{0}/log/train.{1}.%.log'.format(model_dir, _iter))

    if do_average and len(models_to_average) > 1:
        # average the output of the different jobs.
        nets_dirs = []
        for n in models_to_average:
            nets_dirs.append("{0}/model_{1}.{2}".format(model_dir, _iter + 1, n))
        # utils.get_average_nnet_model(
        #     dir=model_dir, iter=_iter,
        #     nnets_list=" ".join(nets_dirs),
        #     run_opts=run_opts)
        output_dir = os.path.join(model_dir,"model_"+str(_iter+1))
        if not os.path.exists(output_dir):
            os.mkdir(output_dir)
        avg_command = "python2 local/contrib/avg_checkpoint.py --checkpoints {checkpoints} --num_last_checkpoints {" \
                      "num_check} --output_path {output_path}".format(checkpoints=",".join(nets_dirs),num_check=int(
                      len(nets_dirs)),output_path=os.path.join(output_dir,'model'))
        print(avg_command)
        p = subprocess.Popen(avg_command, shell=True, stdout=subprocess.PIPE)
        [stdout, stderr] = p.communicate()
        if p.returncode is not 0:
            output = "Command exited with status {0}: {1}".format(p.returncode, stderr.decode('utf-8'))
            raise Exception(output)
        # rename averaged model
        for avg_file in glob.glob(os.path.join(output_dir,'model-*')):
            avg_file_name, avg_file_ext = os.path.splitext(avg_file)
            new_avg_filename = re.sub('/model-.*','/model',avg_file_name)
            os.rename(avg_file,new_avg_filename+avg_file_ext)
        # copy model_*.1 files
        assert os.path.exists(os.path.join(model_dir,"model_"+str(_iter+1)+".1")), "{0} not exists".format(
            os.path.join(model_dir,"model_"+str(_iter+1)+".1"))
        # shutil.copy("{model_dir}/model_{iter}.1/checkpoint".format(model_dir=model_dir, iter=_iter + 1),output_dir)
        shutil.copy("{model_dir}/model_{iter}.1/done".format(model_dir=model_dir, iter=_iter + 1),output_dir)
        # sys.exit()

    else:
        # choose the best model from different jobs
        utils.copy_best_nnet_dir(_dir=model_dir, _iter=_iter, best_model_index=best_model)

    try:
        for i in range(1, num_jobs + 1):
            shutil.rmtree("{0}/model_{1}.{2}".format(model_dir, _iter + 1, i))
    except OSError:
        logger.error("Error while trying to delete the client models.")
        raise

    if not os.path.isfile(new_model):
        raise Exception("Could not find {0}, at the end of iteration {1}".format(new_model, _iter))
    elif os.stat(new_model).st_size == 0:
        raise Exception("{0} has size 0. Something went wrong in iteration {1}".format(new_model, _iter))


def eval_trained_dnn(main_dir, _iter, egs_dir, run_opts):
    input_model_dir = "{dir}/model_{iter}".format(dir=main_dir, iter=_iter)

    # we assume that there are just one tar file for validation
    tar_file = ("{0}/valid_egs.1.tar".format(egs_dir))

    _command = '{command} "{main_dir}/log/compute_prob_valid.{iter}.log" ' \
               'local/tf/eval_dnn.py ' \
               '--tar-file="{tar_file}" --use-gpu=no ' \
               '--log-file="{main_dir}/log/compute_prob_valid.{iter}.log" ' \
               '--input-dir="{input_model_dir}"'.format(command=run_opts.command,
                                                        main_dir=main_dir,
                                                        iter=_iter,
                                                        tar_file=tar_file,
                                                        input_model_dir=input_model_dir)

    utils.background_command(_command)

    # we assume that there are just one tar file for train diagnostics
    tar_file = ("{0}/train_subset_egs.1.tar".format(egs_dir))

    _command = '{command} "{main_dir}/log/compute_prob_train_subset.{iter}.log" ' \
               'local/tf/eval_dnn.py ' \
               '--tar-file="{tar_file}" --use-gpu=no ' \
               '--log-file="{main_dir}/log/compute_prob_train_subset.{iter}.log" ' \
               '--input-dir="{input_model_dir}"'.format(command=run_opts.command,
                                                        main_dir=main_dir,
                                                        iter=_iter,
                                                        tar_file=tar_file,
                                                        input_model_dir=input_model_dir)

    utils.background_command(_command)


def train(args, run_opts):
    """ The main function for training.
    Args:
        args: a Namespace object with the required parameters
            obtained from the function process_args()
        run_opts: RunOpts object obtained from the process_args()
    """

    arg_string = pprint.pformat(vars(args))
    logger.info("Arguments for the experiment\n{0}".format(arg_string))

    default_egs_dir = '{0}/egs'.format(args.dir)
    if args.egs_dir is None:
        egs_dir = default_egs_dir
    else:
        egs_dir = args.egs_dir

    [num_archives, egs_feat_dim, archives_minibatch_count] = utils.verify_egs_dir(egs_dir)

    if args.num_jobs_final > num_archives:
        raise Exception('num_jobs_final cannot exceed the number of archives '
                        'in the egs directory')

    if args.stage <= -1:
        if not os.path.exists('{0}/model_0/done'.format(args.dir)):
            logger.info("Preparing the initial network.")
            num_classes = args.num_targets
            model_name = args.tf_model_class
            if model_name == "ModelWithoutDropoutLadderTdnn":
                model = eval('ladder.%s()' % model_name)
            else:
                model = eval('models.%s()' % model_name)
            logger.info("Start calling build_model to initialize the model %s ..." % model_name)
            model.build_model(num_classes, egs_feat_dim, '{0}/model_0'.format(args.dir), logger=logger)
            with open(os.path.join(args.dir, 'model_name.txt'), 'wt') as fid:
                fid.write(model_name)
        else:
            logger.info("The initial network exist from before.")

    # set num_iters so that as close as possible, we process the data
    # $num_epochs times, i.e. $num_iters * $avg_num_jobs) ==
    # $num_epochs * $num_archives, where
    # avg_num_jobs = (num_jobs_initial + num_jobs_final) / 2.
    num_archives_to_process = int(args.num_epochs * num_archives)
    num_archives_processed = 0
    num_iters = int(((num_archives_to_process * 2) / (args.num_jobs_initial +
        args.num_jobs_final)))
    # If do_final_combination is True, compute the set of models_to_combine.
    # Otherwise, models_to_combine will be none.
    if args.do_final_combination:
        models_to_combine = utils.get_model_combine_iters(num_iters, num_archives, args.max_models_combine,
                                                          args.num_jobs_final)
    else:
        models_to_combine = None

    logger.info("Training will run for {0} epochs = {1} iterations".format(args.num_epochs, num_iters))

    for _iter in range(num_iters):
        current_num_jobs = int(0.5 + args.num_jobs_initial
                               + (args.num_jobs_final - args.num_jobs_initial)
                               * float(_iter) / num_iters)

        if args.stage <= _iter:
            lrate = utils.get_learning_rate(_iter, current_num_jobs,
                                            num_iters,
                                            num_archives_processed,
                                            num_archives_to_process,
                                            args.initial_effective_lrate,
                                            args.final_effective_lrate)

            shrinkage_value = 1.0 - (args.proportional_shrink * lrate)
            if shrinkage_value <= 0.5:
                raise Exception("proportional-shrink={0} is too large, it gives "
                                "shrink-value={1}".format(args.proportional_shrink,
                                                          shrinkage_value))

            percent = num_archives_processed * 100.0 / num_archives_to_process
            epoch = (num_archives_processed * args.num_epochs / num_archives_to_process)
            shrink_info_str = ''
            if shrinkage_value != 1.0:
                shrink_info_str = 'shrink: {0:0.5f}'.format(shrinkage_value)
            logger.info("Iter: {0}/{1}    Epoch: {2:0.2f}/{3:0.1f} ({4:0.1f}% complete)    "
                        "lr: {5:0.6f}    {6}".format(_iter, num_iters - 1, epoch, args.num_epochs,
                                                     percent, lrate, shrink_info_str))
            current_dropout = utils.get_dropout_edit_string(
                args.dropout_schedule, float(num_archives_processed) / num_archives_to_process)
            train_one_iteration(
                model_dir=args.dir,
                _iter=_iter,
                random_seed=args.random_seed,
                egs_dir=egs_dir,
                num_jobs=current_num_jobs,
                num_archives_processed=num_archives_processed,
                num_archives=num_archives,
                learning_rate=lrate,
                current_dropout=current_dropout,
                minibatch_size=args.minibatch_size,
                momentum=args.momentum,
                max_param_change=args.max_param_change,
                shrinkage_value=shrinkage_value,
                feature_dim=egs_feat_dim,
                archives_minibatch_count=archives_minibatch_count,
                run_opts=run_opts)

            if args.cleanup:
                # do a clean up everything but the last 2 models, under certain conditions
                utils.remove_model(args.dir, _iter - 2, models_to_combine, args.preserve_model_interval)

        num_archives_processed = num_archives_processed + current_num_jobs

    if args.stage <= num_iters:
        # TODO now do_final_combination is false but the default value was True
        if args.do_final_combination:
            logger.info("Doing final combination to produce final.raw")
            raise Exception('combine model using average not implemented yet.')
            # train_lib.common.combine_models(
            #     dir=args.dir, num_iters=num_iters,
            #     models_to_combine=models_to_combine, egs_dir=egs_dir,
            #     minibatch_size_str=args.minibatch_size, run_opts=run_opts,
            #     get_raw_nnet_from_am=False,
            #     max_objective_evaluations=args.max_objective_evaluations)
        else:
            utils.force_symlink("model_{1}".format(args.dir, num_iters), "{0}/model_final".format(args.dir))

    if args.cleanup:
        logger.info("Cleaning up the experiment directory {0}".format(args.dir))
        for _iter in range(num_iters):
            utils.remove_model(args.dir, _iter, models_to_combine, args.preserve_model_interval)

    # do some reporting
    [report, _, _] = utils.generate_report(args.dir)
    with open("{dir}/accuracy.report".format(dir=args.dir), "wt") as fid:
        fid.write(report)


def main():
    [args, run_opts] = get_args()
    try:
        train(args, run_opts)
        utils.wait_for_background_commands()
    except BaseException as e:
        # look for BaseException so we catch KeyboardInterrupt, which is
        # what we get when a background thread dies.
        if not isinstance(e, KeyboardInterrupt):
            traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    # sys.argv = ["train_dnn.py","--stage=47","--tf-model-class=ModelWithoutDropoutTdnn",
    #             "--cmd=/home/feit/Tool/kaldi/egs/wsj/s5/utils/parallel/run.pl",
    #         "--num-targets=7323", "--proportional-shrink=10", "--minibatch-size=64", "--max-param-change=2",
    #         "--momentum=0.5", "--num-jobs-initial=1", "--num-jobs-final=2", "--initial-effective-lrate=0.001",
    #         "--final-effective-lrate=0.0001", "--random-seed=2468", "--num-epochs=3",
    #             "--dropout-schedule=0,0@0.20,0.1@0.50,0",
    #             "--egs-dir=/media/feit/Work/Work/SpeakerID/Kaldi_Voxceleb/exp_clean/xvector_tf_but_test/egs",
    #             "--preserve-model-interval=10", "--use-gpu=yes",
    #             "--dir=/media/feit/Work/Work/SpeakerID/Kaldi_Voxceleb/exp_clean/xvector_tf_but_test"]
    main()