#!/usr/bin/env bash

# This script does the RNNLM training.  It assumes you have already run
# 'prepare_rnnlm_dir.sh' to prepare the directory.


#num-jobs-initial, num-jobs-final, max-change, embedding-max-change [initial,final?],
#num-samples, minibatch-size, chunk-length, [and the same for dev data]...
#initial-effective-learning-rate, final-effective-learning-rate, ...
#embedding-learning-rate-factor, num-epochs


stage=0
num_jobs_initial=1
num_jobs_final=1
rnnlm_max_change=0.5
embedding_max_change=0.5
chunk_length=32
num_epochs=100  # maximum number of epochs to train.  later we
                # may find a stopping criterion.
initial_effective_lrate=0.001
final_effective_lrate=0.0001
embedding_l2=0.005
embedding_lrate_factor=0.1  # the embedding learning rate is the
                            # nnet learning rate times this factor.
backstitch_training_scale=0.0    # backstitch training scale
backstitch_training_interval=1   # backstitch training interval
cmd=run.pl  # you might want to set this to queue.pl

# some options passed into rnnlm-get-egs, relating to sampling.
num_samples=512
sample_group_size=2  # see rnnlm-get-egs
num_egs_threads=10  # number of threads used for sampling, if we're using
                    # sampling.  the actual number of threads that runs at one
                    # time, will be however many is needed to balance the
                    # sampling and the actual training, this is just the maximum
                    # possible number that are allowed to run
use_gpu=true  # use GPU for training
use_gpu_for_diagnostics=false  # set true to use GPU for compute_prob_*.log

# optional cleanup options
cleanup=false  # add option --cleanup true to enable automatic cleanup of old models
cleanup_strategy="keep_latest"  # determines cleanup strategy, use either "keep_latest" or "keep_best"
cleanup_keep_iters=3  # number of iterations that will have their models retained

trap 'for pid in $(jobs -pr); do kill -KILL $pid; done' INT QUIT TERM
. utils/parse_options.sh

