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

prune_lattice.pl: Prune edges incommensurate with known subword transcription

Usage
  prune_lattice.pl N MLF DICT < script

Example
  prune_lattice.pl 3 ref_lms.mlf pronlex_lw.dict < prune_lw.scp
  prune_lattice.pl 3 ref_phones.mlf pronlex.dict < prune_pw.scp

prune_lattice.pl removes, from input lattices, any edges whose best
pronunciation differs from the reference subword MLF by more than N
units.  Each edge in each input lattice is expanded according to every
possible dictionary entry in DICT.  Each such pronunciation is
compared to the transcription of the same time period, as provided in
MLF.  If the Levenshtein distance between a pronunciation and the
reference is less than N, the edge is retained; otherwise the edge is
deleted.

The script file should contain input & output filenames in pairs,
i.e., "input1.lat output1.lat input2.lat output2.lat ..."

ENDUSAGE

####################################################################

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

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

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

 
###################################################
# Subroutine to read MLF entry
#
# Syntax: $data = read_MLF_entry("filename", \*INPUTFILE);
#
# Returned value $data is an arrayref of arrayrefs
# $$tran[0][0] is first seg start time, $$tran[0][1] is end time,
# @$$tran[0][2..$#{$$tran[0]}] is the sequence of labels at first time.
#
sub read_MLF_entry {
  my($filename) = $_[0];
  my($fh) = $_[1];
  my(@tran) = ();
  my(@labels) = ();
  
  # Read all continuous lines that begin with a char other than period
  for($l=<$fh>; $l =~ /^[^\.]/; $l=<$fh>) {
    # Option 1: line is in "time time labels" format
    if ($l=~/(\d+)\s+(\d+)\s+(.*)/ ) {
      @labels = split(/\s+/,$3);
      push(@tran, [ $1 * $MLFPERIOD, $2 * $MLFPERIOD, @labels ] );
    }
    # Option 2: line contains any non-whitespace labels, but no times
    elsif ($l =~ /\S/) {
      @labels = split(/\s+/,$l);
      push(@tran, [ 0, 0, @labels ] );
    }
  }
  return(\@tran);
}

###################################################
# Subroutine to read a noway-format (or HTK-format) dictionary
#
# Syntax: $dict = read_dict(\*FILEHANDLE)
#
# $dict = hashref of arrayrefs of arrayrefs
#   @{$$dict{'word'}[0]} is the first pronunciation of "word", e.g., 
#   $$dict{'word'}[0] = [ 0.35 w er d ];
#   where the first entry is the probability of this pronunciation.
#
# If entries have no attached probabilities (e.g., HTK dictionaries),
# all probabilities in dict are set to 1.0.
#
sub read_dict {
  my($fh) = $_[0];
  my(%dict) = ();
  my($word) = '';

  # Read every line, ignoring blank lines
  while(($line=<$fh>) !~ /^$/) {
    chomp($line); chomp($line);
    my(@pron) = split(/\s+/,$line);

    # Case 1: first field is word(prob)
    if($pron[0] =~ /(\S+)\(([\d\.]+)\)/) {
      $word = $1;
      $pron[0] = $2;
    }
    # Case 2: first field is just word; set prob to 1.0
    else {
      $word = $pron[0];
      $pron[0] = 1;
    }
    if(exists($dict{$word})) { push(@{$dict{$word}}, \@pron); }
    else { $dict{$word} = [ \@pron ]; }
  }
  return(\%dict);
}

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

###################################################
# Extract word string from an MLF-style transcription
#
# $wrds = trans2wrds($trans)
#
# $wrds: arrayref of wrds
# $trans: arrayref of transcriptions
#
sub trans2wrds {
  my(@trans) = @{$_[0]};
  my(@wrds) = ();
  foreach $t (@trans) { push(@wrds, $$t[2]); }
  return(\@wrds);
}

