#!/usr/bin/env bash

# This script prepares some things needed by rnnlm/train_rnnlm.sh, e.g. it initializes
# the model and ensures we have the split-up integerized data on disk.

cmd=run.pl
stage=0
sampling=true            # add the option --sampling false to disable creation
                         # of sampling.lm
words_per_split=5000000  # aim for training on 5 million words per job.  the
                         # aim is to have the jobs last at least a couple of
                         # minutes.  If this leads to having just one job,
                         # we repeat the archive as many times as needed to
                         # get the target length.
unigram_factor=100.0     # Option used when pruning the LM used for sampling.
                         # You can increase this, e.g. to 200 or 400, if the LM used
                         # for sampling is too big, causing rnnlm-get-egs to
                         # take up too many CPUs worth of compute.

. utils/parse_options.sh

if [ $# != 3 ]; then
  echo "Usage: $0 [options] <text-dir> <rnnlm-config-dir> <rnnlm-dir>"
  echo "Sets up the directory <rnnlm-dir> for RNNLM training as done by"
  echo "rnnlm/train_rnnlm.sh, and initializes the model."
  echo " <text-dir> is as validated by rnnlm/validate_text_dir.py"
  echo " <rnnlm-config-dir> is as validated by rnnlm/validate_config_dir.sh."
  exit 1
fi


text_dir=$1
config_dir=$2
dir=$3

set -e
. ./path.sh

if [ $stage -le 0 ]; then
  echo "$0: validating input"

  rnnlm/validate_text_dir.py --spot-check=true $text_dir

  rnnlm/validate_config_dir.sh $text_dir $config_dir

  if ! mkdir -p $dir; then
    echo "$0: could not create RNNLM dir $dir"
  fi
fi

if [ $stage -le 1 ]; then
  if [ $config_dir != $dir/config ]; then
    echo "$0: copying config directory"
    mkdir -p $dir/config
    # copy expected things from $config_dir to $dir/config.
    for f in words.txt data_weights.txt oov.txt xconfig; do
      cp $config_dir/$f $dir/config
    done
    # features.txt is optional, check separately
    if [ -f $config_dir/features.txt ]; then
      cp $config_dir/features.txt $dir/config
    fi
  fi

  rnnlm/get_special_symbol_opts.py < $dir/config/words.txt > $dir/special_symbol_opts.txt
fi

if [ $stage -le 2 ]; then
  if [ ! -f $text_dir/dev.counts ] || [ $text_dir/dev.counts -ot $text_dir/dev.txt ]; then
    echo "$0: preparing unigram counts in $text_dir"
    rnnlm/get_unigram_counts.sh $text_dir
  fi
fi


if [ $stage -le 3 ]; then
  if [ -f $dir/config/features.txt ]; then
    # prepare word-level features in $dir/word_feats.txt
    # first we need the appropriately weighted unigram counts.

    if awk '{if($2 == "unigram"){saw_unigram=1;}} END{exit(saw_unigram ? 0 : 1)}' $dir/config/features.txt; then
      # we need the unigram probabilities
      rnnlm/ensure_counts_present.sh $text_dir

      rnnlm/get_unigram_probs.py --vocab-file=$dir/config/words.txt \
        --unk-word=$(cat $dir/config/oov.txt) \
        --data-weights-file=$dir/config/data_weights.txt $text_dir \
        >$dir/unigram_probs.txt
      unigram_opt="--unigram-probs=$dir/unigram_probs.txt"
    else
      unigram_opt=
    fi
    rnnlm/get_word_features.py $unigram_opt --treat-as-bos='#0' \
      $dir/config/words.txt $dir/config/features.txt >$dir/word_feats.txt
  else
    [ -f $dir/word_feats.txt ] && rm $dir/word_feats.txt
  fi
fi

if [ $stage -le 4 ]; then
  echo "$0: preparing split-up integerized text data with weights"
  num_repeats=1
  num_splits=$(rnnlm/get_num_splits.sh $words_per_split $text_dir $dir/config/data_weights.txt)
  if [ $num_splits -lt 1 ]; then
    # the script outputs the negative of the num-repeats if the num-splits is 1
    # and the num-repeats is >1.
    num_repeats=$[-num_splits]
    num_splits=1
  fi

  # note: the python script treats the empty unknown word as a special case,
  # so if oov.txt is empty we don't have to take any special action.
  rnnlm/prepare_split_data.py --unk-word="$(cat $dir/config/oov.txt)" \
     --vocab-file=$dir/config/words.txt --data-weights-file=$dir/config/data_weights.txt \
     --num-splits=$num_splits $text_dir $dir/text
  echo $num_repeats >$dir/text/info/num_repeats
fi

if [ $stage -le 5 ]; then
  echo "$0: initializing neural net"
  mkdir -p $dir/config/nnet

  steps/nnet3/xconfig_to_configs.py --xconfig-file=$dir/config/xconfig \
    --config-dir=$dir/config/nnet

  # initialize the neural net.
  nnet3-init $dir/config/nnet/ref.config $dir/0.raw
fi

embedding_dim=$(rnnlm/get_embedding_dim.py $dir/0.raw)

if [ $stage -le 6 ]; then
  echo "$0: initializing embedding matrix"
  if [ -f $dir/config/features.txt ]; then
    # We are using sparse features.
    feat_dim=$(tail -n 1 $dir/config/features.txt | awk '{print $1 + 1;}')

    first_element_opt=
    if grep -q '0\tconstant' $dir/config/features.txt; then
      first_element_opt="--first-element 1.0"
    fi
    # we'll probably make the stddev configurable soon, or maybe just remove it.
    # At some point stability was an issue.
    rnnlm/initialize_matrix.pl $first_element_opt --stddev 0.001 \
      $feat_dim $embedding_dim > $dir/feat_embedding.0.mat

    [ -f $dir/word_embedding.0.mat ] && rm $dir/word_embedding.0.mat
  else
    vocab_size=$(tail -n 1 $dir/config/words.txt | awk '{print $NF + 1}')
    rnnlm/initialize_matrix.pl --first-column 1.0 $vocab_size $embedding_dim > $dir/word_embedding.0.mat

    [ -f $dir/feat_embedding.0.mat ] && rm $dir/feat_embedding.0.mat
  fi
fi

mkdir -p $dir/log

if [ $stage -le 7 ]; then
  if $sampling; then
    echo "$0: preparing language model for sampling"
    num_splits=$(cat $dir/text/info/num_splits)
    text_files=$(for n in $(seq $num_splits); do echo -n $dir/text/$n.txt ''; done)
    vocab_size=$(tail -n 1 $dir/config/words.txt | awk '{print $NF + 1}')

    special_symbol_opts=$(cat $dir/special_symbol_opts.txt)

    # this prints some nontrivial log information, so run using '$cmd' to ensure
    # the output gets saved.
    # ***NOTE*** we will likely later have to pass in options to this program to control
    # the size of the sampling LM.
    $cmd $dir/log/prepare_sampling_lm.log \
         rnnlm-get-sampling-lm --unigram-factor=$unigram_factor $special_symbol_opts \
              --vocab-size=$vocab_size  "cat $text_files|" $dir/sampling.lm
    echo "$0: done estimating LM for sampling."
  else
    [ -f $dir/sampling.lm ] && rm $dir/sampling.lm
  fi
fi