#!/usr/bin/perl -w

# ./mkrec.pl grammar/NumeralsSwe.gf numerals-swe.dict train_utts.txt train_data

use strict;
use File::Basename;

my $prompts2mlf = check_prog("prompts2mlf");
my $proto2hmmdefs = check_prog("proto2hmmdefs");
my $makespmodel = check_prog("makespmodel.pl");
my $maketrihed = check_prog("maketrihed");
my $mkwordlist = check_prog("mkwordlist.pl");

if (@ARGV != 4) {
    print STDERR "Usage: mkrec.pl <GF grammar> <dictionary> <training utterance file> <training data dir>\n";
    exit 2;
}

my ($gf_grammar,$dictionary,$train_utts,$train_data) = @ARGV;

if (! -f $gf_grammar) {
    print STDERR "GF grammar '$gf_grammar' not found\n";
    exit 2;
}

if (! -f $dictionary) {
    print STDERR "Dictionary '$dictionary' not found\n";
    exit 2;
}

if (! -f $train_utts) {
    print STDERR "Training utterance file '$train_utts' not found\n";
    exit 2;
}

if (! -d $train_data) {
    print STDERR "Training data directory '$train_data' not found\n";
    exit 2;
}

my $tmpdir = "tmp";
run("rm -rf '$tmpdir'");
mkdir $tmpdir;

my $outdir = "final";
run("rm -rf '$outdir'");
mkdir $outdir;

my $current_hmm = 0;

# Create a list of all words used in the grammar.

print "Creating list of all words...\n";

run("$mkwordlist '$gf_grammar' > $tmpdir/wordlist.txt");

# Get the pronunciation of all the words from
# the dictionary to form an application-specific dictionary.

# Create a list of all the phones used in the application
# dictionary.

# Make all words upper case and add word-end pauses 
# for each pronunciation.

open MKDICT, ">$tmpdir/mkdict.ded";
print MKDICT<<END;
AS sp
UW
END
close MKDICT;

print "Creating application dictionary and monophone list...\n";
system "HDMan32 -g $tmpdir/mkdict.ded -m -w $tmpdir/wordlist.txt -n $tmpdir/monophones1 -l $tmpdir/dlog $tmpdir/dict.nosil '$dictionary'";

# Add a special word for silence to the dictionary and monophones list.

open DICT, ">$tmpdir/dict";
print DICT "!SILENCE        sil\n";
open DICT_NOSIL, "$tmpdir/dict.nosil";
while (<DICT_NOSIL>) {
    print DICT;
}
close DICT_NOSIL;
close DICT;

my $outdict = $outdir . "/" . replace_ext($gf_grammar, "dct");
system "cp '$tmpdir/dict' '$outdict'";


if (! grep(/^sil$/,`cat $tmpdir/monophones1`)) {
    run("echo sil >> $tmpdir/monophones1");
}

# Translate the list of training utterances to
# a word-level transcription file.

print "Creating word-level MLF file for training utterances...\n";
run("$prompts2mlf '$train_utts' > $tmpdir/train_utts_words.mlf");

open MKPHONES0, ">$tmpdir/mkphones0.led";
print MKPHONES0<<END;
EX
IS sil sil
DE sp
END
close MKPHONES0;

# Create a phone-level transcription file from the word-level
# transcriptions and the dictionary. The transcriptions use the 
# first available pronunciation for each word. This problem will
# will be taken care of in later steps.

print "Creating phone-level MLF file for training utterances...\n";
system "HLEd32 -l '*' -d $tmpdir/dict -i $tmpdir/train_utts_phone.mlf $tmpdir/mkphones0.led $tmpdir/train_utts_words.mlf";

# Decide on a parametrization configuration. Currently this
# uses the fixed configuration recommended in the HTK tutorial.

print "Parametrizing...\n";

open PARAM_CONFIG, ">$tmpdir/param_config";
print PARAM_CONFIG<<END;
TARGETKIND = MFCC_0
TARGETRATE = 100000.0
SAVECOMPRESSED = T
SAVEWITHCRC = T
WINDOWSIZE = 250000.0
USEHAMMING = T
PREEMCOEF = 0.97
NUMCHANS = 26
CEPLIFTER = 22
NUMCEPS = 12
ENORMALISE = F
END
close PARAM_CONFIG;

