#!/bin/bash
  
  # Copyright 2012-2016   Johns Hopkins University (Author: Daniel Povey).  Apache 2.0.
  # Copyright 2014-2015   Vimal Manohar
  
  # Decodes denlats and dumps egs for discriminative training, in one script
  # (avoids writing the non-compact lattices to disk, which can use a lot of disk
  # space).
  
  
  # Begin configuration section.
  cmd=run.pl
  max_copy_jobs=5  # Limit disk I/O
  
  # feature options
  online_ivector_dir=
  
  # example splitting and context options
  frames_per_eg=150 # number of frames of labels per example.
                    # Note: may in general be a comma-separated string of alternative
                    # durations; the first one (the principal num-frames) is preferred.
  frames_overlap_per_eg=30 # number of supervised frames of overlap that we aim for per eg.
                    # can be useful to avoid wasted data if you're using --left-deriv-truncate
                    # and --right-deriv-truncate.
  looped=false       # Set to true to enable looped decoding [can
                     # be a bit faster, for forward-recurrent models like LSTMs.]
  
  # .. these context options also affect decoding.
  extra_left_context=0    # amount of left-context per eg, past what is required by the model
                          # (only useful for recurrent networks like LSTMs/BLSTMs)
  extra_right_context=0   # amount of right-context per eg, past what is required by the model
                          # (only useful for backwards-recurrent networks like BLSTMs)
  extra_left_context_initial=-1    # if >= 0, the --extra-left-context to use at
                                   # the start of utterances.  Recommend 0 if you
                                   # used 0 for the baseline DNN training; if <0,
                                   # defaults to same as extra_left_context
  extra_right_context_final=-1     # if >= 0, the --extra-right-context to use at
                                   # the end of utterances.  Recommend 0 if you
                                   # used 0 for the baseline DNN training; if <0,
                                   # defaults to same as extra_left_context
  
  compress=true   # set this to false to disable lossy compression of features
                  # dumped with egs (e.g. if you want to see whether results are
                  # affected).
  
  num_utts_subset=80     # number of utterances in validation and training
                         # subsets used for diagnostics.
  num_egs_subset=800     # number of egs (maximum) for the validation and training
                         # subsets used for diagnostics.
  frames_per_iter=1000000 # each iteration of training, see this many frames
                          # per job.  This is just a guideline; it will pick a number
                          # that divides the number of samples in the entire data.
  cleanup=true
  
  stage=0
  nj=200
  
  # By default this script uses final.mdl in <srcdir>, this configures it.
  iter=final
  
  
  # decoding-graph option
  self_loop_scale=0.1  # for decoding graph.. should be 1.0 for chain models.
  
  # options relating to decoding.
  frames_per_chunk_decoding=150
  beam=13.0
  lattice_beam=7.0
  acwt=0.1
  max_active=5000
  min_active=200
  max_mem=20000000 # This will stop the processes getting too large.
  # This is in bytes, but not "real" bytes-- you have to multiply
  # by something like 5 or 10 to get real bytes (not sure why so large)
  num_threads=1
  
  # affects whether we invoke lattice-determinize-non-compact after decoding
  # discriminative-get-supervision.
  determinize_before_split=true
  
  
  # End configuration section.
  
  
  echo "$0 $@"  # Print the command line for logging
  
  if [ -f path.sh ]; then . ./path.sh; fi
  . parse_options.sh || exit 1;
  
  
  if [ $# != 5 ]; then
    echo "Usage: $0 [opts] <data> <lang> <src-dir> <ali-dir> <degs-dir>"
    echo " e.g.: $0 data/train data/lang exp/nnet3/tdnn_a exp/nnet3/tdnn_a_ali exp/nnet3/tdnn_a_degs"
    echo ""
    echo "For options, see top of script file.  Standard options:"
    echo "  --config <config-file>                           # config file containing options"
    echo "  --cmd (utils/run.pl|utils/queue.pl <queue opts>) # how to run jobs (probably would be good to add --max-jobs-run 5 or so if using"
    echo "                                                   # GridEngine (to avoid excessive NFS traffic)."
    echo "  --stage <stage|-8>                               # Used to run a partially-completed training process from somewhere in"
    echo "                                                   # the middle."
    echo "  --online-ivector-dir <dir|"">                    # Directory for online-estimated iVectors, used in the"
    echo "                                                   # online-neural-net setup."
    echo "  --nj <nj|200>                                    # number of jobs to submit to the queue."
    echo "  --num-threads <n|1>                              # number of threads per decoding job"
    exit 1;
  fi
  
  data=$1
  lang=$2
  srcdir=$3
  alidir=$4
  dir=$5
  
  
  extra_files=
  [ ! -z $online_ivector_dir ] && \
    extra_files="$extra_files $online_ivector_dir/ivector_period $online_ivector_dir/ivector_online.scp"
  
  # Check some files.
  for f in $data/feats.scp $lang/L.fst $lang/phones/silence.csl $srcdir/${iter}.mdl $srcdir/tree \
        $srcdir/cmvn_opts $alidir/ali.1.gz $alidir/num_jobs $extra_files; do
    [ ! -f $f ] && echo "$0: no such file $f" && exit 1;
  done
  
  mkdir -p $dir/log $dir/info || exit 1;
  
  utils/lang/check_phones_compatible.sh $lang/phones.txt $srcdir/phones.txt || exit 1;
  utils/lang/check_phones_compatible.sh $lang/phones.txt $alidir/phones.txt || exit 1;
  cp $lang/phones.txt $dir || exit 1;
  
  
  
  utils/split_data.sh --per-utt $data $nj
  sdata=$data/split${nj}utt
  
  
  ## Set up features.
  echo "$0: feature type is raw"
  
  
  cmvn_opts=$(cat $srcdir/cmvn_opts) || exit 1
  
  feats="ark,s,cs:apply-cmvn $cmvn_opts --utt2spk=ark:$sdata/JOB/utt2spk scp:$sdata/JOB/cmvn.scp scp:$sdata/JOB/feats.scp ark:- |"
  
  cp $srcdir/{splice_opts,cmvn_opts} $dir 2>/dev/null || true
  
  ## set iVector options
  if [ ! -z "$online_ivector_dir" ]; then
    online_ivector_period=$(cat $online_ivector_dir/ivector_period)
    ivector_opts="--online-ivectors=scp:$online_ivector_dir/ivector_online.scp --online-ivector-period=$online_ivector_period"
  fi
  
  ## set frame-subsampling-factor option and copy file
  if [ -f $srcdir/frame_subsampling_factor ]; then
    frame_subsampling_factor=$(cat $srcdir/frame_subsampling_factor) || exit 1
    # e.g. for 'chain' systems
    frame_subsampling_opt="--frame-subsampling-factor=$frame_subsampling_factor"
    cp $srcdir/frame_subsampling_factor $dir
    if [ $frame_subsampling_factor -ne 1 ] && [ "$self_loop_scale" == "0.1" ]; then
      echo "$0: warning: frame_subsampling_factor is not 1 (so likely a chain system),"
      echo "...  but self-loop-scale is 0.1.  Make sure this is not a mistake."
      sleep 1
    fi
  else
    frame_subsampling_factor=1
  fi
  
  if [ "$self_loop_scale" == "1.0" ] && [ "$acwt" == 0.1 ]; then
    echo "$0: warning: you set --self-loop-scale=1.0 (so likely a chain system)",
    echo " ... but the acwt is still 0.1 (you probably want --acwt 1.0)"
    sleep 1
  fi
  
  ## Make the decoding graph.
  if [ $stage -le 0 ]; then
    new_lang="$dir/"$(basename "$lang")
    rm -r $new_lang 2>/dev/null
    cp -rH $lang $dir
    echo "$0: Making unigram grammar FST in $new_lang"
    oov=$(cat data/lang/oov.txt)
    cat $data/text | utils/sym2int.pl --map-oov $oov -f 2- $lang/words.txt | \
     awk '{for(n=2;n<=NF;n++){ printf("%s ", $n); } printf("
  "); }' | \
      utils/make_unigram_grammar.pl | fstcompile | fstarcsort --sort_type=ilabel > $new_lang/G.fst \
      || exit 1;
  
    utils/mkgraph.sh --self-loop-scale $self_loop_scale $new_lang $srcdir $dir/dengraph || exit 1;
  fi
  
  # copy alignments into ark,scp format which allows us to use different num-jobs
  # from the alignment, and is also convenient for getting priors.
  if [ $stage -le 1 ]; then
    echo "$0: Copying input alignments"
    nj_ali=$(cat $alidir/num_jobs)
    alis=$(for n in $(seq $nj_ali); do echo -n "$alidir/ali.$n.gz "; done)
    $cmd $dir/log/copy_alignments.log \
       copy-int-vector "ark:gunzip -c $alis|" \
       ark,scp:$dir/ali.ark,$dir/ali.scp || exit 1;
  fi
  
  [ -f $dir/ali.scp ] || { echo "$0: expected $dir/ali.scp to exist"; exit 1; }
  
  if [ $stage -le 2 ]; then
    echo "$0: working out number of frames of training data"
    num_frames=$(steps/nnet2/get_num_frames.sh $data)
    echo $num_frames > $dir/info/num_frames
    echo "$0: working out feature dim"
    feats_one="$(echo $feats | sed s:JOB:1:g)"
    if feat_dim=$(feat-to-dim "$feats_one" - 2>/dev/null); then
      echo $feat_dim > $dir/info/feat_dim
    else # run without stderr redirection to show the error.
      feat-to-dim "$feats_one" -; exit 1
    fi
  else
    num_frames=$(cat $dir/info/num_frames)
  fi
  if ! [ "$num_frames" -gt 0 ]; then
    echo "$0: bad num-frames=$num_frames"; exit 1
  fi
  
  # copy the model to the degs directory.
  cp $srcdir/${iter}.mdl $dir/final.mdl || exit 1
  
  # Create some info in $dir/info
  
  # Work out total number of archives. Add one on the assumption the
  # num-frames won't divide exactly, and we want to round up.
  num_archives=$[num_frames/frames_per_iter+1]
  
  echo $num_archives >$dir/info/num_archives
  echo $frame_subsampling_factor >$dir/info/frame_subsampling_factor
  cp $lang/phones/silence.csl $dir/info/
  
  # the first field in frames_per_eg (which is a comma-separated list of numbers)
  # is the 'principal' frames-per-eg, and for purposes of working out the number
  # of archives we assume that this will be the average number of frames per eg.
  frames_per_eg_principal=$(echo $frames_per_eg | cut -d, -f1)
  
  
  # read 'mof' as max_open_filehandles.
  # When splitting up the scp files, we don't want to have to hold too many
  # files open at once.  If the number of archives we have to write exceeds
  # 256 (or less if unlimit -n is smaller), we split in two stages.
  mof=$(ulimit -n) || exit 1
  # the next step helps work around inconsistency between different machines on a
  # cluster.  It's unlikely that the allowed number of open filehandles would ever
  # be less than 256.
  if [ $mof -gt 256 ]; then mof=256; fi
  # allocate mof minus 3 for the max allowed outputs, because of
  # stdin,stderr,stdout.  this will normally come to 253.  We'll do a two-stage
  # splitting if the needed number of scp files is larger than this.
  num_groups=$[(num_archives+(mof-3)-1)/(mof-3)]
  group_size=$[(num_archives+num_groups-1)/num_groups]
  if [ $num_groups -gt 1 ]; then
    new_num_archives=$[group_size*num_groups]
    [ $new_num_archives -ne $num_archives ] && \
      echo "$0: rounding up num-archives from $num_archives to $new_num_archives for easier splitting"
    num_archives=$new_num_archives
    echo $new_num_archives >$dir/info/num_archives
  fi
  
  
  if [ -e $dir/storage ]; then
    # Make soft links to storage directories, if distributing this way..  See
    # utils/create_split_dir.pl.
    echo "$0: creating data links"
    utils/create_data_link.pl $(for x in $(seq $num_archives); do echo $dir/degs.$x.ark; done)
    utils/create_data_link.pl $(for x in $(seq $num_archives); do echo $dir/degs.$x.scp; done)
    utils/create_data_link.pl $(for y in $(seq $nj); do echo $dir/degs_orig.$y.ark; done)
    utils/create_data_link.pl $(for y in $(seq $nj); do echo $dir/degs_orig.$y.scp; done)
    utils/create_data_link.pl $(for y in $(seq $nj); do echo $dir/degs_orig_filtered.$y.scp; done)
  fi
  
  
  extra_context_opts="--extra-left-context=$extra_left_context --extra-right-context=$extra_right_context --extra-left-context-initial=$extra_left_context_initial --extra-right-context-final=$extra_right_context_final"
  
  # work out absolute context opts, --left-context and so on [need model context]
  model_left_context=$(nnet3-am-info $srcdir/${iter}.mdl | grep "^left-context:" | awk '{print $2}')
  model_right_context=$(nnet3-am-info $srcdir/${iter}.mdl | grep "^right-context:" | awk '{print $2}')
  left_context=$[model_left_context+extra_left_context+frame_subsampling_factor/2]
  right_context=$[model_right_context+extra_right_context+frame_subsampling_factor/2]
  context_opts="--left-context=$left_context --right-context=$right_context"
  if [ $extra_left_context_initial -ge 0 ]; then
    left_context_initial=$[model_left_context+extra_left_context_initial+frame_subsampling_factor/2]
    context_opts="$context_opts --left-context-initial=$left_context_initial"
  fi
  if [ $extra_right_context_final -ge 0 ]; then
    right_context_final=$[model_right_context+extra_right_context_final+frame_subsampling_factor/2]
    context_opts="$context_opts --right-context-final=$right_context_final"
  fi
  
  ##
  if [ $num_threads -eq 1 ]; then
    if $looped; then
      decoder="nnet3-latgen-faster-looped"
      [ $extra_left_context_initial -ge 0 ] && \
        decoder="$decoder --extra-left-context-initial=$extra_left_context_initial"
    else
      decoder="nnet3-latgen-faster $extra_context_opts"
    fi
    threads_cmd_opt=
  else
    $looped && { echo "$0: --num-threads must be one if you use looped decoding"; exit 1; }
    threads_cmd_opt="--num-threads $num_threads"
    decoder="nnet3-latgen-faster-parallel --num-threads=$num_threads $extra_context_opts"
    true
  fi
  
  # set the command to determinize lattices, if specified.
  if $determinize_before_split; then
    lattice_determinize_cmd="lattice-determinize-non-compact --acoustic-scale=$acwt --max-mem=$max_mem --minimize=true --prune=true --beam=$lattice_beam ark:- ark:-"
  else
    lattice_determinize_cmd="cat"
  fi
  
  if [ $stage -le 3 ]; then
    echo "$0: decoding and dumping egs"
    $cmd $threads_cmd_opt JOB=1:$nj $dir/log/decode_and_get_egs.JOB.log \
       $decoder \
       $ivector_opts $frame_subsampling_opt \
      --frames-per-chunk=$frames_per_chunk_decoding \
      --determinize-lattice=false \
      --max-active=$max_active --min-active=$min_active --beam=$beam \
      --lattice-beam=$lattice_beam --acoustic-scale=$acwt --allow-partial=false \
      --word-symbol-table=$lang/words.txt $dir/final.mdl  \
      $dir/dengraph/HCLG.fst "$feats" ark:- \| \
      $lattice_determinize_cmd  \| \
      nnet3-discriminative-get-egs --acoustic-scale=$acwt --compress=$compress \
        $frame_subsampling_opt --num-frames=$frames_per_eg \
        --num-frames-overlap=$frames_overlap_per_eg \
        $ivector_opts $context_opts \
        $dir/final.mdl "$feats"  "ark,s,cs:-" \
        "scp:utils/filter_scp.pl $sdata/JOB/utt2spk $dir/ali.scp |" \
        ark,scp:$dir/degs_orig.JOB.ark,$dir/degs_orig.JOB.scp || exit 1
  fi
  
  
  if [ $stage -le 4 ]; then
    echo "$0: getting validation utterances."
  
    ## Get list of validation utterances.
    awk '{print $1}' $data/utt2spk | utils/shuffle_list.pl | head -$num_utts_subset \
     > $dir/valid_uttlist || exit 1;
  
    if [ -f $data/utt2uniq ]; then  # this matters if you use data augmentation.
      echo "File $data/utt2uniq exists, so augmenting valid_uttlist to"
      echo "include all perturbed versions of the same 'real' utterances."
      mv $dir/valid_uttlist $dir/valid_uttlist.tmp
      utils/utt2spk_to_spk2utt.pl $data/utt2uniq > $dir/uniq2utt
      cat $dir/valid_uttlist.tmp | utils/apply_map.pl $data/utt2uniq | \
        sort | uniq | utils/apply_map.pl $dir/uniq2utt | \
        awk '{for(n=1;n<=NF;n++) print $n;}' | sort  > $dir/valid_uttlist
      rm $dir/uniq2utt $dir/valid_uttlist.tmp
    fi
  
    # the following awk statement turns 'foo123' into something like
    # '^foo123-[0-9]\+ ' which is a grep expression that matches the lines in the
    # .scp file that correspond to an utterance in valid_uttlist.
    cat $dir/valid_uttlist | awk '{printf("^%s-[0-9]\\+ 
  ", $1);}' \
       >$dir/valid_uttlist.regexps || exit 1
  
    # remove the validation utterances from deg_orig.*.scp to produce
    # degs_orig_filtered.*.scp.
    # note: the '||' true is in case the grep returns nonzero status for
    # some splits, because they were all validation utterances.
    $cmd JOB=1:$nj $dir/log/filter_and_shuffle.JOB.log \
       grep -v -f $dir/valid_uttlist.regexps $dir/degs_orig.JOB.scp '>' \
       $dir/degs_orig_filtered.JOB.scp '||' true || exit 1
  
    # extract just the validation utterances from deg_orig.*.scp to produce
    # degs_valid.*.scp.
    $cmd JOB=1:$nj $dir/log/extract_validation_egs.JOB.log \
      grep -f $dir/valid_uttlist.regexps $dir/degs_orig.JOB.scp '>' \
      $dir/degs_valid.JOB.scp '||' true || exit 1
  
    for j in $(seq $nj); do
      cat $dir/degs_valid.$j.scp; rm $dir/degs_valid.$j.scp;
    done | utils/shuffle_list.pl | head -n$num_utts_subset >$dir/valid_diagnostic.scp || exit 1
  
    [ -s $dir/valid_diagnostic.scp ] || { echo "$0: error getting validation egs"; exit 1; }
  fi
  
  
  
  # function/pseudo-command to randomly shuffle input lines using a small buffer size
  function shuffle {
      perl -e ' use List::Util qw(shuffle); srand(0);
         $bufsz=1000; @A = (); while(<STDIN>) { push @A, $_; if (@A == $bufsz) {
         $n=int(rand()*$bufsz); print $A[$n]; $A[$n] = $A[$bufsz-1]; pop @A; }}
         @A = shuffle(@A); print @A; '
  }
  # funtion/pseudo-command to put input lines round robin to command line args.
  function round_robin {
    perl -e '@F=(); foreach $a (@ARGV) { my $f; open($f, ">$a") || die "opening file $a"; push @F, $f; }
           $N=@F; $N>0||die "No output files"; $n=0;
           while (<STDIN>) { $fh=$F[$n%$N]; $n++; print $fh $_ || die "error printing"; } ' $*
  }
  
  
  if [ $stage -le 5 ]; then
    echo "$0: rearranging scp files"
  
    if [ $num_groups -eq 1 ]; then
      # output directly to the archive files.
      outputs=$(for n in $(seq $num_archives); do echo $dir/degs.$n.scp; done)
    else
      # output to intermediate 'group' files.
      outputs=$(for g in $(seq $num_groups); do echo $dir/degs_group.$g.scp; done)
    fi
  
    # We can't use UNIX's split command because of compatibility issues (BSD
    # version very different from GNU version), so we use 'round_robin' which is
    # a bash function that calls an inline perl script.
    for j in $(seq $nj); do cat $dir/degs_orig_filtered.$j.scp; done | \
      shuffle | round_robin $outputs || exit 1
  
    if [ $num_groups -gt 1 ]; then
      for g in $(seq $num_groups); do
        first=$[1+group_size*(g-1)]
        last=$[group_size*g]
        outputs=$(for n in $(seq $first $last); do echo $dir/degs.$n.scp; done)
        cat $dir/degs_group.$g.scp | shuffle | round_robin $outputs
      done
    fi
  fi
  
  if [ $stage -le 6 ]; then
    echo "$0: getting train-subset scp"
    # get degs_train_subset.scp by taking the top and tail of the degs files [quicker
    # than cat'ing all the files, random shuffling and head]
  
    nl=$[$num_egs_subset/$num_archives + 1]
  
    # use utils/shuffle_list.pl because it provides a complete shuffle (ok since
    # the amount of data is small).  note: shuf is not available on mac by
    # default.
    for n in $(seq $num_archives); do
      head -n$nl $dir/degs.$n.scp;  tail -n$nl $dir/degs.$n.scp
    done  | utils/shuffle_list.pl | head -n$num_utts_subset >$dir/train_diagnostic.scp
    [ -s $dir/train_diagnostic.scp ] || { echo "$0: error getting train_diagnostic.scp"; exit 1; }
  fi
  
  if [ $stage -le 7 ]; then
    echo "$0: creating final archives"
    $cmd --max-jobs-run "$max_copy_jobs" \
       JOB=1:$num_archives $dir/log/copy_archives.JOB.log \
       nnet3-discriminative-copy-egs scp:$dir/degs.JOB.scp ark:$dir/degs.JOB.ark || exit 1
  
    run.pl $dir/log/copy_train_subset.log \
        nnet3-discriminative-copy-egs scp:$dir/train_diagnostic.scp \
           ark:$dir/train_diagnostic.degs  || exit 1
  
    run.pl $dir/log/copy_valid_subset.log \
        nnet3-discriminative-copy-egs scp:$dir/valid_diagnostic.scp \
           ark:$dir/valid_diagnostic.degs  || exit 1
  fi
  
  if [ $stage -le 10 ] && $cleanup; then
    echo "$0: cleaning up temporary files."
    for j in $(seq $nj); do
      for f in $dir/degs_orig.$j.{ark,scp} $dir/degs_orig_filtered.$j.scp; do
        [ -L $f ] && rm $(utils/make_absolute.sh $f); rm $f
      done
    done
    rm $dir/degs_group.*.scp $dir/valid_diagnostic.scp $dir/train_diagnostic.scp 2>/dev/null
    rm $dir/ali.ark $dir/ali.scp 2>/dev/null
    for n in $(seq $num_archives); do
      for f in $dir/degs.$n.scp; do
        [ -L $f ] && rm $(utils/make_absolute.sh $f); rm $f
      done
    done
  fi
  
  
  echo "$0: Finished decoding and preparing training examples"
  
  exit 0