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

pinch_lattice.pl: Score lattices against a reference transcription.

Usage:
 pinch_lattice.pl script.scp

Pinch_lattice.pl reads in each (HTK-format) input lattice file
specified in script.scp, finds the maximum probability (MAP) path
through the lattice, and pinches the lattice to its MAP path.
Specifically, any edge that is not part of the MAP path is detached
from its original start and nodes, and re-attached to the pair of
nodes, in the MAP path, that bracket its original temporal midpoint.
After all edges have been aligned, and all extra nodes deleted, edges
spanning the same time segment and carrying the same word label are
merged.

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.

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 likelihood 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[3]*$ACOUSTIC_STREAMWEIGHT +
         $$e[4]*$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 is an arrayref as follows:
#    [ sn, en, word, ap, lp, rem ]
#    sn is the index of the start node
#    en is the index of the end node
#    word is the word on this node
#    ap is the acoustic probability
#    lp is the language model probability
#    rem is the remainder of the line, after other entries removed
#    
# $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>) {
    chomp $line; chomp $line; 

    # 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) = (0,0,'',0,0);
      if($line =~ s/S=(\d+)//) 
      { 
	$edge[0] = $1;
	# Tell the start node that I am one of its forward edges
	push(@{$nodes[$1][1]}, $edgenum);
      }
      if($line =~ s/E=(\d+)//) 
      { 
	$edge[1] = $1; 
	# Tell the end node that I am one of its backward edges
	push(@{$nodes[$1][0]}, $edgenum);
      }
      # Get the word label
      if($line =~ s/W=(\S+)//) { $edge[2] = $1; }
      # Get the acoustic score
      if($line =~ s/a=([\-\.\d]+)//) { $edge[3] = $1; }
      # Get the language model score
      if($line =~ s/l=([\-\.\d]+)//) { $edge[4] = $1; }

      # Get everything else
      $line =~ s/\s+/ /g;
      $edge[5] = $line;

      $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}) {
    print $fh sprintf("J=%-7d S=%-5d E=%-5d W=%-20s a=%7.7f l=%7.7f %s\n",
		      $L, @{$$edges[$L]});
  }
  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];
  my($i); my($ei);

  # $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][0]] + edge_likelihood($$edges[$ei]);
      if($test > $lp[$i]) { 
	$lp[$i] = $test;
	$bp[$i] = $$edges[$ei][0];
	$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"
  $l =~ s/.*LAUGH.*//;
  
  # These words seem to be used for various filled pauses
  #   Replace all with UH or UM
  $l =~ s/.*(EPENTHETIC|FILLEDSIL|FALSESTART|HESITATION|GLOTTAL).*//;
  $l =~ s/.*(TRUNCATEDP).*//;
  $l =~ s/.*(NASAL|MMM).*//;
  $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;
  
  # Return outword
  return($l);
}


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

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

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

#######################
# Main loop: process each pair of files in the script
while($script =~ s/^\s*(\S+)\s+(\S+)//) {
  my($inlat) = $1; my($outlat) = $2;
  unless(open(LATFILE, $inlat)) {
    print STDERR "Unable to read from lattice $inlat: $!\n";
    next;
  }
  my(@graph) = read_HTK_lattice(\*LATFILE);
  close(LATFILE);
  my(@deleted_edges) = ();
  my(@edge_midpoints) = ();
  my($original_nodecount) = $#{$graph[0]}+1;

  # DEBUG: write the lattice
  # print STDERR "Read the following lattice from $latfile:\n";
  # write_HTK_lattice(@graph, \*STDERR);
  
  # find maxprob path
  my(@maxprob) = viterbi(@graph[0..1]);
  
  # DEBUG: write out the edge words on MAP path
  print "MAP PATH: ";
  my($i);
  foreach $i (0..$#{$maxprob[1]}) { print ",",$graph[1][$maxprob[1][$i]][2]; }
  print "\n";

  # Add "n=1" to each edge in the MAP path
  foreach $i (0..$#{$maxprob[1]}) { $graph[1][$maxprob[1][$i]][5] .= " n=1"; }
  
  ################################################
  # First, compute all edge midpoint times, before we delete the extra nodes
  my($ei);
  foreach $ei (0..$#{$graph[1]}+1) { 
    $edge_midpoints[$ei] = ($graph[0][$graph[1][$ei][0]][2] + 
			    $graph[0][$graph[1][$ei][1]][2]) / 2;
  }

  ################################################
  # Re-arrange nodes, and delete most:
  # Move nodes in MAP path to the front of the nodelist; delete all others
  my($mli);
  foreach $mli (0..$#{$maxprob[0]}+1) {
    if($maxprob[0][$mli] != $mli)
      { @{$graph[0][$mli]} = @{$graph[0][$maxprob[0][$mli]]}; }

    # Now delete the forw and back edgelists from this MAP node
    @{$graph[0][$mli][0]} = ();
    @{$graph[0][$mli][1]} = ();
  }
  
  # Delete all nodes past $#$ml_nodes+1
  splice(@{$graph[0]}, $#{$maxprob[0]}+1);

  #################################################
  # Re-arrange edges:
  # Move or delete every edge in the graph
  foreach $ei (0..$#{$graph[1]}) {
    # Find the first end node (between 1 and $#{$maxprob[0]})
    #   such that midpoint time is before the time of the end node
    $n0 = 1; $t0 = $edge_midpoints[$ei];
    while(($n0 < $#{$maxprob[0]}) && ($graph[0][$n0][2] < $t0)) { $n0++; }

    # Check all edges already in the back list of node $n0
    # to find out whether any of them have the same word label
    foreach $fi (@{$graph[0][$n0][0]}) {
      # If edges ei and fi have the same word label, add their probabilities,
      #   mark edge ei for deletion, and move any n= label from ei to fi
      if($graph[1][$ei][2] eq $graph[1][$fi][2]) {
	$graph[1][$fi][4] += log(1+exp($graph[1][$ei][4]-$graph[1][$fi][4]));
	$graph[1][$fi][3] += log(1+exp($graph[1][$ei][3]-$graph[1][$fi][3]));
	$deleted_edges[$ei] = 1;
	if($graph[1][$ei][5] =~ /n=1/) { $graph[1][$fi][5] .= " n=1"; }
	break;
      }
    }
    # If no matching word label was found,
    unless($deleted_edges[$ei]) {
      # Give this node the new start time and end time
      $graph[1][$ei][0] = $n0-1;
      push(@{$graph[0][$n0-1][1]},$ei);
      $graph[1][$ei][1] = $n0;
      push(@{$graph[0][$n0][0]},$ei);
    }
  }

  # Create a completely new edgelist, in order of nodes
  #   This operation invalidates edgelists in the nodes.
  my(@new_edgelist) = ();
  my($ni) = 0;
  foreach $ni (0..($#{$graph[0]}-1)) { 
    # If there are edges connecting $ni to $ni+1, use them
    if($#{$graph[0][$ni][1]}>=0) {
      # First find the "n=1" edge, and move it to the start of this list
      foreach $i (0..$#{$graph[0][$ni][1]}) {
	if($graph[1][$graph[0][$ni][1][$i]][5] =~ /n=1/) {
	  my($tmp)=$graph[0][$ni][1][$i];
	  $graph[0][$ni][1][$i] = $graph[0][$ni][1][0];
	  $graph[0][$ni][1][0] = $tmp;
	}
      }
      # Now push these edges into the @new_edgelist
      foreach $ei (@{$graph[0][$ni][1]}) {
	push(@new_edgelist, [ @{$graph[1][$ei]} ]);
      }}
    else 
      # If not, invent a "sil" node to connect them
      { push(@new_edgelist, [ $ni, $ni+1, "sil", -20000, 0 ]); }
  }
  $graph[1] = \@new_edgelist;

  # Write the modified graph
  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 a summary to STDERR
  print STDERR sprintf("$inlat: (%d,%d) -> (%d,%d)\n",
		       $original_nodecount, $#deleted_edges+1, 
		       $#{$graph[0]}+1, $#{$graph[1]}+1);
}