# Parametrize the recorded data.

mkdir "$tmpdir/train_param";

open CODETR, ">$tmpdir/codetr.scp";
open TRAIN, ">$tmpdir/train.scp";
foreach my $u (glob "$train_data/*") {
    my ($f,$dir,$ext) = fileparse($u, qr/\.[^.]*/);
    my $mfc = "$tmpdir/train_param/$f.mfc";
    print CODETR "$u $mfc\n";
    print TRAIN "$mfc\n";
}
close TRAIN;
close CODETR;

system "HCopy32 -C $tmpdir/param_config -S $tmpdir/codetr.scp";

# Create a simple prototype model to start each model from.

print "Creating prototype model...\n";

open PROTO, ">$tmpdir/proto";
print PROTO<<END;
~o <VecSize> 39 <MFCC_0_D_A>
~h "proto"
<BeginHMM>
     <NumStates> 5
     <State> 2
        <Mean> 39
          0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
        <Variance> 39
          1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0
     <State> 3
        <Mean> 39
          0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
        <Variance> 39
          1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0
     <State> 4
        <Mean> 39
          0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
        <Variance> 39
          1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0
     <TransP> 5
      0.0 1.0 0.0 0.0 0.0
      0.0 0.6 0.4 0.0 0.0
      0.0 0.0 0.6 0.4 0.0
      0.0 0.0 0.0 0.7 0.3
      0.0 0.0 0.0 0.0 0.0
<EndHMM>
END
close PROTO;

print "Creating initial models...\n";

open PARAM_CONFIG, "$tmpdir/param_config";
open TRAIN_CONFIG, ">$tmpdir/train_config";
while (<PARAM_CONFIG>) {
    if (/^TARGETKIND/) {
	print TRAIN_CONFIG "TARGETKIND = MFCC_0_D_A\n";
    } else {
	print TRAIN_CONFIG;
    }
}
close TRAIN_CONFIG;
close PARAM_CONFIG;

# Set the mean and variance of the prototype to be the
# mean and variance of the whole data set.

my $outdir0 = current_hmm_dir();
mkdir $outdir0;
system "HCompV32 -C $tmpdir/train_config -f 0.01 -m -S $tmpdir/train.scp -M '$outdir0' $tmpdir/proto";

open MACROS, ">$tmpdir/hmm0/macros";
print MACROS "~o <VecSize> 39 <MFCC_0_D_A>\n";
open VFLOORS, "$tmpdir/hmm0/vFloors";
while (<VFLOORS>) {
    print MACROS;
}
close VFLOORS;
close MACROS;

open MONOPHONES1, "$tmpdir/monophones1";
open MONOPHONES0, ">$tmpdir/monophones0";
while (<MONOPHONES1>) {
    if (!/^sp$/) {
	print MONOPHONES0;
    }
}
close MONOPHONES0;
close MONOPHONES1;

# Make a copy of the prototype model for each of the phones 
# in the dictionary.

run("$proto2hmmdefs $tmpdir/hmm0/proto $tmpdir/monophones0 > $tmpdir/hmm0/hmmdefs");

# Re-estimate the models three times using the training data.

reestimate("$tmpdir/train_utts_phone.mlf", "$tmpdir/train.scp", "$tmpdir/monophones0");
reestimate("$tmpdir/train_utts_phone.mlf", "$tmpdir/train.scp", "$tmpdir/monophones0");
reestimate("$tmpdir/train_utts_phone.mlf", "$tmpdir/train.scp", "$tmpdir/monophones0");

# Fix the silence model by introducing a new short pause
# model tied to the existing silence model, and making the existing
# silence model more robust by adding backwards transitions to it.

print "Fixing the silence model...\n";

my ($indir1, $outdir1) = new_hmm();

run("cp '$indir1/macros' '$outdir1/macros'");
run("$makespmodel < '$indir1/hmmdefs' > '$outdir1/hmmdefs'");

open SIL_HED, ">$tmpdir/sil.hed";
print SIL_HED<<END;
AT 2 4 0.2 {sil.transP}
AT 4 2 0.2 {sil.transP}
AT 1 3 0.3 {sp.transP}
TI silst {sil.state[3],sp.state[2]}
END
close SIL_HED;

