#!/usr/bin/env python
  
  # Copyright 2017  Vimal Manohar
  #           2018  Capital One (Author: Zhiyuan Guan)
  # Apache 2.0
  
  """
  This script converts frame-level speech activity detection marks (in kaldi
  integer vector text archive format) into kaldi segments and utt2spk.
  The input integer vectors are expected to contain '1' for silence frames
  and '2' for speech frames.
  """
  
  from __future__ import print_function
  import argparse
  import logging
  import sys
  
  sys.path.insert(0, 'steps')
  import libs.common as common_lib
  
  logger = logging.getLogger(__name__)
  logger.setLevel(logging.INFO)
  handler = logging.StreamHandler()
  handler.setLevel(logging.INFO)
  formatter = logging.Formatter("%(asctime)s [%(pathname)s:%(lineno)s - "
                                "%(funcName)s - %(levelname)s ] %(message)s")
  handler.setFormatter(formatter)
  logger.addHandler(handler)
  
  global_verbose = 0
  
  
  def get_args():
      parser = argparse.ArgumentParser(
          description="""
  This script converts frame-level speech activity detection marks (in kaldi
  integer vector text archive format) into kaldi segments and utt2spk.
  The input integer vectors are expected to contain 1 for silence frames
  and 2 for speech frames.
  """,
          formatter_class=argparse.RawTextHelpFormatter)
  
      parser.add_argument("--verbose", type=int, choices=[0, 1, 2, 3],
                          default=0, help="Higher verbosity for more logging")
  
      parser.add_argument("--utt2dur", type=str,
                          help="File containing durations of utterances.")
  
      parser.add_argument("--frame-shift", type=float, default=0.01,
                          help="Frame shift to convert frame indexes to time")
  
      parser.add_argument("--segment-padding", type=float, default=0.2,
                          help="Additional padding on speech segments. But we "
                               "ensure that the padding does not go beyond the "
                               "adjacent segment.")
  
      parser.add_argument("--min-segment-dur", type=float, default=0,
                          help="Minimum duration (in seconds) required for a segment "
                               "to be included. This is before any padding. Segments "
                               "shorter than this duration will be removed.")
  
      parser.add_argument("--merge-consecutive-max-dur", type=float, default=0,
                          help="Merge consecutive segments as long as the merged "
                               "segment is no longer than this many seconds. The segments "
                               "are only merged if their boundaries are touching. "
                               "This is after padding by --segment-padding seconds."
                               "0 means do not merge. Use 'inf' to not limit the duration.")
  
      parser.add_argument("in_sad", type=str,
                          help="Input file containing alignments in "
                               "text archive format")
  
      parser.add_argument("out_segments", type=str,
                          help="Output kaldi segments file")
  
      args = parser.parse_args()
  
      global global_verbose
      global_verbose = args.verbose
  
      logger.info("Setting verbosity to {0}".format(global_verbose))
  
      if args.verbose >= 3:
          logger.setLevel(logging.DEBUG)
          handler.setLevel(logging.DEBUG)
      return args
  
  
  def to_str(segment):
      assert len(segment) == 3
      return "[{0:.3f}, {1:.3f}, {2}]".format(segment[0], segment[1],
                                              segment[2])
  
  
  class SegmenterStats(object):
      """Stores stats about the post-process stages"""
  
      def __init__(self):
          self.num_segments_initial = 0
          self.num_short_segments_filtered = 0
          self.num_merges = 0
          self.num_segments_final = 0
          self.initial_duration = 0.0
          self.padding_duration = 0.0
          self.filter_short_duration = 0.0
          self.final_duration = 0.0
  
      def add(self, other):
          """Adds stats from another object"""
          self.num_segments_initial += other.num_segments_initial
          self.num_short_segments_filtered += other.num_short_segments_filtered
          self.num_merges += other.num_merges
          self.num_segments_final += other.num_segments_final
          self.initial_duration += other.initial_duration
          self.filter_short_duration += other.filter_short_duration
          self.padding_duration += other.padding_duration
          self.final_duration += other.final_duration
  
      def __str__(self):
          return ("num-segments-initial={num_segments_initial}, "
                  "num-short-segments-filtered={num_short_segments_filtered}, "
                  "num-merges={num_merges}, "
                  "num-segments-final={num_segments_final}, "
                  "initial-duration={initial_duration}, "
                  "filter-short-duration={filter_short_duration}, "
                  "padding-duration={padding_duration}, "
                  "final-duration={final_duration}".format(
              num_segments_initial=self.num_segments_initial,
              num_short_segments_filtered=self.num_short_segments_filtered,
              num_merges=self.num_merges,
              num_segments_final=self.num_segments_final,
              initial_duration=self.initial_duration,
              filter_short_duration=self.filter_short_duration,
              padding_duration=self.padding_duration,
              final_duration=self.final_duration))
  
  
  def process_label(text_label):
      """Processes an input integer label and returns a 1 or 2,
      where 1 is for silence and 2 is for speech.
  
      Arguments:
          text_label -- input label (must be integer)
      """
      prev_label = int(text_label)
      if prev_label not in [1, 2]:
          raise ValueError("Expecting label to 1 (non-speech) or 2 (speech); "
                           "got {}".format(prev_label))
  
      return prev_label
  
  
  class Segmentation(object):
      """Stores segmentation for an utterances"""
  
      def __init__(self):
          self.segments = None
          self.stats = SegmenterStats()
  
      def initialize_segments(self, alignment, frame_shift=0.01):
          """Initializes segments from input alignment.
          The alignment is frame-level speech-activity detection marks,
          each of which must be 1 or 2."""
          self.segments = []
  
          assert len(alignment) > 0
  
          prev_label = None
          prev_length = 0
          for i, text_label in enumerate(alignment):
              if prev_label is not None and int(text_label) != prev_label:
                  if prev_label == 2:
                      self.segments.append(
                          [float(i - prev_length) * frame_shift,
                           float(i) * frame_shift, prev_label])
                      self.stats.initial_duration += (prev_length * frame_shift)
                  prev_label = process_label(text_label)
                  prev_length = 0
              elif prev_label is None:
                  prev_label = process_label(text_label)
  
              prev_length += 1
  
          if prev_length > 0 and prev_label == 2:
              self.segments.append(
                  [float(len(alignment) - prev_length) * frame_shift,
                   float(len(alignment)) * frame_shift, prev_label])
              self.stats.initial_duration += (prev_length * frame_shift)
  
          self.stats.num_segments_initial = len(self.segments)
          self.stats.num_segments_final = len(self.segments)
          self.stats.final_duration = self.stats.initial_duration
  
      def filter_short_segments(self, min_dur):
          """Filters out segments with durations shorter than 'min_dur'."""
          if min_dur <= 0:
              return
  
          segments_kept = []
          for segment in self.segments:
              assert segment[2] == 2, segment
              dur = segment[1] - segment[0]
              if dur < min_dur:
                  self.stats.filter_short_duration += dur
                  self.stats.num_short_segments_filtered += 1
              else:
                  segments_kept.append(segment)
          self.segments = segments_kept
          self.stats.num_segments_final = len(self.segments)
          self.stats.final_duration -= self.stats.filter_short_duration
  
      def pad_speech_segments(self, segment_padding, max_duration=float("inf")):
          """Pads segments by duration 'segment_padding' on either sides, but
          ensures that the segments don't go beyond the neighboring segments
          or the duration of the utterance 'max_duration'."""
          if max_duration == None:
              max_duration = float("inf")
          for i, segment in enumerate(self.segments):
              assert segment[2] == 2, segment
              segment[0] -= segment_padding  # try adding padding on the left side
              self.stats.padding_duration += segment_padding
              if segment[0] < 0.0:
                  # Padding takes the segment start to before the beginning of the utterance.
                  # Reduce padding.
                  self.stats.padding_duration += segment[0]
                  segment[0] = 0.0
              if i >= 1 and self.segments[i - 1][1] > segment[0]:
                  # Padding takes the segment start to before the end the previous segment.
                  # Reduce padding.
                  self.stats.padding_duration -= (
                          self.segments[i - 1][1] - segment[0])
                  segment[0] = self.segments[i - 1][1]
  
              segment[1] += segment_padding
              self.stats.padding_duration += segment_padding
              if segment[1] >= max_duration:
                  # Padding takes the segment end beyond the max duration of the utterance.
                  # Reduce padding.
                  self.stats.padding_duration -= (segment[1] - max_duration)
                  segment[1] = max_duration
              if (i + 1 < len(self.segments)
                      and segment[1] > self.segments[i + 1][0]):
                  # Padding takes the segment end beyond the start of the next segment.
                  # Reduce padding.
                  self.stats.padding_duration -= (
                          segment[1] - self.segments[i + 1][0])
                  segment[1] = self.segments[i + 1][0]
          self.stats.final_duration += self.stats.padding_duration
  
      def merge_consecutive_segments(self, max_dur):
          """Merge consecutive segments (happens after padding), provided that
          the merged segment is no longer than 'max_dur'."""
          if max_dur <= 0 or not self.segments:
              return
  
          merged_segments = [self.segments[0]]
          for segment in self.segments[1:]:
              assert segment[2] == 2, segment
              if segment[0] == merged_segments[-1][1] and \
                      segment[1] - merged_segments[-1][0] <= max_dur:
                  # The segment starts at the same time the last segment ends,
                  # and the merged segment is shorter than 'max_dur'.
                  # Extend the previous segment.
                  merged_segments[-1][1] = segment[1]
                  self.stats.num_merges += 1
              else:
                  merged_segments.append(segment)
  
          self.segments = merged_segments
          self.stats.num_segments_final = len(self.segments)
  
      def write(self, key, file_handle):
          """Write segments to file"""
          if global_verbose >= 2:
              logger.info("For key {key}, got stats {stats}".format(
                  key=key, stats=self.stats))
          for segment in self.segments:
              seg_id = "{key}-{st:07d}-{end:07d}".format(
                  key=key, st=int(segment[0] * 100), end=int(segment[1] * 100))
              print("{seg_id} {key} {st:.2f} {end:.2f}".format(
                  seg_id=seg_id, key=key, st=segment[0], end=segment[1]),
                  file=file_handle)
  
  
  def run(args):
      """The main function that does everything."""
      utt2dur = {}
      if args.utt2dur is not None:
          with common_lib.smart_open(args.utt2dur) as utt2dur_fh:
              for line in utt2dur_fh:
                  parts = line.strip().split()
                  if len(parts) != 2:
                      raise RuntimeError("Unable to parse line '{0}' in {1}"
                                         "".format(line.strip(), args.utt2dur))
                  utt2dur[parts[0]] = float(parts[1])
  
      global_stats = SegmenterStats()
      with common_lib.smart_open(args.in_sad) as in_sad_fh, \
              common_lib.smart_open(args.out_segments, 'w') as out_segments_fh:
          for line in in_sad_fh:
              parts = line.strip().split()
              utt_id = parts[0]
  
              if len(parts) < 2:
                  raise RuntimeError("Unable to parse line '{0}' in {1}"
                                     "".format(line.strip(),
                                               in_sad_fh))
  
              segmentation = Segmentation()
              segmentation.initialize_segments(
                  parts[1:], args.frame_shift)
              segmentation.filter_short_segments(args.min_segment_dur)
              segmentation.pad_speech_segments(args.segment_padding,
                                               None if args.utt2dur is None
                                               else utt2dur[utt_id])
              segmentation.merge_consecutive_segments(args.merge_consecutive_max_dur)
              segmentation.write(utt_id, out_segments_fh)
              global_stats.add(segmentation.stats)
      logger.info(global_stats)
  
  
  def main():
      """Parses arguments and calls the run method"""
      args = get_args()
      try:
          run(args)
      except Exception:
          raise
  
  
  if __name__ == '__main__':
      main()