#!/usr/bin/perl -w
use strict;
use Getopt::Long;
use Pod::Usage;
use Data::Dumper;
use POSIX;
use Carp;
use Util;

####################################################################
my $smooth=2; # 0 is no smoothing, 1 is add-one, 2 is simple good-turing
my $help;
my $statfile;
my $modelDir;
my $bigramFname;
GetOptions (
			"modelDir|m=s"  => \$modelDir,
			"smooth|s=i" => \$smooth,
			"adjmat|a=s" => \$bigramFname,
            "help|h" => \$help) 
or pod2usage(2);
$statfile = shift;
pod2usage(-exitstatus => 2, -verbose => 3) if ($help || !$statfile);

$modelDir = "." if (!defined($modelDir));
my $addOne=0;
$addOne=1 if ($smooth ==1);

open(STAT, "<$statfile") ||  confess  "$!: cannot open <$statfile";
#skip the header
my $dummy = <STAT>; 
$dummy = <STAT>;
$dummy = <STAT>;
$dummy = <STAT>;
$dummy = <STAT>;

#read the 1-grams
my %unigrams=();
my $minProb=0; #used to calculate the total number of bigrams
my $unigramProbMass=0;
while(<STAT>){
	chomp;
	last if ($_ eq '');
	(my $unigram, my $triphone, my $weight) = split;
	$unigrams{$triphone}=$unigram;
	$unigramProbMass += exp($unigram*log(10));
	if($triphone ne '!ENTER'){ #!ENTER is special - has 0 probability:   -99.999	!ENTER	-4342944920.0630
		$minProb = $minProb < $unigram ? $minProb : $unigram;
	}
}
my $numObsBigrams = 1/exp($minProb*log(10));
my $exitCount=ceil(exp($unigrams{'!EXIT'}*log(10))*$numObsBigrams);
print STDERR "num of bigrams: $numObsBigrams\n";
print STDERR "num of unique triphones: ".scalar(keys %unigrams)."\n";
print STDERR "total unigram prob mass: $unigramProbMass\n";
print STDERR "exit (word) counts: $exitCount\n";

#cache the legal suffixes of the preceding triphone
my %legalSuffix=();
foreach (keys %unigrams){
	next if ($_ eq 'wb' || $_ eq '!ENTER' || $_ eq '!EXIT'); #things of XX-YY form are legal prefixes for wb
	if(m/^([^\-]+)-([^+]+)(\+.*)?$/){
		my $suffix = "$1+$2";
		push @{$legalSuffix{$suffix}}, $_;
	}
	else{
		push @{$legalSuffix{'wb'}}, $_;
		#print "$_\n";
	}

}

#create all the possible bi-grams
my %bigrams=();
map {$bigrams{'wb'}{$_} = $addOne}  @{$legalSuffix{'wb'}};
foreach (keys %unigrams){
	next if ($_ eq 'wb' || $_ eq '!ENTER' || $_ eq '!EXIT'); 
	if(m/^([^\-]+)-([^+]+)$/){
		$bigrams{$_}{'wb'}=$addOne;
	}
	elsif(/^[^\-+]+$/){
		$bigrams{$_}{'wb'}=$addOne;
	}		
	elsif(m/^([^\-]+-)?([^+]+\+.*)$/){ #by now everything must have a plus in it
		my $suffix = $2;
		my $firstTriphone=$_;
		map {$bigrams{$firstTriphone}{$_}=$addOne} @{$legalSuffix{$suffix}};
	}
	else{
		print STDERR "$_\n";
	}

}
$bigrams{'wb'}{'!EXIT'}=$addOne;
$bigrams{'!ENTER'}{'wb'}=$addOne;

my $possibleBigrams=0;
map {$possibleBigrams += scalar(keys %{$_})} values %bigrams;
print STDERR "possible Bigrams: $possibleBigrams\n";