print "Adding extra transitions and tying sp to sil...\n";
my ($indir2, $outdir2) = new_hmm();
system "HHEd32 -H '$indir2/macros' -H '$indir2/hmmdefs' -M '$outdir2' $tmpdir/sil.hed $tmpdir/monophones1";

# Re-estimate two more times.

reestimate("$tmpdir/train_utts_phone.mlf", "$tmpdir/train.scp", "$tmpdir/monophones1");
reestimate("$tmpdir/train_utts_phone.mlf", "$tmpdir/train.scp", "$tmpdir/monophones1");

# Select the closest pronunciation from the dictionary 
# for each word by using  the recognizer with the current model.

print "Re-aligning the training data...\n";

my $hmm_dir = current_hmm_dir();
system "HVite32 -l '*' -o SWT -b '!SILENCE' -C $tmpdir/train_config -a"
       . " -H '$hmm_dir/macros' -H '$hmm_dir/hmmdefs' -i $tmpdir/train_utts_aligned.mlf"
       . " -m -t 250.0 -y lab -I $tmpdir/train_utts_words.mlf -S $tmpdir/train.scp $tmpdir/dict $tmpdir/monophones1";

# Remove utterances for which the previous step failed from the data 
# set (this seems to be quite infrequent).

print "Dropping data which did not survive re-alignment...\n";

open ALIGNED_MLF, "$tmpdir/train_utts_aligned.mlf";
open TRAIN_ALIGNED, ">$tmpdir/train-aligned.scp";
while (<ALIGNED_MLF>) {
  if (s/^\"\*\/(.*)\.lab\"$/$1/){
      chomp;
      print TRAIN_ALIGNED "$tmpdir/train_param/$_.mfc\n";
  }
}
close TRAIN_ALIGNED;
close ALIGNED_MLF;

print "Dropped data:\n";
system "diff $tmpdir/train.scp $tmpdir/train-aligned.scp";

# Re-estimate two more times.

reestimate("$tmpdir/train_utts_aligned.mlf", "$tmpdir/train-aligned.scp", "$tmpdir/monophones1");
reestimate("$tmpdir/train_utts_aligned.mlf", "$tmpdir/train-aligned.scp", "$tmpdir/monophones1");

# We now have a set of monophone models which can be used for 
# recognition. We will now proceed to build triphone models.

my $monophone_hmm = current_hmm_dir();
system "cp -a '$monophone_hmm' $outdir/hmm_mono";
system "cp $tmpdir/monophones1 $outdir/hmm_mono/hmmlist";

# List all triphones in the dictionary.
# The HTK tutorial uses the list of triphones in the training data instead of 
# the triphones in the dictionary. Since we want to get 
# away with using a small data set, it is likely that not all triphones 
# in the dictionary are encountered in the training data.

print "Creating triphone list from dictionary...\n";

open MKTRIDICT, ">$tmpdir/mktridict.ded";
print MKTRIDICT<<END;
TC
END
close MKTRIDICT;

system "HDMan32 -b sp -g $tmpdir/mktridict.ded -n $tmpdir/triphones1 /dev/null $tmpdir/dict";

# Translate the existing monophone training utterance transcriptions
# to use triphones.

print "Annotating training data with triphones...\n";

open MKTRI_LED, ">$tmpdir/mktri.led";
print MKTRI_LED<<END;
WB sp
WB sil
TC
END
close MKTRI_LED;

system "HLEd32 -l '*' -i $tmpdir/train_utts_wintri.mlf $tmpdir/mktri.led $tmpdir/train_utts_aligned.mlf";

run("$maketrihed $tmpdir/monophones1 $tmpdir/triphones1 > $tmpdir/mktri.hed");

# For each triphone in the dictionary we make a copy of the 
# monophone model for the center sound of the triphone. This also
# includes the biphones seen at word boundaries (and the monophones
# resulting from single-phone words, which are not uncommon in
# Swedish).

print "Cloning monophone models...\n";

my ($indir3, $outdir3) = new_hmm();

system "HHEd32 -H '$indir3/macros' -H '$indir3/hmmdefs' -M '$outdir3' $tmpdir/mktri.hed $tmpdir/monophones1";