###################################################
# Find Levenshtein statistics between two symbol strings
#
# $stats = levenshtein_skipN( $ref, $test, $skippable )
#
# $stats = [ N, S, D, I, E=S+D+I ]
# $ref = reference string
# $test = test string
# $skippable = number of words that can be skipped without loss
#      at start or end of $ref
#
# N counts the number of non-null wrds in $ref.
# S, D, I = number of subs, dels, ins, and total errors
#   in the best alignment of $ref with $test.
#   Insertions and deletions do not count null wrds ('').
#   E=S+D+I.
#
sub levenshtein_skipN {
  my($ref) = $_[0];
  my($test) = $_[1];
  my($skippable) = $_[2];
  my(@stats) = ();
  my($m); my($n);
  my($numref)=0; 
  my($numtest)=0; 

  # If zero wrds observed in test list, (SDIE) = (0,D,0,D)
  foreach $m (0..($#{$ref}+1)) { 
    if( $m>0 && $$ref[$m-1] ne '') { $numref++ };
    # Allow skipping of the first $skippable words in $ref
    $deletions = ($numref>$skippable)?$numref-$skippable:0;
    $stats[$m][0] = [0,$deletions,0,$deletions]; 
  }

  # If zero wrds observed in ref list, (SDIE) = (0,I,I)
  foreach $n (0..($#{$test}+1)) { 
    if( $n>0 && $$test[$n-1] ne '') { $numtest++ };
    $stats[$m][0] = [0,0,$numtest,$numtest]; 
  }

  # m is index into the ref list, n is index into the test list    
  foreach $m (1..($#{$ref}+1)) { foreach $n (1..($#{$test}+1)) {

    # If wrds are equal, there is no cost; 
    #   Otherwise, we suffer at least one substitution
    @{$stats[$m][$n]} = @{$stats[$m-1][$n-1]};
    unless($$ref[$m-1] eq $$test[$n-1]) 
      { $stats[$m][$n][0]++; $stats[$m][$n][3]++; }
    
    # If $ref[$m-1] eq '', we can skip it without cost
    if ($$ref[$m-1] eq '' && $stats[$m-1][$n][3] < $stats[$m][$n][3]) 
      { @{$stats[$m][$n]} = @{$stats[$m-1][$n]}; }
    # Otherwise, if deletion results in minimum cost, use it
    elsif ($stats[$m-1][$n][3] < $stats[$m][$n][3] - 1) {
      @{$stats[$m][$n]} = @{$stats[$m-1][$n]};
      $stats[$m][$n][1]++; $stats[$m][$n][3]++;
    }
    
    # If $test[$n-1] eq '', we can skip it without cost
    if ($$test[$n-1] eq '' && $stats[$m][$n-1][3] < $stats[$m][$n][3]) 
      { @{$stats[$m][$n]} = @{$stats[$m][$n-1]}; }
    # Otherwise, if insertion results in lower cost, use it
    elsif ($stats[$m][$n-1][3] < $stats[$m][$n][3] - 1) {
      @{$stats[$m][$n]} = @{$stats[$m][$n-1]};
      $stats[$m][$n][2]++; $stats[$m][$n][3]++;
    }
  }}

  # Allow skipping of the last $skippable words in $ref
  my($lastword) = @$ref;
  foreach $i (1..$skippable) {
    if($stats[@$ref-$i][$#$test+1][3] < $stats[@$ref][$#$test+1][3])
      { $lastword = @$ref-$i; }
  }
  return( $numref, @{$stats[$lastword][$#{$test}+1]} );
}


###################################################
# Find Levenshtein statistics between two symbol strings
#
# $stats = levenshtein( $ref, $test )
#
# $stats = [ N, S, D, I, E=S+D+I ]
# $ref = reference string
# $test = test string
#
# N counts the number of non-null wrds in $ref.
# S, D, I = number of subs, dels, ins, and total errors
#   in the best alignment of $ref with $test.
#   Insertions and deletions do not count null wrds ('').
#   E=S+D+I.
#
sub levenshtein {
  my($ref) = $_[0];
  my($test) = $_[1];
  my(@stats) = ();
  my($m); my($n);
  my($numref)=0; 
  my($numtest)=0; 

  # If zero wrds observed in test list, (SDIE) = (0,D,0,D)
  foreach $m (0..($#{$ref}+1)) { 
    if( $m>0 && $$ref[$m-1] ne '') { $numref++ };
    $stats[$m][0] = [0,$numref,0,$numref]; 
  }

  # If zero wrds observed in ref list, (SDIE) = (0,I,I)
  foreach $n (0..($#{$test}+1)) { 
    if( $n>0 && $$test[$n-1] ne '') { $numtest++ };
    $stats[$m][0] = [0,0,$numtest,$numtest]; 
  }

  # m is index into the ref list, n is index into the test list    
  foreach $m (1..($#{$ref}+1)) { foreach $n (1..($#{$test}+1)) {

    # If wrds are equal, there is no cost; 
    #   Otherwise, we suffer at least one substitution
    @{$stats[$m][$n]} = @{$stats[$m-1][$n-1]};
    unless($$ref[$m-1] eq $$test[$n-1]) 
      { $stats[$m][$n][0]++; $stats[$m][$n][3]++; }
    
    # If $ref[$m-1] eq '', we can skip it without cost
    if ($$ref[$m-1] eq '' && $stats[$m-1][$n][3] < $stats[$m][$n][3]) 
      { @{$stats[$m][$n]} = @{$stats[$m-1][$n]}; }
    # Otherwise, if deletion results in minimum cost, use it
    elsif ($stats[$m-1][$n][3] < $stats[$m][$n][3] - 1) {
      @{$stats[$m][$n]} = @{$stats[$m-1][$n]};
      $stats[$m][$n][1]++; $stats[$m][$n][3]++;
    }
    
    # If $test[$n-1] eq '', we can skip it without cost
    if ($$test[$n-1] eq '' && $stats[$m][$n-1][3] < $stats[$m][$n][3]) 
      { @{$stats[$m][$n]} = @{$stats[$m][$n-1]}; }
    # Otherwise, if insertion results in lower cost, use it
    elsif ($stats[$m][$n-1][3] < $stats[$m][$n][3] - 1) {
      @{$stats[$m][$n]} = @{$stats[$m][$n-1]};
      $stats[$m][$n][2]++; $stats[$m][$n][3]++;
    }
  }}

  return( $numref, @{$stats[$#{$ref}+1][$#{$test}+1]} );
}

###################################################
# 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 noise and laughter are
# replace with a null string "".
#
# This function differs from the one used in score_lattice in that
# filled pauses are not eliminated.
#
# $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|HESITATION|GLOTTAL).*/UH/;
  $l =~ s/.*(TRUNCATEDP|FALSESTART).*//;
  $l =~ s/.*(NASAL|MMM).*/UM/;
  
  # 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; }

# Get the Levenshtein distance threshold
$LEVTHRESH = $ARGV[0];

# Read in all lines from the reference MLF
%IMLF  = ();
open(MLF_FILE, $ARGV[1]) || die "Unable to open MLF $ARGV[1]: $!";
while (<MLF_FILE>) {
  # If this line starts a transcription, read it
  if (/\".*\/([^\/\.]+)[^\/]*\"/) {
    my($filename) = $1;
    $IMLF{$filename} = read_MLF_entry($filename,\*MLF_FILE);
  }
}
close(MLF_FILE);

# Read in the dictionary
open(DICT_FILE, $ARGV[2]) || die "Unable to open DICT $ARGV[2]: $!";
$DICT = read_dict(\*DICT_FILE);
close(DICT_FILE);

# Read in the script file
$script = ''; while (<STDIN>) { $script .= $_; }

#######################
# Main loop: process each pair of files in the script
while($script =~ s/^\s*(\S+)\s+(\S+)//) {
  $inlat = $1; $outlat = $2;
  unless(open(LATFILE, $inlat)) {
    print STDERR "Unable to read from lattice $inlat: $!";
    next;
  }
  my(@graph) = read_HTK_lattice(\*LATFILE);
  close(LATFILE);
  my(@deleted_edges) = ();
  
  # Get filename from $header{'UTTERANCE'}; strip extensions & dirs
  $filename = $graph[2]{'UTTERANCE'}; 
  $filename =~ s/(\S+\/)|(\.\S+)//g;
  unless(exists($IMLF{$filename})) {
    print STDERR "$filename not found in $ARGV[1]\n";
    next;
  }
  $trans = $IMLF{$filename};
  $words = trans2wrds($trans);

  # Evaluate every edge in the lattice
  foreach $ei (0..$#{$graph[1]}) {
    # Get the word from this edge
    $edgeword = normalize_word($graph[1][$ei][2]);
    # Don't try to align this edge if it was nullified by normalize_word
    next unless($edgeword);

    # Find times of the nodes
    unless(exists($graph[0][$graph[1][$ei][0]][2])) {
      print STDERR "Node times not specified in lattice $inlat\n";
      next;
    }
    $t0 = $graph[0][$graph[1][$ei][0]][2];
    $t1 = $graph[0][$graph[1][$ei][1]][2];
    
    # Extract transcription between $t0 and $t1
    # Extract one extra, on each end, in case of lattice misalignment
    $n0 = 0; while($n0 <= $#$trans && $$trans[$n0][1] < $t0) { $n0++; }
    $n1 = $#$trans; while($n1 >= 0 && $$trans[$n1][0] > $t1) { $n1--; }
    @subwords = @$words[($n0>0?$n0-1:$n0)..($n1<$#$words?$n1+1:$n1)];

    # Search available pronunciations until we find one
    #  with Levenshtein distance not greater than $LEVTHRESH
    #  allow skip of 2 words at start or end of transcription
    $minlev = @subwords;
    foreach $pron (@{$$DICT{$edgeword}}) {
      # Extract the part which is not a probability
      @test_pron = @{$pron}[1..@$pron];
      $stats = levenshtein_skipN(\@subwords, \@test_pron, 2);
      $minlev = $$stats[4] if ($minlev < $$stats[4]);
    }

    # If minlev still larger than LEVTHRESH, mark this edge for deletion
    if($minlev > $LEVTHRESH) { $deleted_edges[$ei] = 1; }
  }

  # Retain only the nodes that can be reached forward from node 0
  my(@retained_nodes) = ( 1 );
  foreach $i (0..$#{$graph[0]}) { if($retained_nodes[$i]) { 
    foreach $j (@{$graph[0][$i][1]}) { unless($deleted_edges[$j]) { 
      $retained_nodes[$graph[1][$j][1]] = 1; 
    }}
  }}
  # For all non-retained nodes: delete its edges, splice out the node
  for($i=$#{$graph[0]}; $i >= 0; $i--) { unless($retained_nodes[$i]) {
    foreach $j (@{$graph[0][$i][0]}, @{$graph[0][$i][1]}) 
      { $deleted_edges[$j]=1; }
    splice(@{$graph[0]}, $i, 1);
  }}
  # Update the node numbers of all edges touching each retained node
  foreach $i (0..$#{$graph[0]}) {
    foreach $j (@{$graph[0][$i][0]}) { $graph[1][$j][1] = $i; }
    foreach $j (@{$graph[0][$i][1]}) { $graph[1][$j][0] = $i; }
  }
  # Splice out all deleted edges
  #   This operation invalidates edgelists in the nodes.
  for($j=$#{$graph[1]}; $j >= 0; $j--) { if($deleted_edges[$j]) {
    splice(@{$graph[1]}, $j, 1); 
  }}

  # Write the modified graph
  unless(open(LAT_OUT, ">$outlat")) {
    print STDERR "Unable to write lattice out to $outlat: $!";
    next;
  }
  write_HTK_lattice(@graph, \*LAT_OUT);
  close(LAT_OUT);

  # Write a summary to STDERR
  print STDERR sprintf("$inlat: (%d,%d) -> (%d,%d)\n",
		       $#retained_nodes+1, $#deleted_edges+1, 
		       $#{$graph[0]}+1, $#{$graph[1]}+1);
}


