#!/usr/bin/perl
# augment_lattice.pl: Mark Hasegawa-Johnson, March 2004
# Illinois Speech and Language Engineering
#
# Usage text (printed if augment_lattice.pl is called with no arguments):
$USAGE=<<ENDUSAGE;

augment_lattice.pl: Add tags to the edges in a lattice

Usage:
 augment_lattice.pl script.scp [p=file.dict] [n=1]

augment_lattice.pl reads in each (HTK-format) input lattice file
specified in script.scp, and adds tags as specified on the command
line.  The script file should contain input & output filenames in
pairs, i.e., "input1.lat output1.lat input2.lat output2.lat ...".
Each input lattice file should be for exactly one utterance.  If
lattice files contain more than one utterance each, split them first
using split_lattice.pl.  

Currently implemented tags include the following:

p=file.dict (Pronunciation model)___________

file.dict is a standard-format dictionary file: each line contains a
word, followed by space-separated phonemes.  For each edge in a
lattice, the word label is sought in the dictionary.  If found, its
pronunciation is added to the edge with a tag in the form
p=ph1:...:phN, where phN is the Nth phoneme in the dictionary entry.

n=1 (Mark the MAP path)____________

Find the maximum probability path through the lattice, and add the tag
"n=1" to all edges on the maximum probability path.  Maximum
probability is computed by scaling the acoustic model (a= tag) by
ACOUSTIC_STREAMWEIGHT, scaling the language model (l= tag) by
LM_STREAMWEIGHT, and adding.

ENDUSAGE

####################################################################
# GLOBAL VARIABLES

# Conversion rates for various file formats
$MLFPERIOD = 0.0000001;

# Default stream weights and insertion penalty
$ACOUSTIC_STREAMWEIGHT=1;
$LM_STREAMWEIGHT=9.24283;
$INSERTION_PENALTY=-3.94958;

# Math functions
sub abs { my($x)=shift; ($x>0) ? $x : -$x ; }
sub round { $_[0] =~ /([^\.]*)(.*)/; ($2 >= 0.5) ? $1+1 : $1; }

# Compute the probability of an edge
#   Weight the acoustic likelihood by the acoustic streamweight,
#   weight the language model likelihood by the LM streamweight,
#   add the insertion penalty,
#   return
sub edge_likelihood { 
  my($e)=shift; 
  return($$e{'a'}*$ACOUSTIC_STREAMWEIGHT +
         $$e{'l'}*$LM_STREAMWEIGHT +
         $INSERTION_PENALTY);
}
  
