#!/bin/bash
  
  # Copyright       2014  Daniel Povey
  # Apache 2.0
  
  #
  # This training script computes some things you will need in order to
  # extract VTLN-warped features.  It takes as input the data directory
  # and an already-trained diagonal-covariance UBM.  Note: although this
  # script is in the lid/ directory, because it is intended to be
  # used in language identification, but it uses features of the
  # same type as those used in the speaker-id scripts (see ../sid/),
  # i.e. double-delta features, rather than the "shifted delta cepstra"
  # features commonly used in language id.
  #
  # This script works with either mfcc or plp features; for plp features, you will
  # need to set the --base-feat-type option.  Regardless, you will need to set the
  # --mfcc-config or --plp-config option if your feature-extraction config is not
  # called conf/${base_feat_type}.conf.  The output of this script will be in
  # $dir/final.lvtln and $dir/final.dubm and $dir/final.ali_dubm; the directory
  # can be passed to ./get_vtln_warps.sh to get VTLN warps for a data directory,
  # or (for data passed to this script) you can use the warping factors this
  # script outputs in $dir/final.warp
  #
  
  # Begin configuration.
  stage=-4 #  This allows restarting after partway, when something when wrong.
  config=
  cmd=run.pl
  num_iters=15    # Number of iterations of training.
  num_utt_lvtln_init=400; # number of utterances (subset) to initialize
                          # LVTLN transform.  Not too critical.
  min_warp=0.85
  max_warp=1.25
  warp_step=0.01
  base_feat_type=mfcc # or could be PLP.
  mfcc_config=conf/mfcc.conf  # default, can be overridden.
  plp_config=conf/plp.conf  # default, can be overridden.
  logdet_scale=0.0
  subsample=5 # We use every 5th frame by default; this is more
              # CPU-efficient.
  min_gaussian_weight=0.0001 # does not matter; inherited from diag-ubm training script.
  nj=4
  cleanup=true
  num_gselect=15
  # End configuration.
  
  echo "$0 $@"  # Print the command line for logging
  
  [ -f path.sh ] && . ./path.sh;
  . parse_options.sh || exit 1;
  
  num_classes=$(perl -e "print int(1.5 + ($max_warp - $min_warp) / $warp_step);") || exit 1;
  default_class=$(perl -e "print int(0.5 + (1.0 - $min_warp) / $warp_step);") || exit 1;
  
  if [ $# != 3 ]; then
     echo "Usage: $0 <data-dir> <diag-ubm-dir> <exp-dir>"
     echo "e.g.: $0 data/train_vtln exp/diag_ubm_vtln exp/vtln"
     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 "  --nj <num-jobs>                                  # number of jobs to use (default 4)"
     echo "  --config <config-file>                           # config containing options"
     echo "  --stage <stage>                                  # stage to do partial re-run from."
     echo "  --num-iters <num-iters>                          # number of iterations of training"
     echo "  --base-feat-type <feat-type>                     # mfcc or plp, mfcc is default"
     echo "  --mfcc-config <config>                           # config for MFCC extraction, default is"
     echo "                                                   # conf/mfcc.conf"
     echo "  --plp-config <config>                            # config for PLP extraction, default is"
     echo "                                                   # conf/plp.conf"
     exit 1;
  fi
  
  data=$1
  ubmdir=$2
  dir=$3
  
  for f in $data/feats.scp $ubmdir/final.dubm; do
    [ ! -f $f ] && echo "train_deltas.sh: no such file $f" && exit 1;
  done
  
  
  mkdir -p $dir/log
  echo $nj > $dir/num_jobs
  
  sdata=$data/split$nj;
  split_data.sh $data $nj || exit 1;
  
  cmvn_sliding_opts="--norm-vars=false --center=true --cmn-window=300"
  # don't change $cmvn_sliding_opts, it should probably match the
  # options used in ../sid/train_diag_ubm.sh.
  sifeats="ark,s,cs:add-deltas scp:$sdata/JOB/feats.scp ark:- | apply-cmvn-sliding $cmvn_sliding_opts ark:- ark:- | select-voiced-frames ark:- scp,s,cs:$sdata/JOB/vad.scp ark:- | subsample-feats --n=$subsample ark:- ark:- |"
  
  
  # for the subsets of features that we use to estimate the linear transforms, we
  # don't bother with CMN.  This will give us wrong offsets on the transforms,
  # but it won't matter because we will allow an arbitrary bias term when we apply
  # these transforms.
  
  # you need to define CLASS when invoking $cmd on featsub_warped.
  featsub_warped="ark:add-deltas ark:$dir/feats.CLASS.ark ark:- | select-voiced-frames ark:- scp,s,cs:$data/vad.scp ark:- | subsample-feats --n=$subsample ark:- ark:- |"
  featsub_unwarped="ark:add-deltas ark:$dir/feats.$default_class.ark ark:- | select-voiced-frames ark:- scp,s,cs:$data/vad.scp ark:- | subsample-feats --n=$subsample ark:- ark:- |"
  
  
  if [ -f $data/utt2warp ]; then
    echo "$0: source data directory $data appears to already have VTLN.";
    exit 1;
  fi
  
  # create a small subset of utterances for purposes of initializing the LVTLN transform
  # utils/shuffle_list.pl is deterministic, unlike sort -R.
  cat $data/utt2spk | awk '{print $1}' | utils/shuffle_list.pl | \
    head -n $num_utt_lvtln_init > $dir/utt_subset
  
  if [ $stage -le -4 ]; then
    echo "$0: computing warped subset of features"
    if [ -f $data/segments ]; then
      echo "$0 [info]: segments file exists: using that."
      subset_feats="utils/filter_scp.pl $dir/utt_subset $data/segments | extract-segments scp:$data/wav.scp - ark:- "
    else
      echo "$0 [info]: no segments file exists: using wav.scp directly."
      subset_feats="utils/filter_scp.pl $dir/utt_subset $data/wav.scp | wav-copy scp:- ark:- "
    fi
    rm $dir/.error 2>/dev/null
    for c in $(seq 0 $[$num_classes-1]); do
      this_warp=$(perl -e "print ($min_warp + ($c*$warp_step));")
      config_name=${base_feat_type}_config # e.g. mfcc_config or plp_config
      this_config=$(eval echo \$$config_name)  #  e.g. conf/mfcc.conf or conf/plp.conf by default.
      $cmd $dir/log/compute_warped_feats.$c.log \
        $subset_feats \| compute-${base_feat_type}-feats --verbose=2 \
        --config=$this_config --vtln-warp=$this_warp ark:- ark:- \| \
        copy-feats --compress=true ark:- ark:$dir/feats.$c.ark || touch $dir/.error &
    done
    wait;
    if [ -f $dir/.error ]; then
      echo "$0: Computing warped features failed: check $dir/log/compute_warped_feats.*.log"
      exit 1;
    fi
  fi
  
  if ! utils/filter_scp.pl $dir/utt_subset $data/feats.scp | \
    compare-feats --threshold=0.98 scp:-  ark:$dir/feats.$default_class.ark >&/dev/null; then
    echo "$0: features stored on disk differ from those computed with no warping."
    echo "    Possibly your feature type is wrong (--base-feat-type option)"
    exit 1;
  fi
    
  if [ -f $data/segments ]; then
    subset_utts="ark:extract-segments scp:$sdata/JOB/wav.scp $sdata/JOB/segments ark:- |"
  else
    echo "$0 [info]: no segments file exists: using wav.scp directly."
    subset_utts="ark:wav-copy scp:$sdata/JOB/wav.scp ark:- |"
  fi
  
  if [ $stage -le -3 ]; then
    echo "$0: initializing base LVTLN transforms in $dir/0.lvtln (ignore warnings below)"
    dim=$(feat-to-dim "$featsub_unwarped" - ) || exit 1;
  
    $cmd $dir/log/init_lvtln.log \
      gmm-init-lvtln --dim=$dim --num-classes=$num_classes --default-class=$default_class \
        $dir/0.lvtln || exit 1;
  
    for c in $(seq 0 $[$num_classes-1]); do
      this_warp=$(perl -e "print ($min_warp + ($c*$warp_step));")
      orig_feats=ark:$dir/feats.$default_class.ark
      warped_feats=ark:$dir/feats.$c.ark
      logfile=$dir/log/train_special.$c.log
      this_featsub_warped="$(echo $featsub_warped | sed s/CLASS/$c/)"
      if ! gmm-train-lvtln-special --warp=$this_warp --normalize-var=true \
        $c $dir/0.lvtln $dir/0.lvtln \
        "$featsub_unwarped" "$this_featsub_warped" 2>$logfile; then
        echo "$0: Error training LVTLN transform, see $logfile";
        exit 1;
      fi
    done  
  fi
  
  cp $ubmdir/final.dubm $dir/0.dubm
  
  if [ $stage -le -2 ]; then
    echo "$0: computing Gaussian selection info."
  
    $cmd JOB=1:$nj $dir/log/gselect.JOB.log \
      gmm-gselect --n=$num_gselect $ubmdir/final.dubm "$sifeats" \
        "ark:|gzip -c >$dir/gselect.JOB.gz" || exit 1;
  fi
    
  
  if [ $stage -le -1 ]; then
    echo "$0: computing initial LVTLN transforms"  # do this per-utt.
  
    $cmd JOB=1:$nj $dir/log/lvtln.0.JOB.log \
      gmm-global-gselect-to-post $dir/0.dubm "$sifeats" \
        "ark,s,cs:gunzip -c $dir/gselect.JOB.gz|" ark:- \| \
      gmm-global-est-lvtln-trans --logdet-scale=$logdet_scale --verbose=1 \
        $dir/0.dubm $dir/0.lvtln "$sifeats" ark,s,cs:- ark:$dir/trans.0.JOB ark,t:$dir/warp.0.JOB || exit 1
    
    # consolidate the warps into one file.
    for j in $(seq $nj); do cat $dir/warp.0.$j; done > $dir/warp.0
    rm $dir/warp.0.*
  fi
  
  
  x=0
  while [ $x -lt $num_iters ]; do
    feats="$sifeats transform-feats ark:$dir/trans.$x.JOB ark:- ark:- |"
  
    # First update the model.
    if [ $stage -le $x ]; then
      echo "$0: Updating model on pass $x"
      # Accumulate stats.
      $cmd JOB=1:$nj $dir/log/acc.$x.JOB.log \
        gmm-global-acc-stats "--gselect=ark,s,cs:gunzip -c $dir/gselect.JOB.gz|" \
         $dir/$x.dubm "$feats" $dir/$x.JOB.acc || exit 1;
  
      $cmd $dir/log/update.$x.log \
        gmm-global-est --remove-low-count-gaussians=false --min-gaussian-weight=$min_gaussian_weight \
          $dir/$x.dubm "gmm-global-sum-accs - $dir/$x.*.acc|" \
        $dir/$[$x+1].dubm || exit 1;
      $cleanup && rm $dir/$x.*.acc $dir/$x.dubm
    fi
  
    # Now update the LVTLN transforms (and warps.)
    if [ $stage -le $x ]; then
      echo "$0: re-estimating LVTLN transforms on pass $x"
      $cmd JOB=1:$nj $dir/log/lvtln.$x.JOB.log \
        gmm-global-gselect-to-post $dir/$[$x+1].dubm "$feats" \
          "ark,s,cs:gunzip -c $dir/gselect.JOB.gz|" ark:- \| \
        gmm-global-est-lvtln-trans --logdet-scale=$logdet_scale --verbose=1 \
          $dir/$[$x+1].dubm $dir/0.lvtln "$sifeats" ark,s,cs:- \
          ark:$dir/trans.$[$x+1].JOB ark,t:$dir/warp.$[$x+1].JOB || exit 1
  
      # consolidate the warps into one file.
      for j in $(seq $nj); do cat $dir/warp.$[$x+1].$j; done > $dir/warp.$[$x+1]
      rm $dir/warp.$[$x+1].*
      $cleanup && rm $dir/trans.$x.*
    fi
    x=$[$x+1]
  done
  
  feats="$sifeats transform-feats ark:$dir/trans.$x.JOB ark:- ark:- |"
  
  if [ $stage -le $x ]; then
    # Accumulate stats for "alignment model"-- this model is computed with the
    # speaker-independent features, but matches Gaussian-for-Gaussian with the
    # final speaker-adapted model.
    $cmd JOB=1:$nj $dir/log/acc_alimdl.JOB.log \
      gmm-global-acc-stats-twofeats "--gselect=ark,s,cs:gunzip -c $dir/gselect.JOB.gz|" \
        $dir/$x.dubm "$feats" "$sifeats" $dir/$x.JOB.acc || exit 1
    [ `ls $dir/$x.*.acc | wc -w` -ne "$nj" ] && echo "$0: Wrong #accs" && exit 1;
    # Update model.
    $cmd $dir/log/est_alimdl.log \
      gmm-global-est --min-gaussian-weight=$min_gaussian_weight \
        --remove-low-count-gaussians=false $dir/$x.dubm \
       "gmm-global-sum-accs - $dir/$x.*.acc|" $dir/$x.ali_dubm  || exit 1;
    $cleanup && rm $dir/$x.*.acc
  fi
  
  if true; then # Diagnostics
    ln -sf warp.$x $dir/final.warp
    if [ -f $data/spk2gender ]; then 
      # To make it easier to eyeball the male and female speakers' warps
      # separately, separate them out.
      for g in m f; do # means: for gender in male female
        cat $dir/final.warp | \
          utils/filter_scp.pl <(grep -w $g $data/spk2gender | awk '{print $1}') > $dir/final.warp.$g
        echo -n "The last few warp factors for gender $g are: "
        tail -n 10 $dir/final.warp.$g | awk '{printf("%s ", $2);}'; 
        echo
      done
    fi
  fi
  
  ln -sf $x.dubm $dir/final.dubm
  ln -sf $x.ali_dubm $dir/final.ali_dubm
  ln -sf 0.lvtln $dir/final.lvtln
  
  # Summarize warning messages...
  utils/summarize_warnings.pl  $dir/log
  
  echo "$0: Done training LVTLN model in $dir"