if($bigramFname){
	print STDERR "writing out bigram adjacency matrix...\n";
	open (BIGRAM, ">$bigramFname") || confess "cannot open >$bigramFname";
	print BIGRAM Dumper(\%bigrams);
	close BIGRAM;
}

#now add in the true bigram counts
$dummy = <STAT>;
my $obsBigrams=0;
my $bigramProbMass=0;
#my %bGivenA=();
while(<STAT>){
	chomp;
	last if ($_ eq '');
	$obsBigrams++;
	#print STDERR "$_\n";
	(my $bigramProb, my $firstTriphone, my $secondTriphone) = split;
	if (defined($bigrams{$firstTriphone}{$secondTriphone})){
		$bigrams{$firstTriphone}{$secondTriphone} += ceil($numObsBigrams*exp(($bigramProb+$unigrams{$firstTriphone})*log(10)));
		$bigramProbMass += exp(($bigramProb+$unigrams{$firstTriphone})*log(10));
		#$bGivenA{$firstTriphone}+=exp(($bigramProb)*log(10));
	}
	else{
		print STDERR "!!!bigram{$firstTriphone}{$secondTriphone} undefined.\n"
	}
}

close STAT;
print STDERR "num of unique observed bigrams: $obsBigrams\n";
print STDERR "total bigram Prob mass: ".($bigramProbMass+$exitCount/$numObsBigrams)."\n";
#print Dumper(\%bGivenA);
#print Dumper(\%bigrams);

my $smoothedBigramCount;
if ($smooth ==0){
	print STDERR "Not performing smoothing.\n";	
	$smoothedBigramCount = $numObsBigrams;
}
elsif ($smooth ==1){ #add-one smooth
	print STDERR "Performing add-one smoothing.\n";
	$smoothedBigramCount = $numObsBigrams+$possibleBigrams-$obsBigrams;
	print STDERR "Allocating ".($possibleBigrams-$obsBigrams)/$smoothedBigramCount." prob mass to unseen bigrams\n";
}
elsif ($smooth ==2){ #simple good-turing smooth
	print STDERR "Performing simple Good-Turing smoothing.\n";
	$smoothedBigramCount = $numObsBigrams;
	open(BIGRAMCOUNTS, ">$modelDir/bigramCounts.txt") || confess "$!: cannot open >$modelDir/bigramCounts.txt";
	foreach my $firstTriphone (keys %bigrams){
		foreach my $secondTriphone ( keys %{$bigrams{$firstTriphone}}){
			if ($bigrams{$firstTriphone}{$secondTriphone}>0 && $secondTriphone ne '!EXIT'){
				print BIGRAMCOUNTS "$bigrams{$firstTriphone}{$secondTriphone} $firstTriphone $secondTriphone\n";
			}
		}
	}
	close BIGRAMCOUNTS;

	print STDERR "running the simple good-turing estimator\n";
	ecsystem("smooth.func $modelDir/bigramCounts.txt >$modelDir/bigramSGTSmoothingFreqs.txt");

	my %smoothedFreqs=();
	#do not use it in the SGT estimate or take weight from wb !EXIT. Because 
	#it's so common - it throws off the regression line in a big way.
	$smoothedFreqs{$bigrams{'wb'}{'!EXIT'}}=$bigrams{'wb'}{'!EXIT'}; 
	open(SMOOTHEDFREQ, "<$modelDir/bigramSGTSmoothingFreqs.txt") || confess "$!: cannot open <$modelDir/bigramSGTSmoothingFreqs.txt";
	while(<SMOOTHEDFREQ>){
		chomp;
		(my $origFreq, my $smoothedFreq)=split;
		$smoothedFreqs{$origFreq}=$smoothedFreq;
	}
	close SMOOTHEDFREQ;

	#adjust the weight for the existing freqs
	my $totalCountBefore=0;
	my $totalCountAfter=0;
	foreach my $firstTriphone (keys %bigrams){
		foreach my $secondTriphone ( keys %{$bigrams{$firstTriphone}}){
			if ($bigrams{$firstTriphone}{$secondTriphone}>0){
				$totalCountBefore+=$bigrams{$firstTriphone}{$secondTriphone};
				$totalCountAfter+=$smoothedFreqs{$bigrams{$firstTriphone}{$secondTriphone}};
			}
		}
	}
	print STDERR "totalCountBefore: $totalCountBefore totalCountAfter $totalCountAfter exitCount $exitCount\n";
	print STDERR "Allocating ".(($numObsBigrams-$totalCountAfter)/$numObsBigrams)." prob mass to unseen bigrams\n";
	if($possibleBigrams-$obsBigrams>0){
		$smoothedFreqs{0}=($numObsBigrams-$totalCountAfter)/($possibleBigrams-$obsBigrams);
		print STDERR "Each unseen bigram is counted as seen $smoothedFreqs{0} times.\n";
		foreach my $firstTriphone (keys %bigrams){
			foreach my $secondTriphone ( keys %{$bigrams{$firstTriphone}}){
				$bigrams{$firstTriphone}{$secondTriphone} = $smoothedFreqs{$bigrams{$firstTriphone}{$secondTriphone}};
			}
		}
	}
	else{
		print STDERR "all possible bigrams were seen in the original dataset! Will not smooth bigrams.\n";
	}
}

