#!/usr/bin/env perl

# Copyright 2014  Johns Hopkins University (Author: Sanjeev Khudanpur), Vimal Manohar
# Apache 2.0

################################################################################
#
# This script was written to check the goodness of automatic segmentation tools
# It assumes input in the form of two Kaldi segments files, i.e. a file each of
# whose lines contain four space-separated values:
#
#    UtteranceID  FileID  StartTime EndTime
#
# It computes # missed frames, # false positives and # overlapping frames.
#
################################################################################

if ($#ARGV == 1) {
  $ReferenceSegmentation = $ARGV[0];
  $HypothesizedSegmentation = $ARGV[1];
  printf STDERR ("Comparing reference segmentation\n\t%s\nwith proposed segmentation\n\t%s\n",
    $ReferenceSegmentation,
    $HypothesizedSegmentation);
} else {
  printf STDERR "This program compares the reference segmenation with the proposted segmentation\n";
  printf STDERR "Usage: $0 reference_segments_filename proposed_segments_filename\n";
  printf STDERR "e.g. $0 data/dev10h/segments data/dev10h.seg/segments\n";
  exit (0);
}

################################################################################
# First read the reference segmentation, and
# store the start- and end-times of all segments in each file.
################################################################################

open (SEGMENTS, "cat $ReferenceSegmentation | sort -k2,2 -k3n,3 -k4n,4 |")
  || die "Unable to open $ReferenceSegmentation";
$numLines = 0;
while ($line=<SEGMENTS>) {
  chomp $line;
  @field = split("[ \t]+", $line);
  unless ($#field == 3) {
    exit (1);
    printf STDERR "Skipping unparseable line in file $ReferenceSegmentation\n\t$line\n";
    next;
  }
  $fileID = $field[1];
  unless (exists $firstSeg{$fileID}) {
    $firstSeg{$fileID} = $numLines;
    $actualSpeech{$fileID} = 0.0;
    $hypothesizedSpeech{$fileID} = 0.0;
    $foundSpeech{$fileID} = 0.0;
    $falseAlarm{$fileID} = 0.0;
    $minStartTime{$fileID} = 0.0;
    $maxEndTime{$fileID} = 0.0;
  }
  $refSegName[$numLines] = $field[0];
  $refSegStart[$numLines] = $field[2];
  $refSegEnd[$numLines] = $field[3];
  $actualSpeech{$fileID} += ($field[3]-$field[2]);
  $minStartTime{$fileID} = $field[2] if ($minStartTime{$fileID}>$field[2]);
  $maxEndTime{$fileID} = $field[3] if ($maxEndTime{$fileID}<$field[3]);
  $lastSeg{$fileID} = $numLines;
  ++$numLines;
}
close(SEGMENTS);
print STDERR "Read $numLines segments from $ReferenceSegmentation\n";

################################################################################
# Process hypothesized segments sequentially, and gather speech/nonspeech stats
################################################################################

open (SEGMENTS, "cat $HypothesizedSegmentation | sort -k2,2 -k1,1 |")
  # Kaldi segments files are sorted by UtteranceID, but we re-sort them here
  # so that all segments of a file are read together, sorted by start-time.
  || die "Unable to open $HypothesizedSegmentation";
