#!/usr/bin/env bash
# Copyright 2012  Johns Hopkins University (Author: Daniel Povey).  Apache 2.0.
#           2013  Daniel Povey
#           2014  David Snyder

# This trains a full-covariance UBM from an existing (diagonal or full) UBM,
# for a specified number of iterations.  This is for speaker-id systems
# (we use features specialized for that, and vad).

# Begin configuration section.
nj=16
cmd=run.pl
stage=-2
num_gselect=20 # cutoff for Gaussian-selection that we do once at the start.
subsample=5
num_iters=4
min_gaussian_weight=1.0e-04
remove_low_count_gaussians=true # set this to false if you need #gauss to stay fixed.
cleanup=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 [ $# != 3 ]; then
  echo "Usage: steps/train_full_ubm.sh <data> <old-ubm-dir> <new-ubm-dir>"
  echo "Trains a full-covariance UBM starting from an existing diagonal or"
  echo "full-covariance UBM system."
  echo " e.g.: steps/train_full_ubm.sh --num-iters 8 data/train exp/diag_ubm exp/full_ubm"
  echo "main options (for others, see top of script file)"
  echo "  --config <config-file>                           # config containing options"
  echo "  --cmd (utils/run.pl|utils/queue.pl <queue opts>) # how to run jobs."
  echo "  --nj <n|16>                                      # number of parallel training jobs"
  echo "  --num-gselect <n|20>                             # Number of Gaussians to select using"
  echo "                                                   # initial model (diagonalized if needed)"
  echo "  --subsample <n|5>                                # Take every n'th sample, for efficiency"
  echo "  --num-iters <n|4>                                # Number of iterations of E-M"
  echo "  --min-gaussian-weight <weight|1.0e-05>           # Minimum Gaussian weight (below this,"
  echo "                                                   # we won't update, and will remove Gaussians"
  echo "                                                   # if --remove-low-count-gaussians is true"
  echo "  --remove-low-count-gaussians <true,false|true>   # If true, remove Gaussians below min-weight"
  echo "                                                   # (will only happen on last iteration, in any case"
  echo "  --cleanup <true,false|true>                      # If true, clean up accumulators, intermediate"
  echo "                                                   # models and gselect info"
  exit 1;
fi

data=$1
srcdir=$2
dir=$3

for f in $data/feats.scp $data/vad.scp; do
  [ ! -f $f ] && echo "No such file $f" && exit 1;
done


mkdir -p $dir/log
echo $nj > $dir/num_jobs
sdata=$data/split$nj;
utils/split_data.sh $data $nj || exit 1;

## Set up features.
feats="ark,s,cs:apply-cmvn-sliding --norm-vars=false --center=true --cmn-window=300 scp:$sdata/JOB/feats.scp ark:- | add-deltas-sdc ark:- ark:- | select-voiced-frames ark:- scp,s,cs:$sdata/JOB/vad.scp ark:- | subsample-feats --n=$subsample ark:- ark:- |"

if [ $stage -le -2 ]; then
  if [ -f $srcdir/final.dubm ]; then # diagonal-covariance in $srcdir
    $cmd $dir/log/convert_diag_to_full \
      gmm-global-to-fgmm $srcdir/final.dubm $dir/0.ubm || exit 1;
  elif [ -f $srcdir/final.ubm ]; then
    cp $srcdir/final.ubm $dir/0.ubm || exit 1;
  else
    echo "$0: in $srcdir, expecting final.ubm or final.dubm to exist"
    exit 1;
  fi
fi

if [ $stage -le -1 ]; then
  echo "$0: doing Gaussian selection (using diagonal form of model; selecting $num_gselect indices)"
  $cmd JOB=1:$nj $dir/log/gselect.JOB.log \
    gmm-gselect --n=$num_gselect "fgmm-global-to-gmm $dir/0.ubm - |" "$feats" \
    "ark:|gzip -c >$dir/gselect.JOB.gz" || exit 1;
fi


x=0
while [ $x -lt $num_iters ]; do
  echo "Pass $x"
  if [ $stage -le $x ]; then
    $cmd JOB=1:$nj $dir/log/acc.$x.JOB.log \
      fgmm-global-acc-stats "--gselect=ark,s,cs:gunzip -c $dir/gselect.JOB.gz|" $dir/$x.ubm "$feats" \
      $dir/$x.JOB.acc || exit 1;

    if [ $[$x+1] -eq $num_iters ];then
      lowcount_opt="--remove-low-count-gaussians=$remove_low_count_gaussians" # as specified by user.
    else
    # On non-final iters, we in any case can't remove low-count Gaussians because it would
    # cause the gselect info to become out of date.
      lowcount_opt="--remove-low-count-gaussians=false"
    fi
    $cmd $dir/log/update.$x.log \
    fgmm-global-est $lowcount_opt --min-gaussian-weight=$min_gaussian_weight --verbose=2 $dir/$x.ubm "fgmm-global-sum-accs - $dir/$x.*.acc |" \
      $dir/$[$x+1].ubm || exit 1;
    $cleanup && rm $dir/$x.*.acc $dir/$x.ubm
  fi
  x=$[$x+1]
done

$cleanup && rm $dir/gselect.*.gz

rm $dir/final.ubm 2>/dev/null
mv $dir/$x.ubm $dir/final.ubm || exit 1;
