#!/bin/bash

# Copyright 2012  Johns Hopkins University (Author: Daniel Povey).  Apache 2.0.

# mix up (or down); do 3 iters of model training; realign; then do two more
# iterations of model training.

# Begin configuration section.
cmd=run.pl
beam=10
retry_beam=40
boost_silence=1.0 # Factor by which to boost silence likelihoods in alignment
scale_opts="--transition-scale=1.0 --acoustic-scale=0.1 --self-loop-scale=0.1"
num_iters=5
realign_iters=3 # Space-separated list of iterations to realign on.
stage=0
# End configuration section.

echo "$0 $@"  # Print the command line for logging

[ -f path.sh ] && . ./path.sh;
. parse_options.sh || exit 1;

if [ $# != 5 ]; then
   echo "Usage: steps/mixup.sh <num-gauss> <data-dir> <lang-dir> <old-exp-dir> <exp-dir>"
   echo " e.g.: steps/mixup.sh 20000 data/train_si84 data/lang exp/tri3b exp/tri3b_20k"
   echo "main options (for others, see top of script file)"
   echo "  --cmd (utils/run.pl|utils/queue.pl <queue opts>) # how to run jobs."
   echo "  --config <config-file>                           # config containing options"
   echo "  --stage <stage>                                  # stage to do partial re-run from."
   exit 1;
fi

numgauss=$1
data=$2
lang=$3
srcdir=$4
dir=$5

for f in $data/feats.scp $srcdir/final.mdl $srcdir/final.mat; do
  [ ! -f $f ] && echo "mixup_lda_etc.sh: no such file $f" && exit 1;
done

nj=`cat $srcdir/num_jobs` || exit 1;
sdata=$data/split$nj;

splice_opts=`cat $srcdir/splice_opts 2>/dev/null`

mkdir -p $dir/log
cp $srcdir/splice_opts $dir 2>/dev/null
cp $srcdir/final.mat $dir
echo $nj > $dir/num_jobs
[[ -d $sdata && $data/feats.scp -ot $sdata ]] || split_data.sh $data $nj || exit 1;

cp $srcdir/tree $dir


## Set up features.
if [ -f $srcdir/final.mat ]; then feat_type=lda; else feat_type=delta; fi
echo "$0: feature type is $feat_type"

case $feat_type in
  delta) sifeats="ark,s,cs:apply-cmvn --norm-vars=false --utt2spk=ark:$sdata/JOB/utt2spk scp:$sdata/JOB/cmvn.scp scp:$sdata/JOB/feats.scp ark:- | add-deltas ark:- ark:- |";;
  lda) sifeats="ark,s,cs:apply-cmvn --norm-vars=false --utt2spk=ark:$sdata/JOB/utt2spk scp:$sdata/JOB/cmvn.scp scp:$sdata/JOB/feats.scp ark:- | splice-feats $splice_opts ark:- ark:- | transform-feats $srcdir/final.mat ark:- ark:- |"
    cp $srcdir/final.mat $dir    
    ;;
  *) echo "Invalid feature type $feat_type" && exit 1;
esac
if [ -f $srcdir/trans.1 ]; then
  echo Using transforms from $srcdir;
  rm $dir/trans.* 2>/dev/null
  ln.pl $srcdir/trans.* $dir  # Link those transforms to current directory.
  feats="$sifeats transform-feats --utt2spk=ark:$sdata/JOB/utt2spk ark,s,cs:$dir/trans.JOB ark:- ark:- |"
else
  feats="$sifeats"
fi
## Done setting up features.

rm $dir/fsts.*.gz 2>/dev/null
ln.pl $srcdir/fsts.*.gz $dir  # Link training-graph FSTs to current directory.

## Mix up old model
if [ $stage -le 0 ]; then
  echo Mixing up old model to $numgauss Gaussians
# Note: this script also works for mixing down.
  $cmd $dir/log/mixup.log \
    gmm-mixup --mix-up=$numgauss --mix-down=$numgauss \
    $srcdir/final.mdl $srcdir/final.occs $dir/1.mdl || exit 1;
fi
## Done.

cur_alidir=$srcdir # dir to find alignments.
[ -z "$realign_iters" ] && ln.pl $srcdir/ali.*.gz $dir; # link alignments, if
 # we won't be generating them.

x=1
while [ $x -le $num_iters ]; do
  echo "$0: iteration $x"
  if echo $realign_iters | grep -w $x >/dev/null; then
    if [ $stage -le $x ]; then
      echo "$0: realigning data"
      mdl="gmm-boost-silence --boost=$boost_silence `cat $lang/phones/optional_silence.csl` $dir/$x.mdl - |"
      $cmd JOB=1:$nj $dir/log/align.$x.JOB.log \
        gmm-align-compiled $scale_opts --beam=10 --retry-beam=40 "$mdl" \
        "ark:gunzip -c $dir/fsts.JOB.gz|" "$feats" \
        "ark:|gzip -c >$dir/ali.JOB.gz" || exit 1;
    fi
    cur_alidir=$dir
  fi
  if [ $stage -le $x ]; then
    echo "$0: accumulating statistics"
    $cmd JOB=1:$nj $dir/log/acc.$x.JOB.log \
      gmm-acc-stats-ali  $dir/$x.mdl "$feats" \
      "ark,s,cs:gunzip -c $cur_alidir/ali.JOB.gz|" $dir/$x.JOB.acc || exit 1;
    echo "$0: re-estimating model"
    [ "`ls $dir/$x.*.acc | wc -w`" -ne $nj ] && echo "$0: wrong #accs" && exit 1;
    $cmd $dir/log/update.$x.log \
      gmm-est --write-occs=$dir/$[$x+1].occs $dir/$x.mdl \
      "gmm-sum-accs - $dir/$x.*.acc |" $dir/$[$x+1].mdl || exit 1;
    rm $dir/$x.mdl $dir/$x.*.acc
    rm $dir/$x.occs  2>/dev/null
  fi
  x=$[$x+1]
done

rm $dir/final.mdl $dir/final.occs 2>/dev/null
ln -s $x.mdl $dir/final.mdl
ln -s $x.occs $dir/final.occs

if [ -f $dir/trans.1 ]; then 
  echo "$0: accumulating stats for alignment model."
  $cmd JOB=1:$nj $dir/log/acc_alimdl.JOB.log \
    ali-to-post "ark:gunzip -c $dir/ali.JOB.gz|" ark:-  \| \
    gmm-acc-stats-twofeats $dir/$x.mdl "$feats" "$sifeats" \
    ark,s,cs:- $dir/$x.JOB.acc || exit 1;
  [ "`ls $dir/$x.*.acc | wc -w`" -ne $nj ] && echo "$0: wrong #accs" && exit 1;  
  echo "$0: Re-estimating alignment model."
  $cmd $dir/log/est_alimdl.log \
    gmm-est --write-occs=$dir/final.occs --remove-low-count-gaussians=false $dir/$x.mdl \
    "gmm-sum-accs - $dir/$x.*.acc|" $dir/$x.alimdl  || exit 1;
  rm $dir/$x.*.acc
  rm $dir/final.alimdl 2>/dev/null
  ln -s $x.alimdl $dir/final.alimdl 
fi

utils/summarize_warnings.pl $dir/log

echo Done