# Re-estimate the triphone models twice using the 
# triphone-transcribed data.

reestimate("$tmpdir/train_utts_wintri.mlf", "$tmpdir/train-aligned.scp", "$tmpdir/triphones1");
reestimate("$tmpdir/train_utts_wintri.mlf", "$tmpdir/train-aligned.scp", "$tmpdir/triphones1");

my $triphone_hmm = current_hmm_dir();

# We now have a set of triphone models which can be used for 
# recognition.

# Here the HTK tutorial goes on to create tied-state triphone models,
# in order to make more robust models. However, this would 
# seem to require manually deciding which triphones to cluster together
# and has been skipped to keep the process simple.

run("cp -a '$triphone_hmm' $outdir/hmm_tri");
run("cp $tmpdir/triphones1 $outdir/hmm_tri/hmmlist");

my $network = $outdir . "/" . replace_ext($gf_grammar, "net");
print "Creating recognition network '$network'...\n";
open GF, "|gf -s '$gf_grammar' > /dev/null";
print GF "pg -printer=slf | wf '$network'\n";
close GF;

print "\n";
print "Done.\n";
print "\n";
print "Dictionary: $outdict\n";
print "Recognition network: $network\n";
print "\n";
print "Monophone model:\n";
print "HMMLIST = $outdir/hmm_mono/hmmlist\n";
print "MMF0    = $outdir/hmm_mono/macros\n";
print "MMF1    = $outdir/hmm_mono/hmmdefs\n";
print "Evaluate with:\n";
print "eval_rec.pl $outdict $network <test utterance file> <test data dir> "
    . "$outdir/hmm_mono/hmmlist $outdir/hmm_mono/macros $outdir/hmm_mono/hmmdefs\n";
print "\n";
print "Triphone model:\n";
print "HMMLIST = $outdir/hmm_tri/hmmlist\n";
print "MMF0    = $outdir/hmm_tri/macros\n";
print "MMF1    = $outdir/hmm_tri/hmmdefs\n";
print "Evaluate with:\n";
print "eval_rec.pl $outdict $network <test utterance file> <test data dir> "
    . "$outdir/hmm_tri/hmmlist $outdir/hmm_tri/macros $outdir/hmm_tri/hmmdefs\n";
print "\n";


sub current_hmm_dir {
    return "$tmpdir/hmm$current_hmm";
}

sub new_hmm {
    my $indir = current_hmm_dir();
    $current_hmm++;
    my $outdir = current_hmm_dir(); 
    mkdir $outdir;
    return ($indir, $outdir);
}

# reestimate(mlf, file list file, model list file)
sub reestimate {
    print "Re-estimating...\n";
    my ($mlf, $file_list, $model_list) = @_;
    my ($indir, $outdir) = new_hmm();
    system("HERest32 -C $tmpdir/train_config -I '$mlf' -t 250.0 150.0 1000.0"
	   . " -S '$file_list' -H '$indir/macros' -H '$indir/hmmdefs'"
	   . " -M '$outdir' '$model_list'") == 0 
	   || die "Re-estimation failed: $?";
}

sub run {
    my $command = shift;
    print "$command\n";
    system $command;
}

sub check_prog {
    my $prog_dir = (fileparse($0, qr/\.[^.]*/))[1];
    my $prog = ($prog_dir ne "") ? 
	         (($prog_dir =~ /\/$/) ? $prog_dir.$_[0] : $prog_dir."/".$_[0])
		 : $_[0];

    if (! -x $prog) {
	print STDERR "$prog is not an executable file.\n";
	print STDERR "Maybe you need to run 'make'?\n";
	exit 1;
    }

    return $prog;
}

sub allwords {
    my $grammar = $_[0];

    my @r = ();

    open GF, "echo 'pg -printer=fullform' | gf -s -batch -nocpu '$grammar' |";
    while (<GF>) {
	if (!/^\s*$/) {
	    chomp;
	    s/([^:]*)\s:.*/$1/;
	    push @r, $_;
	}
    }
    close GF;

    return @r;
}

sub replace_ext {
  return (fileparse($_[0], qr/\.[^.]*/))[0] . "." . $_[1];
}
