#!/bin/bash

# currently assumes: phone_list, substitutions,  train.scp, aligned.0.mlf
#

root=$1

if [ ! -d "$root" ] || [ $BASH_ARGC -lt 2 ] ; then
    echo
    echo
    echo "Usage: $0  root_directory  N  [init_hmm]"
    echo "                                          V1.0, 2005/12/13"
    echo
    echo
    echo "Parallel version (using Sun grid engine) of do_alignment,"
    echo "which is the CSTR recipie to do phone alignment."
    echo
    echo "Goes to root_directory and expects to find there the"
    echo "files phone_list, substitutions, train.scp, aligned.0.mlf"
    echo "resources/* and proto/*."
    echo
    echo "N is the number of parallel jobs."
    echo
    echo "If init_hmm is given and contains MMF and vFloors,"
    echo "these master model file and variance floors will"
    echo "be used as initial ones, i.e. copied into hmm0."
    echo
    echo "Executables HCompV, HERest, HHEd, HVite, make_mfcc_list,"
    echo "and Split_stp need to be in your PATH."
    exit 1
fi

N=$2  # number of parallel jobs

for exe in qsub HCompV HERest HHEd HVite Split_stp; do

    which $exe   &> /dev/null
    if [ $? -gt 0 ] ;then
       echo "`which $exe`, Aborted!"
      exit 1
    fi
done


## The general idea how to parallelize the original script "do_alignment"
## is to divide it into segments, create a temporary shell script for 
## each segment (in $root/scripts), even if it is just one HTK command 
## and submit it then to SGE using qsub directly (-o and -e do not work
## within submit scripts).  The first section (general setup) is a single 
## thread, parameter re-estimations and alignments are $N threads.
##
## The dependencies are taken care of by "-sync y" which causes qsub 
## to wait for all array jobs to be finished.

echo "START_TIME is `date`"

echo "CHANGING DIRECTORY"
cd $root

##  SETTING UP OPTIONS FOR QSUB


## Stderr and stdout also goes to special dirs:
script_dir=$root/scripts; if [ ! -d $script_dir ] ; then mkdir $script_dir; fi
stdout_dir=$root/stdout;  if [ ! -d $stdout_dir ] ; then mkdir $stdout_dir; fi
# stderr is merged into stdout through the "-j y" switch
#job_root_name="do_alignment" too long, total length limited to 10
job_root_name="alignmt"
section=1
job_name=$job_root_name.$section
script=$script_dir/$job_name
out=$stdout_dir/$job_name

std_qsub_args="-S /bin/bash -v PATH=$PATH -cwd -sync y"
qsub_args="$std_qsub_args -N $job_name -j y -o $out $script"



## SECTION 1:  GENERAL SETUP
##
## in this section, every backtick and dollar is escaped except 
## the argv-related: $root, $BASH_ARGC and $2

cat << EOF > $script; chmod u+x $script
#!/bin/bash

echo "CHECKING FOR FILES"
if [ ! -r "phone_list" ] ; then
    echo "phone_list not found"
    exit 1
fi
if [ ! -r "phone_substitutions" ] ; then
    echo "phone_substitutions not found"
    exit 1
fi
if [ ! -r "train.scp" ] ; then
    echo "train.scp not found"
    exit 1
fi
if [ ! -r "aligned.0.mlf" ] ; then
    echo "aligned.0.mlf not found"
    echo "(I am no longer using ../utts.mlf)"
    exit 1
fi
if [ ! -d "proto" ] ; then
    echo "proto not found"
    exit 1
fi
if [ ! -d "resources" ] ; then
    echo "resources not found"
    exit 1
fi

# create dict and dict2
echo "CREATING DICT"
awk '{print \$1 " " \$1}' phone_list > dict
echo "CREATING DICT2"
cat dict phone_substitutions > dict2

# Generate Master Model file
echo "CREATING MASTER MODEL FILE"
mkdir -p hmm0

if [ $BASH_ARGC -eq 3 ] ; then
    echo " COPYING VARIENCE FLOOR"
    if [ ! -r \$3/vFloors ] ; then
        echo "Copying \$3/vFloors failed, Aborted!" ; exit 1
    fi
    cp $3/vFloors hmm0

    echo " COPYING MODELS"
    if [ ! -r \$3/MMF ] ; then
        echo "Copying \$3/MMF failed, Aborted!" ; exit 1
    fi
    cp $3/MMF hmm0
