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

align_edges.pl: DP-Align alternate edges in a word graph

Usage:
 align_edges.pl script.scp

Align_edges.pl reads in HTK-format word graphs, aligns landmark-based
pronunciation models of all coterminous edgesets to the pronunciation
model of the maximum-posterior-probability edge. Each pronunciation
model is tagged with numbers describing the alignment.  A "coterminous
edgeset" is a set of edges all of which have the same starting node
and the same ending node.

Pronunciation models should be specified in the p= tag of the edge.
Pronunciation models should consist of landmarks tagged with
distinctive features.  All distinctive features are optional, except
that a landmark must specify either the feature syllabic or the
feature sonorant.

Alignment is computed using a minimum-feature-distance DP alignment,
with the following constraint: +syllabic landmarks (syllable nuclei)
may only align with other +syllabic landmarks.  All non-syllabic
landmarks may align with one another.

In the output, pronunciation models are tagged with the alignment
number of each landmark, e.g.

  2(+syllabic,-front,-high,+low,+CP,+stress)

The tag number ("2", in the example above) identifies the landmark
number in the maximum-posterior-probability edge that matches or most
recently precedes the current landmark.  If a landmark is to be
inserted before any of the landmarks in the MAP edge, it will be
tagged with a -1.

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;

# Deletion penalty for alignment of landmarks -- number of features
$LM_INSERTION_PENALTY=10;
$LM_DELETION_PENALTY=10;

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

###################################################
# 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);
}

###################################################
# Expand landmark-based transcription
#
# $aref = expand_transcription($string)
#
# $string contains colon-separated landmark specifications
# Each landmark specification is a set of comma-separated feature specs
# Each feature spec is a +, -, -+, or +- followed by the name of
#   the feature.
#
# $aref is an arrayref of hashrefs.
# $aref[$n] is the $nth expanded landmark specification
# $aref[$n]{'sonorant'} exists only if the feature 'sonorant'
#    was specified for the $nth landmark; if so, $aref[$n]{'sonorant'}
#    contains the value of the feature (-, +, -+, or +-).
#
sub expand_transcription {
  my(@lms) = split(/:/,$_[0]);
  my(@aref) = ();

  foreach $lm (@lms) {
    my(@feats) = split(/,/,$lm);
    my(%feattab) = ();
    
    # Value is the series of + and - specifiers at feature start,
    # Key is all remaining characters in the feature specification
    foreach $f (@feats) { 
       $f=~/([\+\-]+)(\S+)/; 
       $feattab{$2}=$1; 
     }
    
    push(@aref, \%feattab);
  }
  return(\@aref);
}
      
###################################################
# Compress landmark-based transcription
#
# $string = compress_transcription($aref)
#
# Opposite of expand_transcription
#
sub compress_transcription {
  my($string) = "";
  my(@aref) = @{$_[0]};
  # Set interpolator to comma
  $"=',';

  foreach $lm (@aref) {
    my(@feats) = ();

    # Print syllabic, sonorant, continuant first if they exist
    foreach $f ('syllabic', 'sonorant', 'continuant') {
      if(exists($$lm{$f})) { push(@feats, $$lm{$f}.$f) ; }
    }
    # Then print the other features
    foreach $f (keys %$lm) {
      if($f !~ /^(syllabic|sonorant|continuant)$/) {push(@feats,$$lm{$f}.$f);}
    }

    # Then concatenate them all to $string
    if(@feats > 0) { 
      if(length($string)>0) { $string .= ':' };
      $string .= "@feats";
    }
  }
  # Reset interpolator, and return
  $"=' ';
  return($string);
}

###################################################
# Compute feature distance between two feature tables
#
# $dist = feature_distance($t0, $t1)
#
# $t0 and $t1 are hashrefs
# $$t0{'sonorant'} is the value of feature 'sonorant' in 
#   hashref $t0 (either +, -, -+, +-, or nonexistent)
#
# $dist is a number: the distance between $t0 and $t1
#
# $dist counts the number of characters in either feature
#   table that are not matched by characters in the other
#   feature table.  Thus, for example, if 'feat' has
#   the following values in $t0 and $t1, $dist is incremented
#   by the amounts shown:
#
# $$t0{'feat'}    $$t1{'feat'}     $dist +=
#     +               +               0
#     +               unspecified     1
#     +-              -               1
#     +               -               2
#     +-              -+              2
#     +-              unspecified     2
#
sub feature_distance {
  my($t0) = $_[0];
  my($t1) = $_[1];
  my($dist) = 0;
  my(@k0) = keys %$t0;
  my(@k1) = keys %$t1;

  foreach $f (@k0) {
    # If one is a substring of the other, add length diff
    if (index($$t0{$f}, $$t1{$f})>=0) 
      { $dist += length($$t0{$f})-length($$t1{$f}); }
    elsif (index($$t1{$f}, $$t0{$f})>=0) 
      { $dist += length($$t1{$f})-length($$t0{$f}); }
    # If neither is a substring, add length of $$t0{$f}
    else { $dist += length($$t0{$f}); }
  }
  foreach $f (@k1) {
    # Increment $dist only if this feature hasn't yet been seen
    unless (exists($$t0{$f})) { $dist += length($$t1{$f}); }
  }

  # Special penalty guards against aligning any +syllabic landmark
  #    with any nonsyllabic landmark, if both landmarks are non-null
  if(((@k0 > 0) && (@k1 > 0)) &&
     ((($$t1{'syllabic'} eq '+') && ($$t0{'syllabic'} ne '+')) ||
      (($$t0{'syllabic'} eq '+') && ($$t1{'syllabic'} ne '+'))))
    { $dist += 1000000; }

  return($dist);
}

