#!/usr/bin/perl -w

# ./mkmodel.pl train_data

use strict;
use File::Basename;

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

if (@ARGV != 1) {
    print STDERR "Usage: mkmodel.pl <training data dir>\n";
    exit 2;
}

my ($train_data) = @ARGV;

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

print "Clearing old results...\n";

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

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

my $current_hmm = 0;

# Figure out which label and sound files to use

print "Figuring out which sound files are transcribed...\n";

my @labelfiles = glob "$train_data/*.lab";
my @allsoundfiles = glob "$train_data/*.wav";
my @soundfiles = ();
foreach (@allsoundfiles) {
    (my $l = $_) =~ s/\.wav$/.lab/;
    if (! -f $l) {
	print STDERR "No transcription for $_, skipping.\n";
    } elsif ((stat($l))[7] == 0) {
	print STDERR "Transcription for $_ is empty, skipping.\n";
    } else {
	push @soundfiles, $_;
    }
}

printf("Using %d out of %d files.\n", 
       scalar(@soundfiles), 
       scalar(@allsoundfiles));

# Create a list of all the phones used in the training data.

print "Creating monophone list from training data...\n";

my %monophones = ();
foreach my $labf (@labelfiles) {
    open LABF, $labf;
    while (<LABF>) {
	chomp;
	s/^\d+\s+\d+\s+(\S+).*$/$1/;
	$monophones{$_} = 1;
    }
    close LABF;
}

open MONO, ">$tmpdir/monophones1";
foreach (sort (keys %monophones)) {
    print MONO "$_\n";
}
close MONO;

# Create a monophone list without the short pause model

run("grep -v '^sp\$' '$tmpdir/monophones1' > '$tmpdir/monophones0'");

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

print "Parametrizing training data...\n";

open PARAM_CONFIG, ">$tmpdir/param_config";
print PARAM_CONFIG<<END;
SOURCEFORMAT = WAV
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 (@soundfiles) {
    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;

run("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;

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";
    } elsif (!/^SOURCEFORMAT/) {
	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;
run("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;

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

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

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

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

reestimate("$tmpdir/train.scp", "$tmpdir/monophones0");
reestimate("$tmpdir/train.scp", "$tmpdir/monophones0");
reestimate("$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();
run("HHEd32 -H '$indir2/macros' -H '$indir2/hmmdefs' -M '$outdir2' $tmpdir/sil.hed $tmpdir/monophones1");

# Re-estimate two more times.

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


# FIXME: do some more stuff here?


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

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


# FIXME: build triphone models?

print "\n";
print "Done.\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";



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(file list file, model list file)
sub reestimate {
    print "Re-estimating...\n";
    my ($file_list, $model_list) = @_;
    my ($indir, $outdir) = new_hmm();
    run("HERest32 -T 1 -C $tmpdir/train_config -t 250.0 150.0 1000.0"
	   . " -L '$train_data' -S '$file_list' -H '$indir/macros' -H '$indir/hmmdefs'"
	   . " -M '$outdir' '$model_list'");
}

sub run {
    my $command = shift;
    print "$command\n";
    system($command) == 0 || die "Failed: $?";
}



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