#!/usr/bin/env perl
  use warnings; #sed replacement for -w perl parameter
  
  # Copyright 2012  Arnab Ghoshal
  
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #  http://www.apache.org/licenses/LICENSE-2.0
  #
  # THIS CODE IS PROVIDED *AS IS* BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  # KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION ANY IMPLIED
  # WARRANTIES OR CONDITIONS OF TITLE, FITNESS FOR A PARTICULAR PURPOSE,
  # MERCHANTABLITY OR NON-INFRINGEMENT.
  # See the Apache 2 License for the specific language governing permissions and
  # limitations under the License.
  
  
  # This script builds a lower-order LM.
  
  my $usage = "Usage: subset_lm.pl -i input[.gz] -n order -o output[.gz]
   Builds a lower-oder LM from the input (can read & write gzipped files).
  ";
  
  use strict;
  use Getopt::Long;
  my ($in_lm, $order, $out_lm, @ngram_counts);
  GetOptions ("i=s" => \$in_lm,   # Input language model file
  	    "n=i" => \$order,   # Max n-gram order of the output
              "o=s" => \$out_lm); # Output language model file
  
  die "$usage" unless (defined($in_lm) && defined($order) && defined($out_lm));
  die "Output order must be positive." unless ($order > 0);
  
  my $cmd = ($in_lm =~ /\.gz$/)? "gzip -dc $in_lm |" : "<$in_lm";
  open(IN, "$cmd") or die "Cannot read from language model file '$in_lm': $!";
  $cmd = ($out_lm =~ /\.gz$/)? "| gzip -c > $out_lm" : ">$out_lm";
  open(OUT, "$cmd") or die "Cannot write to language model file '$out_lm': $!";
  
  while (<IN>) { last if /^\\data\\/; }  # Skip till header
  
  # Read counts for various n-gram orders
  while (<IN>) {
    next if /^$/;
    last if /^\\1-grams\:/;
    m/^ngram (\d+)\s*\=\s*(\d+)$/ or die "Bad line: $_";
    $ngram_counts[$1] = $2;  # Not 0-indexed!
  }
  
  # Handle degenerate case.
  if ($#ngram_counts <= $order) {
    warn "Input LM order ($#ngram_counts) <= required LM order ($order): input and output will be identical.";
    seek(IN, 0, 0);
    while (<IN>) { print OUT; }  # Don't have to worry about gz 
    close IN; close OUT;
    exit 0;
  } else {  # Print the header (with lower n-gram order)
    print OUT "\\data\\
  ";
    for my $i (1..$order) {
      print OUT "ngram $i=$ngram_counts[$i]
  ";
    }
    print OUT "
  ";  
  }
  
  my $nproc;
  for my $i (1..$order) {
    $nproc = 0;  # Number of n-grams actually found.
    print OUT "
  \\${i}-grams:
  ";
    while (<IN>) {
      next if ($_ =~ /^$/);
      if ($_ =~ /^\\(\d+)\-grams\:/) {
        die "Expecting $ngram_counts[$i] $i-grams; found $nproc!" 
  	unless ($nproc == $ngram_counts[$i]);
        last; 
      }
      my @fields = split;
      die "Bad line: $_" if ($#fields<$i || $#fields>$i+1);
      # For the highest order of n-gram, remove back-off weight, if present.
      pop(@fields) if ($i==$order && $#fields==$i+1);
      print OUT join(" ", @fields), "
  ";
      $nproc += 1;
    }
  }
  
  print OUT "\\end\\
  ";
  close IN;
  close OUT;