###################################################
# Compute the minimum-feature-distance alignment
# between one or more strings of landmarks
#
# $a = align_lms( $t0, $t1 )
#
# $t0 is an arrayref of landmarks
#  $$t0[$m] is a hashref
#   $$t0[$m]{'feat'} is the value of 'feat' for the $mth landmark
#
# $t1 is similarly defined.
#
# $a is an arrayref of the same size as $t0;
#   $$t1[$$a[$n]-1] aligns with $$t0[$n] in the minimum-distance alignment
#
sub align_lms {
  my($t0)= $_[0];
  my($t1)= $_[1];
  my(@a)=(); 
  my(@cost)=();
  my(@net)=();
  my($m); my($n);

  # Null alignment
  if (@$t0==0) { return(\@a); }

  # Compute substitution costs
  foreach $m (1..(1+$#$t0)) { foreach $n (1..(1+$#$t1)) {
    $cost[$m][$n] = feature_distance($$t0[$m-1], $$t1[$n-1]);
  }}

  # Next, compute deletion and insertion costs
  foreach $m (1..(1+$#$t0)) { 
    $cost[$m][0] = feature_distance($$t0[$m-1],{}) + $LM_DELETION_PENALTY; 
  }
  foreach $n (1..(1+$#$t1)) { 
    $cost[0][$n] = feature_distance($$t1[$n-1],{}) + $LM_INSERTION_PENALTY; 
  }

  # Next, initialize row zero and column zero of net array
  $net[0][0] = [ 0, 0, 0 ];
  foreach $m (1..(1+$#$t0)) 
    { $net[$m][0] = [ $net[$m-1][0]+$cost[$m][0], $m-1, 0 ]; }
  foreach $n (1..(1+$#$t1)) 
    { $net[0][$n] = [ $net[0][$n-1]+$cost[0][$n], 0, $n-1]; }

  # Now compute DP cost to each ($m,$n) pair
  foreach $m (1..(1+$#$t0)) { foreach $n (1..(1+$#$t1)) {
    $sc=$cost[$m][$n] + $net[$m-1][$n-1][0];
    $dc=$cost[$m][0] + $net[$m-1][$n][0];
    $ic=$cost[0][$n] + $net[$m][$n-1][0];
    if($sc <= $dc && $sc <= $ic) { $net[$m][$n] = [ $sc, $m-1, $n-1 ]; }
    elsif($ic <= $dc) { $net[$m][$n] = [ $ic, $m, $n-1 ]; }
    else { $net[$m][$n] = [ $dc, $m-1, $n ]; }
  }}

  # Accumulate the backpointers, working backward through @net
  my(@bpm) = ( 1+$#$t0 );
  my(@bpn) = ( 1+$#$t1 );
  # Continue until we get back past all nonzero $m's
  while ($bpm[$#bpm] > 0) {
    $n=$#bpm;
    push(@bpm, $net[$bpm[$n]][$bpn[$n]][1]);
    push(@bpn, $net[$bpm[$n]][$bpn[$n]][2]);
  }
  # Assign to $a[$m] the last $n matching $m in a backpointer
  foreach $m (0..($#bpm-1)) { $a[$bpm[$m]-1] = $bpn[$m]; }
  return(\@a);
}

#######################
#
# align_edgeset(@edgeset)
#
# Input @edgeset is assumed to be a coterminous edgeset.
# Pronunciation models from each edgeset are extracted,
# and aligned using align_landmarks.
# The resulting pronunciations are compressed back to the 'p'
# tags in the edges in @edgeset.
#
sub align_edgeset {
  my(@edgeset) = ( @_ );
  my(@transcriptions) = ();
  my($maxlogprobi) = -1;
  my($logprob) = 0;
  my($m); my($n); my($a); my($t);

  # Find the edge that was on the ML path
  foreach $i (0..$#edgeset) {
    if($edgeset[$i]{'n'} == 1) { $maxlogprobi = $i; }
  }

  # Expand the maxlogprob transcription
  my($tmaxlp) = expand_transcription($edgeset[$maxlogprobi]{'p'});

  # Align the transcriptions
  foreach $m (0..$#edgeset) {
    # Expand the $m'th transcription
    $t = expand_transcription($edgeset[$m]{'p'});
    # Compute its alignment to $tmaxlp
    $a = align_lms($t, $tmaxlp);

    # Split the landmarks again, ...
    my(@lms) = split(/:/,$edgeset[$m]{'p'});

    # ... wrap them in alignment(), and append them back to the transcription
    $edgeset[$m]{'p'} = $$a[0].'('.$lms[0].')';
    foreach $n (1..$#lms) {
      $edgeset[$m]{'p'} .= ':'.$$a[$n].'('.$lms[$n].')'; 
    }
  }

  return(@edgeset);
}
  
################################################################

# 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+)//) {
  $inlat = $1; $outlat = $2;

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

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

  # Process each potential starting node
  foreach $startnum (0..$#{$graph[0]}) {
    @edgesets = ();

    # Sort, by ending node, all edges that start with $startnum
    foreach $edgenum (@{$graph[0][$startnum][1]}) {
      $endnum = $graph[1][$edgenum]{'E'};
      push(@{$edgesets[$endnum]}, $edgenum);
    }

    # Process each edgeset
    foreach $edgeset (@edgesets) {
      # If edgeset is larger than one element, align
      if(@{$edgeset} > 1) {
	@{$graph[1]}[@$edgeset] = align_edgeset(@{$graph[1]}[@$edgeset]); 
      }
    }
  }
  
  # 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);

}



