#!/usr/bin/env bash


set -e

. path.sh


if [ $# -ne 6 ]; then
    echo "usage: <model-conf> <uttids> <dataset> <epochs> <njobs> <out-dir>"
    exit 1
fi

modelconf=$1
uttids=$2
dataset=$3
epochs=$4
njobs=$5
outdir=$6

mkdir -p $outdir

# Create the units' HMM.
if [ ! -f $outdir/hmms.mdl ]; then
    beer hmm mkphones -d $dataset $modelconf $outdir/hmms.mdl || exit 1
else
    echo "units' HMM already created. Skipping."
fi


# Create the phone-loop model.
if [ ! -f $outdir/0.mdl ]; then
    beer hmm phonelist $outdir/hmms.mdl | \
        beer hmm mkphoneloopgraph - \
        $outdir/ploop_graph.pkl || exit 1
    beer hmm mkdecodegraph $outdir/ploop_graph.pkl $outdir/hmms.mdl \
        $outdir/decode_graph.pkl || exit 1
    beer hmm mkphoneloop $outdir/decode_graph.pkl $outdir/hmms.mdl \
        $outdir/0.mdl || exit 1

    # Create the optimizer of the training.
    beer hmm optimizer $outdir/0.mdl $outdir/optim_0.mdl || exit 1
else
    echo "Phone Loop model already created. Skipping."
fi

function split_() {
    njobs=$1
    outdir=$2
    uttids=$3
    epoch_num=$4

    mkdir -p $outdir/epoch${epoch_num}/{split,log}

    # randomize the data 
    shuf_uttids=$(mktemp)
    shuf $uttids > $shuf_uttids
    #shuf $uttids --random-source=$uttids > $shuf_uttids # with a seed 

    # split will give njobs files with a numerical name
    # used on submit_gnu_parallel, the number of files
    # generated by split_ is the same than the number
    # of processes that will be run in the computer
    split --numeric-suffixes=1 -a ${#njobs} \
        -n l/$njobs $shuf_uttids $outdir/epoch${epoch_num}/split/

    rm -rf $shuf_utts

}
export -f split_;


function submit_gnu_parallel() {
    njobs=$1
    outdir=$2
    mdl=$3
    dataset=$4
    uttids=$5
    epoch_num=$6
    
    split_ $njobs $outdir $uttids $epoch_num

    ls -AF1 $outdir/epoch${epoch_num}/split/* | \
        parallel --eta -j $njobs --tmpdir $outdir/epoch${epoch_num} --files \
        "cat {} | beer hmm accumulate $outdir/$mdl $dataset $outdir/epoch${epoch_num}/elbo{/.}.pkl"
        
}
export -f submit_gnu_parallel;

# Training.
if [ ! -f $outdir/final.mdl ]; then
    echo "training..."

    # Retrieve the last model.
    mdl=$(find $outdir -name "[0-9]*mdl" -exec basename {} \; | \
        sort -t '.' -k 1 -g | tail -1)
    epoch="${mdl%.*}"

    # ...and the optimizer.
    optim=optim_${epoch}.mdl

    # Accumulate the statistics in parallel.    
    while [ $((++epoch)) -le $epochs ]; do
        echo "epoch - $epoch"

        # Accumulate the statistics in parallel.
        submit_gnu_parallel \
            "$njobs" \
            "$outdir" \
            "$mdl" \
            "$dataset" \
            "$uttids" \
            "$epoch"  || exit 1

        # Update the model' parameters.
        find $outdir/epoch${epoch} -name '*pkl' | \
            beer hmm update $outdir/$mdl $outdir/$optim \
                $outdir/${epoch}.mdl $outdir/optim_${epoch}.pkl || exit 1

        mdl=${epoch}.mdl
        optim=optim_${epoch}.pkl
    done

    cp $outdir/$mdl $outdir/final.mdl

else
    echo "Model already trained. Skipping."
fi

# Generating labels.
if [ ! -f $outdir/trans.txt ]; then
    # Creating the most likely transcription.
    # TODO: can the decoder be paralellized? 
    echo "generating transcription for the $dataset dataset..."
    beer hmm decode --per-frame $outdir/final.mdl \
        $dataset > $outdir/trans.txt || exit 1
else
    echo "transcription already generated. Skipping."
fi