else

    echo " SETTING VARIENCE FLOOR"
    HCompV -C config -f 0.01 -m -S train.scp -M hmm0 proto/5states
    if [ \$? -gt 0 ] ; then echo "Set varience floor failed, Aborted!" ; exit 1 ; fi
    
    # create models
    echo " GENERATING MODELS"
    for m in \`cat phone_list\` ; do
        if [ "\$m" != "sp" ] ; then
    	grep -v "~h" hmm0/5states > hmm0/\$m
        else 
    	cp proto/3states hmm0/\$m
        fi
    done
    
        HHEd -d hmm0 -w hmm0/MMF resources/tie_silence.hed phone_list
        if [ \$? -gt 0 ] ; then echo "Build master model file failed, Aborted!" ; exit 1 ; fi

fi  

EOF

echo "qsub $qsub_args" | tee $out
qsub $qsub_args
if [ $? -gt 0 ] ; then echo "Initial setup failed, aborted!"; exit 1; fi

###
### INITIAL TRAINING
###

# Quoting the HTK book:
#
#   The second way of speeding-up the operation of HEREST is to use more than
# one computer in parallel. The way that this is done is to divide the
# training data amongst the available machines and then to run HEREST on each
# machine such that each invocation of HEREST uses the same initial set of
# models but has its own private set of data. By setting the option -p N where
# N is an integer, HEREST will dump the contents of all its accumulators  into
# a file called HERN.acc rather than updating and outputing a new set of
# models. These dumped files are collected together and input to a new
# invocation of HEREST with the option -p 0 set. HEREST then reloads the
# accumulators from all of the dump files and updates the models in the normal
# way.

#   HERest -S trlist1 -I labs -H dir1/hmacs -M dir2 -p 1 hmmlist
#   HERest -S trlist2 -I labs -H dir1/hmacs -M dir2 -p 2 hmmlist
#   ...
#   HERest -H dir1/hmacs -M dir2 -p 0 hmmlist dir2/*.acc


i=0
Split_stp train.scp $N  # train.scp_1_$N, train.scp_3_$N, train.scp_3_$N ...

# Re-estimation
for j in 1 2 3 4 5; do
    echo "RE-ESTIMATING MODEL PARAMETERS (ITERATION $i)"
    mkdir -p hmm$[$i +1]

    section=$[$section + 1]
    job_name=$job_root_name.$section
    script=$script_dir/$job_name
    out=$stdout_dir/$job_name
    qsub_args="$std_qsub_args -N $job_name -j y -o $out.\$TASK_ID -t 1-$N $script"


    cat << EOF > $script; chmod +x $script
#!/bin/bash
    HERest -C config -T 1023 -t 250.0 150.0 1000.0 -H hmm${i}/MMF  \
           -H hmm0/vFloors -I aligned.0.mlf -M hmm$[$i + 1] \
           -S train.scp.\${SGE_TASK_ID}_$N -p \$SGE_TASK_ID phone_list
EOF

    echo "qsub $qsub_args" | tee $out
    qsub $qsub_args
    if [ $? -gt 0 ] ; then echo "Re-estimation $i failed, Aborted!" ; exit 1 ; fi

    job_name=$job_root_name.$section.collect
    script=$script_dir/$job_name
    out=$stdout_dir/$job_name
    qsub_args="$std_qsub_args -N $job_name -j y -o $out $script"

    cat << EOF > $script; chmod +x $script
#!/bin/bash

    HERest -H hmm${i}/MMF -H hmm0/vFloors -M hmm$[$i + 1] -p 0 \
           phone_list hmm$[$i + 1]/*.acc
EOF

    echo "qsub $qsub_args" | tee $out
    qsub $qsub_args
    if [ $? -gt 0 ] ; then echo "Collecting re-estimation $i failed, Aborted!" ; exit 1 ; fi

    i=$[$i + 1]
done


# First alignment to correct labelling:
section=$[$section + 1]
job_name=$job_root_name.$section
script=$script_dir/$job_name
out=$stdout_dir/$job_name
qsub_args="$std_qsub_args -N $job_name -j y -o $out.\$TASK_ID -t 1-$N $script"

echo "FIRST ALIGNMENT AND VOWEL REDUCTION"

cat << EOF > $script; chmod u+x $script
#!/bin/bash

HVite -l \* -C config  -a -m -I aligned.0.mlf -H hmm${i}/MMF \
      -i aligned.1.\${SGE_TASK_ID}_$N.mlf -m \
      -S train.scp.\${SGE_TASK_ID}_$N -y lab dict2 phone_list
EOF

echo "qsub $qsub_args" | tee $out
qsub $qsub_args
if [ $? -gt 0 ] ; then echo "First alignment failed, Aborted!" ; exit 1 ; fi

cat  aligned.1.*_$N.mlf > aligned.1.mlf


# Reestimate a few more times
for j in 1 2 3 4 ; do
echo "RE-ESTIMATING MODEL PARAMETERS (ITERATION $i)"
    mkdir -p hmm$[$i +1]

    section=$[$section + 1]
    job_name=$job_root_name.$section
    script=$script_dir/$job_name
    out=$stdout_dir/$job_name
    qsub_args="$std_qsub_args -N $job_name -j y -o $out.\$TASK_ID -t 1-$N $script"

    cat << EOF > $script; chmod +x $script
#!/bin/bash

    HERest -C config -T 1 -t 250.0 150.0 1000.0 -H hmm${i}/MMF \
           -H hmm0/vFloors -I aligned.1.mlf -M hmm$[$i + 1] \
           -S train.scp.\${SGE_TASK_ID}_$N -p \$SGE_TASK_ID phone_list
EOF
    echo "qsub $qsub_args" | tee $out
    qsub $qsub_args

    job_name=$job_root_name.$section.collect
    script=$script_dir/$job_name
    out=$stdout_dir/$job_name
    qsub_args="$std_qsub_args -N $job_name -j y -o $out $script"

    cat << EOF > $script; chmod +x $script
#!/bin/bash
    HERest -H hmm${i}/MMF -H hmm0/vFloors -M hmm$[$i + 1] -p 0 \
           phone_list hmm$[$i + 1]/*.acc
EOF

    echo "qsub $qsub_args" | tee $out
    qsub $qsub_args
    if [ $? -gt 0 ] ; then echo "Collecting re-estimation $i failed, Aborted!" ; exit 1 ; fi

    i=$[$i + 1]

done

# Realignment to correct labelling:
section=$[$section + 1]
job_name=$job_root_name.$section
script=$script_dir/$job_name
out=$stdout_dir/$job_name
qsub_args="$std_qsub_args -N $job_name -j y -o $out.\$TASK_ID -t 1-$N $script"

echo "REALIGNMENT"

cat << EOF > $script; chmod u+x $script
#!/bin/bash

HVite -l \* -C config  -a -m -I aligned.0.mlf -H hmm${i}/MMF \
      -i aligned.2.\${SGE_TASK_ID}_$N.mlf -m \
      -S train.scp.\${SGE_TASK_ID}_$N -y lab dict2 phone_list
EOF

echo "qsub $qsub_args" | tee $out
qsub $qsub_args
if [ $? -gt 0 ] ; then echo "Second alignment failed, Aborted!" ; exit 1 ; fi

cat  aligned.2.*_$N.mlf > aligned.2.mlf


###
### INCREASE MISTURES
###

# Increase mixtures.

for m in 2 3 5 8 ; do
    echo "INCREASING MIXTURES TO $m"  
    mkdir -p hmm$[$i +1]
    HHEd -C config -H hmm${i}/MMF -M hmm$[$i + 1] resources/mixup${m}.hed phone_list
    if [ $? -gt 0 ] ; then echo "Mixup to $m mixtures failed, Aborted!" ; exit 1 ; fi
    i=$[$i + 1]
    for j in 1 2 3; do
        echo "RE-ESTIMATING MODEL PARAMETERS (ITERATION $i)"
	mkdir -p hmm$[$i +1]

        section=$[$section + 1]
        job_name=$job_root_name.$section
        script=$script_dir/$job_name
        out=$stdout_dir/$job_name
        qsub_args="$std_qsub_args -N $job_name -j y -o $out.\$TASK_ID -t 1-$N $script"
       
        cat << EOF > $script; chmod +x $script
#!/bin/bash

        HERest -C config -T 1 -t 250.0 150.0 1000.0 -H hmm${i}/MMF  \
               -H hmm0/vFloors -I aligned.2.mlf -M hmm$[$i + 1] \
               -S train.scp.\${SGE_TASK_ID}_$N -p \$SGE_TASK_ID phone_list


EOF
        echo "qsub $qsub_args" | tee $out
        qsub $qsub_args

        job_name=$job_root_name.$section.collect
        script=$script_dir/$job_name
        out=$stdout_dir/$job_name 
        qsub_args="$std_qsub_args -N $job_name -j y -o $out $script"
    
        cat << EOF > $script; chmod +x $script
#!/bin/bash

        HERest -H hmm${i}/MMF -H hmm0/vFloors -M hmm$[$i + 1] -p 0 \
               phone_list hmm$[$i + 1]/*.acc
EOF
    
        echo "qsub $qsub_args" | tee $out
        qsub $qsub_args
        if [ $? -gt 0 ] ; then echo "Collecting re-estimation $i failed, Aborted!" ; exit 1 ; fi
	i=$[$i + 1]
    done

done

# Final alignment

section=$[$section + 1]
job_name=$job_root_name.$section
script=$script_dir/$job_name
out=$stdout_dir/$job_name
qsub_args="$std_qsub_args -N $job_name -j y -o $out.\$TASK_ID -t 1-$N $script"

echo "FINAL ALIGNMENT"

cat << EOF > $script; chmod u+x $script
#!/bin/bash

HVite -l \* -C config  -a -m -I aligned.0.mlf -H hmm${i}/MMF \
      -i aligned.3.\${SGE_TASK_ID}_$N.mlf -T 1 \
      -S train.scp.\${SGE_TASK_ID}_$N -y lab dict2 phone_list

EOF

echo "qsub $qsub_args" | tee $out
qsub $qsub_args
if [ $? -gt 0 ] ; then echo "Final alignment failed, Aborted!" ; exit 1 ; fi

cat  aligned.3.*_$N.mlf > aligned.3.mlf

echo "END_TIME is `date`"