$numLines = 0;
$totalHypSpeech = 0.0;
$totalFoundSpeech = 0.0;
$totalFalseAlarm = 0.0;
$numShortSegs = 0;
$numLongSegs = 0;
while ($line=<SEGMENTS>) {
  chomp $line;
  @field = split("[ \t]+", $line);
  unless ($#field == 3) {
    exit (1);
    printf STDERR "Skipping unparseable line in file $HypothesizedSegmentation\n\t$line\n";
    next;
  }
  $fileID = $field[1];
  $segStart = $field[2];
  $segEnd = $field[3];
  if (exists $firstSeg{$fileID}) {
    # This FileID exists in the reference segmentation
    # So gather statistics for this UtteranceID
    $hypothesizedSpeech{$fileID} += ($segEnd-$segStart);
    $totalHypSpeech += ($segEnd-$segStart);
    if (($segStart>=$maxEndTime{$fileID}) || ($segEnd<=$minStartTime{$fileID})) {
      # This entire segment is a false alarm
      $falseAlarm{$fileID} += ($segEnd-$segStart);
      $totalFalseAlarm += ($segEnd-$segStart);
    } else {
      # This segment may overlap one or more reference segments
      $p = $firstSeg{$fileID};
      while ($refSegEnd[$p]<=$segStart) {
        ++$p;
      }
      # The overlap, if any, begins at the reference segment p
      $q = $lastSeg{$fileID};
      while ($refSegStart[$q]>=$segEnd) {
        --$q;
      }
      # The overlap, if any, ends at the reference segment q
      if ($q<$p) {
        # This segment sits entirely in the nonspeech region
        # between the two reference speech segments q and p
        $falseAlarm{$fileID} += ($segEnd-$segStart);
        $totalFalseAlarm += ($segEnd-$segStart);
      } else {
        if (($segEnd-$segStart)<0.20) {
          # For diagnosing Pascal's VAD segmentation
          print STDOUT "Found short speech region $line\n";
          ++$numShortSegs;
        } elsif (($segEnd-$segStart)>60.0) {
          ++$numLongSegs;
          # For diagnosing Pascal's VAD segmentation
          print STDOUT "Found long speech region $line\n";
        }
        # There is some overlap with segments p through q
        for ($s=$p; $s<=$q; ++$s) {
          if ($segStart<$refSegStart[$s]) {
            # There is a leading false alarm portion before s
            $falseAlarm{$fileID} += ($refSegStart[$s]-$segStart);
            $totalFalseAlarm += ($refSegStart[$s]-$segStart);
            $segStart=$refSegStart[$s];
          }
          $speechPortion = ($refSegEnd[$s]<$segEnd) ?
          ($refSegEnd[$s]-$segStart) : ($segEnd-$segStart);
          $foundSpeech{$fileID} += $speechPortion;
          $totalFoundSpeech += $speechPortion;
          $segStart=$refSegEnd[$s];
        }
        if ($segEnd>$segStart) {
          # There is a trailing false alarm portion after q
          $falseAlarm{$fileID} += ($segEnd-$segStart);
          $totalFalseAlarm += ($segEnd-$segStart);
        }
      }
    }
  } else {
    # This FileID does not exist in the reference segmentation
    # So all this speech counts as a false alarm
    exit (1);
    printf STDERR ("Unexpected fileID in hypothesized segments: %s", $fileID);
    $totalFalseAlarm += ($segEnd-$segStart);
  }
  ++$numLines;
}
close(SEGMENTS);
print STDERR "Read $numLines segments from $HypothesizedSegmentation\n";

################################################################################
# Now that all hypothesized segments have been processed, compute needed stats
################################################################################

$totalActualSpeech = 0.0;
$totalNonSpeechEst = 0.0; # This is just a crude estimate of total nonspeech.
foreach $fileID (sort keys %actualSpeech) {
  $totalActualSpeech += $actualSpeech{$fileID};
  $totalNonSpeechEst += $maxEndTime{$fileID} - $actualSpeech{$fileID};
  #######################################################################
  # Print file-wise statistics to STDOUT; can pipe to /dev/null is needed
  #######################################################################
  printf STDOUT ("%s: %.2f min actual speech, %.2f min hypothesized: %.2f min overlap (%d\%), %.2f min false alarm (~%d\%)\n",
    $fileID,
    ($actualSpeech{$fileID}/60.0),
    ($hypothesizedSpeech{$fileID}/60.0),
    ($foundSpeech{$fileID}/60.0),
    ($foundSpeech{$fileID}*100/($actualSpeech{$fileID}+0.01)),
    ($falseAlarm{$fileID}/60.0),
    ($falseAlarm{$fileID}*100/($maxEndTime{$fileID}-$actualSpeech{$fileID}+0.01)));
}

################################################################################
# Finally, we have everything needed to report the segmentation statistics.
################################################################################

printf STDERR ("------------------------------------------------------------------------\n");
printf STDERR ("TOTAL: %.2f hrs actual speech, %.2f hrs hypothesized: %.2f hrs overlap (%d\%), %.2f hrs false alarm (~%d\%)\n",
  ($totalActualSpeech/3600.0),
  ($totalHypSpeech/3600.0),
  ($totalFoundSpeech/3600.0),
  ($totalFoundSpeech*100/($totalActualSpeech+0.000001)),
  ($totalFalseAlarm/3600.0),
  ($totalFalseAlarm*100/($totalNonSpeechEst+0.000001)));
printf STDERR ("\t$numShortSegs segments < 0.2 sec and $numLongSegs segments > 60.0 sec\n");
printf STDERR ("------------------------------------------------------------------------\n");
