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

lat2praat.pl: Extract a Praat transcription file from an HTK-format lattice

Usage: lat2praat.pl SCP_file maxdepth
 SCP_file: filenames, ordered in pairs: lattice_file praat_file

Each line in SCP_file is assumed to contain one input filename and one
output filename.  The input file is assumed to be an HTK-format
lattice file for one utterance.  The output file is created as a Praat
TextGrid file.  

The TextGrid file contains up to "maxdepth" tiers (default: 8).  Edges
in the lattice are processed in numerical order.  Each edge defines a
Praat segment that begins at the time of the start node, ends at the
time of the end node, and is labeled with the nodes word label.  Each
edge is written to the first label tier to which it can be written
without over-writing any existing segment.  If there is no such tier,
the edge is omitted from the output.

ENDUSAGE

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

# Maximum depth of the TextGrid file (default: 8)
$MAXDEPTH=8;

# Math functions
sub abs { my($x)=shift; ($x>0) ? $x : -$x ; }

###################################################
# 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 a TextGrid header
#
# Syntax: write_textgrid_header($fh, $firsttime, $lastttime, $maxdepth)
#
sub write_textgrid_header {
  my($fh) = $_[0];
  my($firstttime) = $_[1];
  my($lasttime) = $_[2];
  my($maxdepth) = $_[3];

  # Print the TEXTGRID header
  print $fh 'File type = "ooTextFile"',"\n";
  print $fh 'Object class = "TextGrid"',"\n\n";
  print $fh "xmin = $firstttime\n";
  print $fh "xmax = $lasttime\n";
  print $fh 'tiers? <exists>',"\n";
  print $fh "size = $maxdepth\n";
  print $fh 'item []: ',"\n";
}

###################################################
# Print a TextGrid tier
#
# Syntax: sprint_textgrid_tier(\@data, $number, $starttime, $endtime)
#
sub sprint_textgrid_tier {
  # Use a private version of phndata array,
  # so we can insert elements w/o global consequences
  my(@phndata) = @{$_[0]};  
  my($number) = $_[1];
  my($starttime) = $_[2];
  my($endtime) = $_[3];

  # If tier is null, create a one-segment tier
  unless(@phndata) { @phndata = ([$starttime, $endtime, ""]); }
    
  # First: insert segments so that segments are collectively exhaustive
  # This must be done first, so that array size is known in advance
  if($phndata[$#phndata][1] < $endtime) 
    { push(@phndata,[$phndata[$#phndata][1],$endtime,""]); }
  for($nt=$#phndata; $nt>0; $nt--) {
    if($phndata[$nt-1][1] < $phndata[$nt][0]) 
      { splice(@phndata,$nt,0,[$phndata[$nt-1][1],$phndata[$nt][0],"" ]) }
  }
  if($phndata[0][0] > $starttime) 
    { unshift(@phndata,[$starttime,$phndata[0][0],""]); }

  # Start with "item [$number]"
  my($output) = "    item [$number]:\n";
  $output .= "        class = \"IntervalTier\"\n"; 
  $output .= "        name = \"tier${number}\"\n"; 
  $output .= "        xmin = $starttime\n";
  $output .= "        xmax = $endtime\n";
  $output .= sprintf("        intervals: size = %d\n", $#phndata + 1);

  # Print out the converted data 
  foreach $nt (0..$#phndata) {
    $output .= sprintf("        intervals [%d]:\n",$nt+1);
    $output .= "            xmin = $phndata[$nt][0]\n";
    $output .= "            xmax = $phndata[$nt][1]\n";
    $output .= "            text = \"$phndata[$nt][2]\"\n";
  }
  return($output);
}

################################################################
# Extract tiers from a lattice file
#
# @tiers = extract_tiers_from_lattice($graph, $maxdepth)
#
# $graph -- reference to an array of the type returned by read_HTK_lattice
#
# $maxdepth -- maximum depth of the output transcription file
#
# Assume that the node times are sorted, so that we can just 
#  seek open slots in order, without searching all processed edges
#  in order to seek an open slot for each new edge
#
sub extract_tiers_from_lattice {
  my(@graph) = @{$_[0]};
  my($maxdepth) = $_[1];
  my(@tiers) = ();
  my($i); my($m);

  # Initialize @tiers with empty arrays
  foreach $i (0..($maxdepth-1)) { $tiers[$i] = []; }

  # Process each possible start node, in order
  foreach $m (0..$#{$graph[0]}) {
    # Time of current node
    my($node_time) = $graph[0][$m][2];

    # Create an array containing all of the edge hashrefs
    #   that start from the current node
    my(@edges)=@{$graph[1]}[@{$graph[0][$m][1]}];

    # Find the MAP edge, if any, and put it first
    #  If each MAP edge starts when the previous ended,
    #  this method guarantees that MAP edges will be on first tier
    foreach $i (0..$#edges) {
      if($edges[$i]{'n'} eq "1") {
	# Put the edge first
	unshift(@edges, $edges[$i]);
	# Delete it from its previous spot
	splice(@edges, $i+1, 1);
      }
    }

    # Try to insert edges into all tiers
    my($t)=-1;
    while ($t < $#tiers && @edges > 0) {
      $t++;

      # If tier has no segments, insert the edge here
      if(@{$tiers[$t]}==0) {
	push(@{$tiers[$t]}, [ $graph[0][$edges[0]{'S'}][2], 
			      $graph[0][$edges[0]{'E'}][2], $edges[0]{'W'} ]); 
	# Remove the edge from the list of available edges
	shift(@edges);
      }

      # If last segment on this tier ends before current node time, 
      #   insert the edge here.  This assumes that
      #   all previously inserted segments start
      #   before the current node time, and that segments
      #   have been inserted in order of increasing start time.
      elsif($tiers[$t][$#{$tiers[$t]}][1] <= $node_time) {
	push(@{$tiers[$t]}, [ $graph[0][$edges[0]{'S'}][2], 
			      $graph[0][$edges[0]{'E'}][2], $edges[0]{'W'} ]); 
	# Remove the edge from the list of available edges
	shift(@edges);
      }
    }
  }

  return(@tiers);
}
    

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

if (@ARGV==0) { print STDERR $USAGE; exit; }
if (@ARGV>1) { $MAXDEPTH=$ARGV[1]; }
 
# 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);

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

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

  # Assume that nodes are sorted
  my($firsttime) = $graph[0][0][2];
  my($lasttime) = $graph[0][$#{$graph[0]}][2];

  # Open the TEXTGRID file and write its header
  open(TEXTGRID, ">$outfile") || die "Unable to write to $outfile: $!";
  write_textgrid_header(\*TEXTGRID, $firsttime, $lasttime, $MAXDEPTH);

  # Extract the tiers
  my(@tiers) = extract_tiers_from_lattice(\@graph, $MAXDEPTH);

  # Write them to textgrid file
  my($m);
  foreach $m (0..$#tiers) {
    print TEXTGRID sprint_textgrid_tier($tiers[$m], $m, $firsttime, $lasttime)
  }

  # Print summary to STDERR
  print STDERR "$inlat -> $outfile\n";
}