###################################################
# Subroutine to read an HTK lattice file
#
# Syntax: ($nodes, $edges, $header) = read_HTK_lattice(\*FILEHANDLE)
#
# $nodes is an arrayref of nodes
#    each node is an arrayref containing the following data:
#    [ back, forw, time ]
#    where forw is an arrayref of indices of forward edges,
#    back is an arrayref of indices of backward edges.
#    time is the time-alignment of this node
# $edges is an arrayref of edges
#    each edge, $e, is a hashref
#    tags on input edge are interpreted as key=val pairs,
#    and loaded as $$e{$key}=$val
#    For example, $$e{'W'} is set to the word label
#    
# $header is a hashref containing the header information
#
sub read_HTK_lattice {
  my($fh) = $_[0];
  my(@nodes) = ();
  my(@edges) = ();
  my(%header) = ();
  my($line);

  # Read lines of the file
  while($line = <$fh>) {
    # Node number given by I=
    if($line =~ /I=(\d+)/) {
      my($nodenum) = $1;
      unless(exists($nodes[$nodenum])) { $nodes[$nodenum] = [ [],[],0 ]; }
      if($line =~ /t=([\.\d]+)/) { $nodes[$nodenum][2] = $1; }
    }
    # Edge number given by J=
    elsif($line =~ s/J=(\d+)//) {
      my($edgenum) = $1;
      my(%edge)=();

      # Look for the startnode marker
      if($line =~ s/S=(\d+)//) { 
	$edge{'S'} = $1;
	# Tell the start node that I am one of its forward edges
	push(@{$nodes[$1][1]}, $edgenum);
      }
      
      # Look for the endnode marker
      if($line =~ s/E=(\d+)//) { 
	$edge{'E'} = $1; 
	# Tell the end node that I am one of its backward edges
	push(@{$nodes[$1][0]}, $edgenum);
      }
      
      # Read all other field-value pairs
      while($line =~ s/(\S+)=(\S+)//) { $edge{$1}=$2; }
      # Store the edge
      $edges[$edgenum] = \%edge;
    }
    # Otherwise, try to read header information
    else {
      while ($line =~ s/(\S+)=(\S+)//) { $header{$1} = $2; }
    }    
  }
  
  return(\@nodes, \@edges, \%header);
}

###################################################
# Subroutine to write an HTK lattice file
#
# Syntax: ($N, $L) = write_HTK_lattice($nodes, $edges, $header, \*FILEHANDLE)
#
# $nodes is an arrayref of nodes
# $edges is an arrayref of edges
# $header is a hashref containing HTK lattice header information
# FILEHANDLE is the opened output file
#
# $N is the number of nodes actually written to the file
# $L is the number of edges actually written to the file
#
sub write_HTK_lattice {
  my($nodes) = $_[0];
  my($edges) = $_[1];
  my($header) = $_[2];
  my($fh) = $_[3];
  my($N) = 0;
  my($L) = 0;

  # Print header information
  print $fh "VERSION=1.0\n";
  print $fh "UTTERANCE=",$$header{'UTTERANCE'},"\n";
  print $fh "lmname=",$$header{'lmname'},"\n";
  print $fh "lmscale=",$$header{'lmscale'}," ";
  print $fh "wdpenalty=",$$header{'wdpenalty'},"\n";
  print $fh "vocab=",$$header{'vocab'},"\n";
  print $fh sprintf("N=%-5d L=%-7d\n",0+@{$nodes},0+@{$edges});

  # Print nodes
  foreach $N (0..$#{$nodes}) {
    print $fh sprintf("I=%-5d ",$N);
    if (exists($$nodes[$N][2])) 
      { print $fh sprintf("t=%4.2f",$$nodes[$N][2]); }
    print $fh "\n";
  }
  # Print edges
  foreach $L (0..$#{$edges}) {
    my($e)=$$edges[$L];

    # Print the standard key-value pairs
    print $fh sprintf("J=%-7d S=%-5d E=%-5d W=%-20s ",$L, 
		      $$e{'S'}, $$e{'E'}, $$e{'W'});
    # Print all other key-value pairs
    foreach (keys %$e) {
      unless(/^[SEW]$/) { print $fh sprintf("%s=%s ", $_, $$e{$_}); }
    }
    # Print newline
    print $fh "\n";
  }
  return($N,$L);
}

###################################################
# Find the maximum-likelihood path through a lattice
#
# ( $optnodes, $optedges, $logprob ) = viterbi($nodes, $edges)
#
# $nodes - arrayref of nodes
# $edges - arrayref of edges
# $optnodes - arrayref of node indices along the optimum path
# $optedges - arrayref of edge indices along the optimum path
# $logprob - the log probability of the optimum path
# 
# This function makes a few critical assumptions:
#  1. the start node of any edge has lower number than its end node
#  2. a complete path through the lattice is any path connecting
#     node 0 to node $#node.
#  3. there is a path from node 0 forward to any node in the network
#
sub viterbi {
  my($nodes) = $_[0];
  my($edges) = $_[1];

  # $bp[$i] is the optimum backpointer from node $i
  my(@bp) = ( 0 ); 
  # $ep[$i] is the edge leading backward from node $i
  my(@ep) = ( 0 ); 
  # $lp[$i] is the maximum log probability along any path to node $i
  my(@lp) = ( 0 ); 

  # Find log probability and backpointer at each node
  foreach $i (1..$#{$nodes}) {
    $lp[$i] = -10000000000000;
    # Search through all edges going backward from this node
    foreach $ei (@{$$nodes[$i][0]}) {
      # Test value: logprob of previous node, plus acoust, plus lm
      my($test) = $lp[$$edges[$ei]{'S'}] + edge_likelihood($$edges[$ei]);
      if($test > $lp[$i]) { 
	$lp[$i] = $test;
	$bp[$i] = $$edges[$ei]{'S'};
	$ep[$i] = $ei;
      }
    }
  }
  
  # Now perform traceback, and put edge indices into result
  my(@optnodes) = ( $#{$nodes} );
  my(@optedges) = ();
  while ($optnodes[0] > 0) {
    unshift(@optedges, $ep[$optnodes[0]]); 
    unshift(@optnodes, $bp[$optnodes[0]]); 
  }
  return(\@optnodes, \@optedges, $lp[$#nodes]);
}


###################################################
# normalize_word:
# Remove distinctions not useful for scoring a lattice
# from a word in input transcription or on a lattice edge.
#
# This function should be same as normalize_words.pl,
# except that filled pauses, noise, and laughter are
# replace with a null string "".
#
# $outword = normalize_word($inword)
#
sub normalize_word {
  # Set the word to upper case
  my($l) = uc($_[0]);
    
  # Any compound including the word "laugh" means "laugh"
  # For rough phonetic transcription, let's replace this with HAHA
  $l =~ s/.*LAUGH.*/HAHA/;
  
  # These words seem to be used for various filled pauses
  #   Replace all with UH or UM
  $l =~ s/.*(EPENTHETIC|FILLEDSIL|FALSESTART|HESITATION|GLOTTAL).*/UH/;
  $l =~ s/.*(TRUNCATEDP).*/UH/;
  $l =~ s/.*(NASAL|MMM).*/HMM/;
  # $l =~ s/\b(UH|UM|UMM)\b//;
  
  # Eliminate various background noise notations
  $l =~ s/\bSIL\b//;
  
  # These words mean silence even if extended, e.g. SILENCEDUETOSPLICE
  $l =~ s/.*(H\#|CROFF|UNKNOWN|CLICK|SILENCE|BREATH).*//;
  $l =~ s/.*(CROSSTALK|STATIC|NOISE|SMACK|BUZZ|HUMM).*//;
  
  # Replace periods by apostrophes
  $l =~ s/\./\'/g;
  # Eliminate second in any series of apostrophes
  $l =~ s/\'\'/\'/g;    
  
  # Eliminate commentary between brackets, braces, or parentheses,
  #  after underscore, or after ? or ! or #
  # This should come last, so that # doesn't keep H# from being used
  $l =~ s/\[.*\]//;
  $l =~ s/\{.*\}//;
  $l =~ s/\(.*\)//;
  $l =~ s/(_|\!|\?|\#).*//;
  
  # Eliminate quotation marks
  $l =~ s/\"//g;

  # Eliminate plus
  $l =~ s/\+//g;

  # Return outword
  return($l);
}


##################################################
# Main program

if (@ARGV==0) { print $USAGE; exit; }

# Which tags need to be computed?
%TAGS=();

# Read in the script file
$arg=shift;
open(F,$arg) || die "Unable to read from script file $arg: $!\n";
$script = ''; while (<F>) { $script .= $_; }
close(F);

# Read in the tag options
while ($arg=shift) {

  # p tag: read in the dictionary 
  if ($arg =~ /p=(\S+)/) {
    $DICTFILE=$1;
    $TAGS{'p'}=1;
    open(DICTFILE, $DICTFILE) || die "Unable to read dict from $DICTFILE\n";
    foreach (<DICTFILE>) { 
      chomp; chomp;
      my(@entry)=split(/\s+/,$_);
      $word=uc(shift(@entry));
      $DICTIONARY{$word}=\@entry;
    }
    close(DICTFILE);
  }

  # n tag
  if ($arg =~ /n=(\S+)/) {
    $TAGS{'n'}=$1;
  }
}

#######################
# Main loop: process each pair of lattice files in the script
while($script =~ s/^\s*(\S+)\s+(\S+)//) {
  $inlat = $1; $outlat = $2;

  # Read the input lattice
  unless(open(LATFILE, $inlat)) {
    print STDERR "Unable to read from lattice $inlat: $!\n";
    next;
  }
  @graph = read_HTK_lattice(\*LATFILE);
  close(LATFILE);

  # Process the p tag option if specified
  if($TAGS{'p'}) {
    # Set the array interpolator to colon
    $"=':';

    # Process each edge
    foreach $e (@{$graph[1]}) {
      # Split multiwords
      @w=split(/_/,$$e{'W'});
      @pron=();
      foreach (@w) {
	# Get rid of symbols used to mean "silence", "laughter", etc
	$w=normalize_word($_);
	# If DICTIONARY knows this word, add it to @pron
	if(exists($DICTIONARY{$w})) {push(@pron, @{$DICTIONARY{$w}});}
      }
      # Add colon-separated pronunciation as the p-tag
      $$e{'p'} = "@pron";
    }
    # Reset the array interpolator
    $"=' ';
  }

  # Process the n tag option if specified
  if($TAGS{'n'}) {
    # find maxprob path
    @maxprob = viterbi(@graph[0..1]);

    # DEBUG: write out the edge words on MAP path
    print STDERR "MAP PATH: ";
    foreach $i (0..$#{$maxprob[1]}) 
      { print STDERR ",",$graph[1][$maxprob[1][$i]]{'W'}; }
    print STDERR "\n";

    # Add "n=1" to each edge in the MAP path
    foreach $i (0..$#{$maxprob[1]}) 
      { $graph[1][$maxprob[1][$i]]{'n'} = 1; }
  }
    
  # Write the modified lattice
  unless(open(LAT_OUT, ">$outlat")) {
    print STDERR "Unable to write lattice out to $outlat: $!\n";
    next;
  }
  write_HTK_lattice(@graph, \*LAT_OUT);
  close(LAT_OUT);

  # Write summary to STDERR
  $nn=0+@{$graph[0]};
  $ne=0+@{$graph[1]};
  print STDERR "$inlat -> $outlat, $nn nodes, $ne edges\n";
}
