#!/bin/bash
  # Copyright 2016  Xiaohui Zhang  Apache 2.0.
  # Copyright 2019  SmartAction (kkm)
  
  # This script combines alignment directories, such as exp/tri4a_ali, and
  # validates matching of the utterances and alignments after combining.
  
  # Begin configuration section.
  cmd=run.pl
  nj=4
  combine_lat=true
  combine_ali=true
  tolerance=10
  # End configuration section.
  echo "$0 $@"  # Print the command line for logging.
  
  [[ -f path.sh ]] && . ./path.sh
  . parse_options.sh || exit 1
  
  export LC_ALL=C
  
  if [[ $# -lt 3 ]]; then
    cat >&2 <<EOF
  Usage: $0 [options] <data> <dest-dir> <src-dir1> <src-dir2> ...
   e.g.: $0 --nj 32 data/train exp/tri3_ali_combined exp/tri3_ali_1 exp_tri3_ali_2
  Options:
   --nj <nj>              # number of jobs to split combined archives [4]
   --combine_ali false    # merge ali.*.gz if present [true]
   --combine_lat false    # merge lat.*.gz if present [true]
   --tolerance <int,%>    # maximum percentage of missing alignments or lattices
                          # w.r.t. total utterances in <data> before error is
                          # reported [10]
  
  The script checks that certain important files are present and compatible in all
  source directories (phones.txt, tree); other are copied from the first source
  (cmvn_opts, final.mdl) without much checking.
  
  Both --combine_ali and --combine_lat are true by default, but the script
  proceeds with a warning if directories do not contain either alignments or
  alignment lattices. Check for files ali.1.gz and/or lat.1.gz in the <dest-dir>
  after the script completes if additional programmatic check is required.
  EOF
    exit 1;
  fi
  
  if [[ ! $combine_lat && ! $combine_ali ]]; then
    echo "$0: at least one of --combine_lat and --combine_ali must be true"
    exit 1
  fi
  
  data=$1
  dest=$2
  shift 2
  first_src=$1
  
  do_ali=$combine_ali
  do_lat=$combine_lat
  
  # Check if alignments and/or lattices are present. Since we combine both,
  # whichever present, issue a warning only. Also verify that the target is
  # different from any source; we cannot combine in-place, and a lot of damage
  # could result.
  for src in $@; do
    if [[ "$(cd 2>/dev/null -P -- "$src" && pwd)" = \
          "$(cd 2>/dev/null -P -- "$dest" && pwd)" ]]; then
      echo "$0: error: Source $src is same as target $dest."
      exit 1
    fi
    if $do_ali && [[ ! -f $src/ali.1.gz ]]; then
      echo "$0: warning: Alignments (ali.*.gz) are not present in $src, not" \
           "combining. Consider '--combine_ali false' to suppress this warning."
      do_ali=false
    fi
    if $do_lat && [[ ! -f $src/lat.1.gz ]]; then
      echo "$0: warning: Alignment lattices (lat.*.gz) are not present in $src,"\
        "not combining. Consider '--combine_lat false' to suppress this warning."
      do_lat=false
    fi
  done
  
  if ! $do_ali && ! $do_lat; then
    echo "$0: error: Cannot combine directories."
    exit 1
  fi
  
  # Verify that required files are present in the first directory.
  for f in cmvn_opts final.mdl num_jobs phones.txt tree; do
    if [ ! -f $first_src/$f ]; then
      echo "$0: error: Required source file $first_src/$f is missing."
      exit 1
    fi
  done
  
  # Verify that phones and trees are compatible in all directories, and than
  # num_jobs files are present, too.
  for src in $@; do
    if [[ $src != $first_src ]]; then
      if [[ ! -f $src/num_jobs ]]; then
        echo "$0: error: Required source file $src/num_jobs is missing."
        exit 1
      fi
      if ! cmp -s $first_src/tree $src/tree; then
        echo "$0: error: tree $src/tree is either missing or not the" \
             "same as $first_src/tree."
        exit 1
      fi
      if [[ ! -f $src/phones.txt ]]; then
        echo "$0: error: Required source file $src/phones.txt is missing."
        exit 1
      fi
      utils/lang/check_phones_compatible.sh $first_src/phones.txt \
                                            $src/phones.txt || exit 1
    fi
  done
  
  # All checks passed, ok to prepare directory. Copy model and other files from
  # the first source, as they either checked to be compatible, or we do not care
  # if they are.
  mkdir -p $dest || exit 1
  rm -f $dest/{cmvn_opts,final.mdl,num_jobs,phones.txt,tree}
  $do_ali && rm -f $dest/ali.*.{gz,scp}
  $do_lat && rm -f $dest/lat.*.{gz,scp}
  cp $first_src/{cmvn_opts,final.mdl,phones.txt,tree} $dest/ || exit 1
  cp $first_src/frame_subsampling_factor $dest/ 2>/dev/null  # If present.
  echo $nj > $dest/num_jobs || exit 1
  
  # Make temporary directory, delete on signal, but not on 'exit 1'.
  temp_dir=$(mktemp -d $dest/temp.XXXXXX) || exit 1
  cleanup() { rm -rf "$temp_dir"; }
  trap cleanup HUP INT TERM
  echo "$0: note: Temporary directory $temp_dir will not be deleted in case of" \
       "script failure, so you could examine it for troubleshooting."
  
  
  # This function may be called twice, once to combine alignments and the second
  # time to combine lattices. The two invocations are as follows:
  #   do_combine ali alignments copy-int-vector $@
  #   do_combine lat lattices   lattice-copy $@
  # where 'ali'/'lat' is a prefix to archive name, 'alignments'/'lattices' go into
  # log messages and logfile names, and 'copy-int-vector'/'lattice-copy' is the
  # program used to copy corresponding objects.
  do_combine() {
    local ark=$1 entities=$2 copy_program=$3
    shift 3
  
    echo "$0: Gathering $entities from each source directory."
    # Assign all source gzipped archive names to an exported variable, one each
    # per source directory, so that we can copy archives in a job per source.
    src_id=0
    for src in $@; do
      src_id=$((src_id + 1))
      nj_src=$(cat $src/num_jobs) || exit 1
      # Create and export variable src_arcs_${src_id} for the job runner.
      # Each numbered variable will contain the list of archives, e. g.:
      # src_arcs_1="exp/tri3_ali/ali.1.gz exp/tri3_ali/ali.1.gz ..."
      # ('printf' repeats its format as long as there are more arguments).
      printf -v src_arks_${src_id} "$src/$ark.%d.gz " $(seq $nj_src)
      export src_arks_${src_id}
    done
  
    # Gather archives in parallel jobs.
    $cmd JOB=1:$src_id $dest/log/gather_$entities.JOB.log \
      $copy_program \
        "ark:gunzip -c \${src_arks_JOB} |" \
        "ark,scp:$temp_dir/$ark.JOB.ark,$temp_dir/$ark.JOB.scp" || exit 1
  
    # Merge (presumed already sorted) scp's into a single script.
    sort -m $temp_dir/$ark.*.scp > $temp_dir/$ark.scp || exit 1
  
    echo "$0: Splitting combined $entities into $nj archives on speaker boundary."
    $cmd JOB=1:$nj $dest/log/chop_combined_$entities.JOB.log \
      $copy_program \
        "scp:utils/split_scp.pl --utt2spk=$data/utt2spk --one-based -j $nj JOB $temp_dir/$ark.scp |" \
        "ark:| gzip -c > $dest/$ark.JOB.gz" || exit 1
  
    # Get some interesting stats, and signal an error if error threshold exceeded.
    n_utt=$(wc -l <$data/utt2spk)
    n_ali=$(wc -l <$temp_dir/$ark.scp)
    n_ali_no_utt=$(join -j1 -v2 $data/utt2spk $temp_dir/$ark.scp | wc -l)
    n_utt_no_ali=$(join -j1 -v1 $data/utt2spk $temp_dir/$ark.scp | wc -l)
    n_utt_no_ali_pct=$(perl -e "print int($n_utt_no_ali/$n_utt * 100 + .5);")
    echo "$0: Combined $n_ali $entities for $n_utt utterances." \
         "There were $n_utt_no_ali utterances (${n_utt_no_ali_pct}%) without" \
         "$entities, and $n_ali_no_utt $entities not matching any utterance."
  
    if (( $n_utt_no_ali_pct >= $tolerance )); then
      echo "$0: error: Percentage of utterances missing $entities," \
           "${n_utt_no_ali_pct}%, is at or above error tolerance ${tolerance}%."
      exit 1
    fi
  
    return 0
  }
  
  # Do the actual combining. Do not check returned exit code, as
  # the function always calls 'exit 1' on failure.
  $do_ali && do_combine ali 'alignments' copy-int-vector "$@"
  $do_lat && do_combine lat 'lattices' lattice-copy "$@"
  
  # Delete the temporary directory on success.
  cleanup
  
  what=
  $do_ali && what+='alignments '
  $do_ali && $do_lat && what+='and '
  $do_lat && what+='lattices '
  echo "$0: Stored combined ${what}in $dest"  # No period, interferes with
                                              # copy/paste from tty emulator.
  exit 0