if [ $# != 1 ]; then
  echo "Usage: $0 [options] <rnnlm-dir>"
  echo "Trains an RNNLM, assuming the things needed for training have already been"
  echo "set up by prepare_rnnlm_dir.sh."
  exit 1
fi


dir=$1


set -e
. ./path.sh


for f in $dir/config/{words,data_weights,oov}.txt \
              $dir/text/1.txt $dir/text/dev.txt $dir/0.raw \
              $dir/text/info/num_splits $dir/text/info/num_repeats \
              $dir/special_symbol_opts.txt; do
  [ ! -f $f ] && echo "$0: expected $f to exist" && exit 1
done

# set some variables and check more files.
num_splits=$(cat $dir/text/info/num_splits)
num_repeats=$(cat $dir/text/info/num_repeats)
text_files=$(for n in $(seq $num_splits); do echo $dir/text/$n.txt; done)
vocab_size=$(tail -n 1 $dir/config/words.txt | awk '{print $NF + 1}')
embedding_type=

if [ -f $dir/feat_embedding.0.mat ]; then
  sparse_features=true
  embedding_type=feat
  if [ -f $dir/word_embedding.0.mat ]; then
    echo "$0: error: $dir/feat_embedding.0.mat and $dir/word_embedding.0.mat both exist."
    exit 1;
  fi
  ! [ -f $dir/word_feats.txt ] && echo "$0: expected $0/word_feats.txt to exist" && exit 1;
else
  sparse_features=false
  embedding_type=word
  ! [ -f $dir/word_embedding.0.mat ] && \
    echo "$0: expected $dir/word_embedding.0.mat to exist" && exit 1
fi

if [ $num_jobs_initial -gt $num_splits ] || [ $num_jobs_final -gt $num_splits ]; then
  echo -n "$0: number of initial or final jobs $num_jobs_initial/$num_jobs_final"
  echo "exceeds num-splits=$num_splits; reduce number of jobs"
  exit 1
fi

num_splits_to_process=$[($num_epochs*$num_splits)/$num_repeats]
num_split_processed=0
num_iters=$[($num_splits_to_process*2)/($num_jobs_initial+$num_jobs_final)]


# this string will combine options and arguments.
train_egs_args="--vocab-size=$vocab_size $(cat $dir/special_symbol_opts.txt)"
if [ -f $dir/sampling.lm ]; then
  # we are doing sampling.
  train_egs_args="$train_egs_args --num-samples=$num_samples --sample-group-size=$sample_group_size --num-threads=$num_egs_threads $dir/sampling.lm"
fi

echo "$0: will train for $num_iters iterations"

# recording some configuration information
cat >$dir/info.txt <<EOF
num_iters=$num_iters
num_epochs=$num_epochs
num_jobs_initial=$num_jobs_initial
num_jobs_final=$num_jobs_final
rnnlm_max_change=$rnnlm_max_change
embedding_max_change=$embedding_max_change
chunk_length=$chunk_length
initial_effective_lrate=$initial_effective_lrate
final_effective_lrate=$final_effective_lrate
embedding_lrate_factor=$embedding_lrate_factor
sample_group_size=$sample_group_size
num_samples=$num_samples
backstitch_training_scale=$backstitch_training_scale
backstitch_training_interval=$backstitch_training_interval
EOF


x=0
num_splits_processed=0
while [ $x -lt $num_iters ]; do

  this_num_jobs=$(perl -e "print int(0.5+$num_jobs_initial+($num_jobs_final-$num_jobs_initial)*$x/$num_iters);")
  ilr=$initial_effective_lrate; flr=$final_effective_lrate; np=$num_splits_processed; nt=$num_splits_to_process;
  this_learning_rate=$(perl -e "print (($x + 1 >= $num_iters ? $flr : $ilr*exp($np*log($flr/$ilr)/$nt))*$this_num_jobs);");
  embedding_lrate=$(perl -e "print ($this_learning_rate*$embedding_lrate_factor);")

  if [ $stage -le $x ]; then

    # Set off the diagnostic job in the background.
    if $sparse_features; then
      word_embedding="rnnlm-get-word-embedding $dir/word_feats.txt $dir/feat_embedding.$x.mat -|"
    else
      word_embedding="$dir/word_embedding.$x.mat"
    fi
    if $use_gpu_for_diagnostics; then queue_gpu_opt="--gpu 1"; gpu_opt="--use-gpu=yes";
    else gpu_opt=''; queue_gpu_opt=''; fi
    backstitch_opt="--rnnlm.backstitch-training-scale=$backstitch_training_scale \
      --rnnlm.backstitch-training-interval=$backstitch_training_interval \
      --embedding.backstitch-training-scale=$backstitch_training_scale \
      --embedding.backstitch-training-interval=$backstitch_training_interval"
    [ -f $dir/.error ] && rm $dir/.error
    $cmd $queue_gpu_opt $dir/log/compute_prob.$x.log \
       rnnlm-get-egs $(cat $dir/special_symbol_opts.txt) \
                     --vocab-size=$vocab_size $dir/text/dev.txt ark:- \| \
       rnnlm-compute-prob $gpu_opt $dir/$x.raw "$word_embedding" ark:- || touch $dir/.error &

    if [ $x -gt 0 ]; then
      $cmd $dir/log/progress.$x.log \
        nnet3-show-progress --use-gpu=no $dir/$[$x-1].raw $dir/$x.raw '&&' \
          nnet3-info $dir/$x.raw &
    fi

    echo "Training neural net (pass $x)"


    ( # this sub-shell is so that when we "wait" below,
      # we only wait for the training jobs that we just spawned,
      # not the diagnostic jobs that we spawned above.

      # We can't 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.
      [ -f $dir/.train_error ] && rm $dir/.train_error
      for n in $(seq $this_num_jobs); do
        k=$[$num_splits_processed + $n - 1]; # k is a zero-based index that we'll derive
                                               # the other indexes from.
        split=$[($k%$num_splits)+1]; # work out the 1-based split index.

        src_rnnlm="nnet3-copy --learning-rate=$this_learning_rate $dir/$x.raw -|"
        if $sparse_features; then
          sparse_opt="--read-sparse-word-features=$dir/word_feats.txt";
          embedding_type=feat
        else
          sparse_opt=''; embedding_type=word
        fi
        if $use_gpu; then gpu_opt="--use-gpu=yes"; queue_gpu_opt="--gpu 1";
        else gpu_opt="--use-gpu=no"; queue_gpu_opt=""; fi
        if [ $this_num_jobs -gt 1 ]; then dest_number=$[x+1].$n
        else dest_number=$[x+1]; fi
        # in the normal case $repeated data will be just one copy.
        repeated_data=$(for n in $(seq $num_repeats); do echo -n $dir/text/$split.txt ''; done)

        rnnlm_l2_factor=$(perl -e "print (1.0/$this_num_jobs);")
        embedding_l2_regularize=$(perl -e "print ($embedding_l2/$this_num_jobs);")

        # allocate queue-slots for threads doing sampling,
        num_threads_=$[$num_egs_threads*2/3]
        [ -f $dir/sampling.lm ] && queue_thread_opt="--num-threads $num_threads_" || queue_thread_opt=

        # Run the training job or jobs.
        $cmd $queue_gpu_opt $queue_thread_opt $dir/log/train.$x.$n.log \
           rnnlm-train \
             --rnnlm.max-param-change=$rnnlm_max_change \
             --rnnlm.l2_regularize_factor=$rnnlm_l2_factor \
             --embedding.max-param-change=$embedding_max_change \
             --embedding.learning-rate=$embedding_lrate \
             --embedding.l2_regularize=$embedding_l2_regularize \
             $sparse_opt $gpu_opt $backstitch_opt \
             --read-rnnlm="$src_rnnlm" --write-rnnlm=$dir/$dest_number.raw \
             --read-embedding=$dir/${embedding_type}_embedding.$x.mat \
             --write-embedding=$dir/${embedding_type}_embedding.$dest_number.mat \
             "ark,bg:cat $repeated_data | rnnlm-get-egs --chunk-length=$chunk_length --srand=$num_splits_processed $train_egs_args - ark:- |" || touch $dir/.train_error &
      done
      wait # wait for just the training jobs.
      [ -f $dir/.train_error ] && \
        echo "$0: failure on iteration $x of training, see $dir/log/train.$x.*.log for details." && exit 1
      if [ $this_num_jobs -gt 1 ]; then
        # average the models and the embedding matrces.  Use run.pl as we don\'t
        # want this to wait on the queue (if there is a queue).
        src_models=$(for n in $(seq $this_num_jobs); do echo $dir/$[x+1].$n.raw; done)
        src_matrices=$(for n in $(seq $this_num_jobs); do echo $dir/${embedding_type}_embedding.$[x+1].$n.mat; done)
        run.pl $dir/log/average.$[x+1].log \
          nnet3-average $src_models $dir/$[x+1].raw '&&' \
          matrix-sum --average=true $src_matrices $dir/${embedding_type}_embedding.$[x+1].mat
      fi
      # optionally, perform cleanup after training
      if [ "$cleanup" = true ] ; then
        python3 rnnlm/rnnlm_cleanup.py $dir --$cleanup_strategy --iters_to_keep $cleanup_keep_iters
      fi
    )
    # the error message below is not that informative, but $cmd will
    # have printed a more specific one.
    [ -f $dir/.error ] && echo "$0: error with diagnostics on iteration $x of training" && exit 1;
  fi

  x=$[x+1]
  num_splits_processed=$[num_splits_processed+this_num_jobs]
done

wait # wait for diagnostic jobs in the background.

if [ $stage -le $num_iters ]; then
  # link the best model we encountered during training (based on
  # dev-set probability) as the final model.
  best_iter=$(rnnlm/get_best_model.py $dir)
  echo "$0: best iteration (out of $num_iters) was $best_iter, linking it to final iteration."
  train_best_log=$dir/log/train.$best_iter.1.log
  ppl_train=`grep 'Overall objf' $train_best_log | awk '{printf("%.1f",exp(-$10))}'`
  dev_best_log=$dir/log/compute_prob.$best_iter.log
  ppl_dev=`grep 'Overall objf' $dev_best_log | awk '{printf("%.1f",exp(-$NF))}'`
  echo "$0: train/dev perplexity was $ppl_train / $ppl_dev."
  ln -sf ${embedding_type}_embedding.$best_iter.mat $dir/${embedding_type}_embedding.final.mat
  ln -sf $best_iter.raw $dir/final.raw
fi

# Now get some diagnostics about the evolution of the objective function.
if [ $stage -le $[num_iters+1] ]; then
  (
    logs=$(for iter in $(seq 1 $[$num_iters-1]); do echo -n $dir/log/train.$iter.1.log ''; done)
    # in the non-sampling case the exact objf is printed and we plot that
    # in the sampling case we print the approximated objf for training.
    grep 'Overall objf' $logs | awk 'BEGIN{printf("Train objf: ")} /exact/{printf("%.2f ", $NF);next} {printf("%.2f ", $10)} END{print "";}'
    logs=$(for iter in $(seq 1 $[$num_iters-1]); do echo -n $dir/log/compute_prob.$iter.log ''; done)
    grep 'Overall objf' $logs | awk 'BEGIN{printf("Dev objf:   ")} {printf("%.2f ", $NF)} END{print "";}'
  ) > $dir/report.txt
  cat $dir/report.txt
fi