#now print out the new smoothed matrix file 
print STDERR "writing out the smoothed matrix file...\n";
print "\\data\\\n";
my $uniqUnigramCount=scalar(keys %bigrams)+1;
my $uniqBigramCount=0;
map {$uniqBigramCount += scalar(keys %$_)} values %bigrams;
print "ngram 1=$uniqUnigramCount\n";
print "ngram 2=$uniqBigramCount\n";
print "\n";
print "\\1-grams:\n";
#!ENTER and !EXIT are special
print "-99.999	!ENTER	-4342944920.0630\n";

my %smoothedMonograms=();
foreach (sort grep {$_ ne '!ENTER'} keys %bigrams){
	my $smoothUnigramProb=0;
	map {$smoothUnigramProb += $_} values %{$bigrams{$_}};
	$smoothUnigramProb = log($smoothUnigramProb / $smoothedBigramCount)/log(10);
	$smoothedMonograms{$_}=$smoothUnigramProb;
	print "$smoothUnigramProb	$_	-4342944920.0630\n";
}
my $exitProb= log($exitCount/$smoothedBigramCount)/log(10);
print "$exitProb\t!EXIT\n";
print "\n";
print "\\2-grams:\n";
print "0.0000\t!ENTER\twb\n";
foreach my $firstTriphone (sort grep {$_ ne '!ENTER'} keys %bigrams){
	foreach my $secondTriphone ( sort keys %{$bigrams{$firstTriphone}}){
		my $smoothBigramProb;
		if ($bigrams{$firstTriphone}{$secondTriphone}>0){
			$smoothBigramProb = log($bigrams{$firstTriphone}{$secondTriphone}/ $smoothedBigramCount)/log(10)-$smoothedMonograms{$firstTriphone};
		}
		else{
			$smoothBigramProb = -99.9999;
		}
		print "$smoothBigramProb\t$firstTriphone\t$secondTriphone\n";
	}
}
print "\n";
print "\\end\\\n";

=head1 NAME

smoothedBigram.pl - creates a smoothed bigram language model over triphones.  Output and input are in the ARPA MIT-LL format.
Input filename is passed on the command line, output is printed on STDOUT.

=head1 SYNOPSIS

smoothedBigram.pl [options] statFile.mtr

 Required:
   none
 Options:
   -h,	--help			brief help message
   -s,  --smooth=0|1|2  0 is no smoothing, 1 is add-one smoothing , 3 is simple good-turing smoothing
   -a,	--adjmat=<fname> if specified, writes the bigram hashtable to fname
   -m,  --modelDir=<dirname> where the intermediate files are stored default is ./

=head1 OPTIONS

=over 4

=item B<--help>

Prints this help message and exits.


=back

=head1 DESCRIPTION
