#!/usr/bin/env python3
  
  # Copyright 2016   Vimal Manohar
  #           2016   Johns Hopkins University (author: Daniel Povey)
  # Apache 2.0
  
  from __future__ import print_function
  import sys, operator, argparse
  
  # Modify the CTM to include for each token the information from Levenshtein
  # alignment of 'hypothesis' and 'reference'
  # (i.e. the output of 'align-text'.
  
  # The information added to each token in the CTM is the reference word and one
  # of the following edit-types:
  #  'cor' = correct  [note: as a special case we count as correct cases where
  #                    the hypothesis word is the OOV symbol and the reference
  #                    word is OOV w.r.t. the supplied vocabulary.]
  #  'sub' = substitution
  #  'del' = deletion
  #  'ins' = insertion
  #  'sil' = (silence in ctm; does not consume a reference word)
  # note: the script modify_ctm_edits.py will add the new
  # note: the following extra edit-type may be added by modify_ctm_edits.py:
  #  'fix'  ... this is like 'cor', but it means the reference has been modified
  #             to fix non-scoreable errors [typically errors that don't change the
  #             meaning], so we don't trust the word or value it as much as a 'cor'.
  #
  
  # Note: Additional lines are added to the CTM to account for deletions.
  
  # Input CTM:
  # (note: the <eps> is for silence in the input CTM that comes from
  # optional-silence in the graph.  However, the input edits don't have anything
  # for these silences.
  # We assume (and check) that the channel will always be '1', because the
  # input CTMs are expected to be 'per utterance', not including real
  # recording-ids.
  
  # Input ctm format:
  # <file-id> <channel> <start-time> <duration> <hyp-word> [<confidence>]
  # note, the confidence defaults to 1 if not provided (these
  # scripts don't actually use the confidence field).
  
  ## TimBrown_2008P-0007226-0007620 1 0.000 0.100 when
  ## TimBrown_2008P-0007226-0007620 1 0.100 0.090 i
  ## TimBrown_2008P-0007226-0007620 1 0.190 0.300 some
  ## TimBrown_2008P-0007226-0007620 1 0.490 0.110 when
  ## TimBrown_2008P-0007226-0007620 1 0.600 0.060 i
  ## TimBrown_2008P-0007226-0007620 1 0.660 0.190 say
  ## TimBrown_2008P-0007226-0007620 1 0.850 0.450 go
  ## TimBrown_2008P-0007226-0007620 1 1.300 0.310 [COUGH]
  ## TimBrown_2008P-0007226-0007620 1 1.610 0.130 you
  ## TimBrown_2008P-0007226-0007620 1 1.740 0.180 got
  ## TimBrown_2008P-0007226-0007620 1 1.920 0.370 thirty
  ## TimBrown_2008P-0007226-0007620 1 2.290 0.830 seconds
  ## TimBrown_2008P-0007226-0007620 1 3.120 0.330 <eps>
  ## TimBrown_2008P-0007226-0007620 1 3.450 0.040 [BREATH]
  ## TimBrown_2008P-0007226-0007620 1 3.490 0.110 to
  ## TimBrown_2008P-0007226-0007620 1 3.600 0.320 [NOISE]
  
  # Input Levenshtein edits : (the output of 'align-text' post-processed by 'wer_per_utt_details.pl')
  
  # AJJacobs_2007P-0001605-0003029 i i ; thought thought ; i'd i'd ; tell tell ; you you ; a a ; little little ; about about ; [UH] [UH] ; what what ; i i ; like like ; to to ; write write ; and and ; [UH] [UH] ; i i ; like like ; to to ; [UH] [UH] ; immerse immerse ; myself myself ; [SMACK] [SMACK] ; in in ; my my ; topics topics ; [UM] [UM] ; i i ; just just ; like like ; to to ; [UH] [UH] ; dive dive ; [SMACK] [SMACK] ; right right ; in in ; and and ; become become ; [UH] [UH] ; sort sort ; of of ; a a ; human human ; guinea guinea ; pig pig ; [BREATH] [BREATH] ; and and ; [UH] [UH]
  # AJJacobs_2007P-0003133-0004110 i i ; see see ; my my ; life life ; as as ; a a ; series series ; of of ; experiments experiments ; [BREATH] [BREATH] ; so so ; [UH] [UH] ; i i ; [NOISE] [NOISE] ; work work ; for for ; esquire esquire ; magazine magazine ; <eps> and ; a a ; couple couple ; of of ; years years ; ago ago ; [BREATH] [BREATH] ; i i ; wrote wrote ; an an ; article article ; called called ; [NOISE] [NOISE] ; my my ; outsourced outsourced ; life life
  
  
  # Output format:
  # <file-id> <channel> <start-time> <duration> <hyp-word> <confidence> <ref-word> <edit-type>
  
  # AJJacobs_2007P-0001605-0003029 1 0 0.09 <eps> 1.0 <eps> sil
  # AJJacobs_2007P-0001605-0003029 1 0.09 0.15 i 1.0 i cor
  # AJJacobs_2007P-0001605-0003029 1 0.24 0.25 thought 1.0 thought cor
  # AJJacobs_2007P-0001605-0003029 1 0.49 0.14 i'd 1.0 i'd cor
  # AJJacobs_2007P-0001605-0003029 1 0.63 0.22 tell 1.0 tell cor
  # AJJacobs_2007P-0001605-0003029 1 0.85 0.11 you 1.0 you cor
  # AJJacobs_2007P-0001605-0003029 1 0.96 0.05 a 1.0 a cor
  # AJJacobs_2007P-0001605-0003029 1 1.01 0.24 little 1.0 little cor
  # AJJacobs_2007P-0001605-0003029 1 1.25 0.5 about 1.0 about cor
  # AJJacobs_2007P-0001605-0003029 1 1.75 0.48 [UH] 1.0 [UH] cor
  # AJJacobs_2007P-0001605-0003029 1 2.23 0.34 <eps> 1.0 <eps> sil
  # AJJacobs_2007P-0001605-0003029 1 2.57 0.21 what 1.0 what cor
  # AJJacobs_2007P-0001605-0003029 1 2.78 0.1 i 1.0 i cor
  # AJJacobs_2007P-0001605-0003029 1 2.88 0.22 like 1.0 like cor
  # AJJacobs_2007P-0001605-0003029 1 3.1 0.13 to 1.0 to cor
  # AJJacobs_2007P-0001605-0003029 1 3.23 0.37 write 1.0 write cor
  # AJJacobs_2007P-0001605-0003029 1 3.6 0.03 <eps> 1.0 <eps> sil
  # AJJacobs_2007P-0001605-0003029 1 3.63 0.36 and 1.0 and cor
  
  
  
  parser = argparse.ArgumentParser(
      description = "Append to the CTM the Levenshtein alignment of 'hypothesis' and 'reference'; "
      "creates augmented CTM with extra fields (see script for details)")
  
  parser.add_argument("--oov", type = int, default = -1,
                      help = "The integer representation of the OOV symbol; substitutions "
                      "by the OOV symbol for out-of-vocabulary reference words are treated "
                      "as correct, if you also supply the --symbol-table option.")
  parser.add_argument("--symbol-table", type = str,
                      help = "The words.txt your system used; if supplied, it is used to "
                      "determine OOV words (and such words will count as correct if "
                      "substituted by the OOV symbol).  See also the --oov option")
  # Required arguments
  parser.add_argument("edits_in", metavar = "<edits-in>",
                      help = "Filename of output of 'align-text', which this program reads. "
                      "Use /dev/stdin for standard input.")
  parser.add_argument("ctm_in", metavar = "<ctm-in>",
                      help = "Filename of input hypothesis in ctm format")
  parser.add_argument("ctm_edits_out", metavar = "<ctm-edits-out>",
                      help = "Filename of output (CTM appended with word-edit information)")
  args = parser.parse_args()
  
  
  
  def OpenFiles():
      global ctm_edits_out, edits_in, ctm_in, symbol_table, oov_word
      try:
          ctm_edits_out = open(args.ctm_edits_out, 'w', encoding='utf-8')
      except:
          sys.exit("get_ctm_edits.py: error opening ctm-edits file {0} for output".format(
                  args.ctm_edits_out))
      try:
          edits_in = open(args.edits_in, encoding='utf-8')
      except:
          sys.exit("get_ctm_edits.py: error opening edits file {0} for input".format(
                  args.edits_in))
      try:
          ctm_in = open(args.ctm_in, encoding='utf-8')
      except:
          sys.exit("get_ctm_edits.py: error opening ctm file {0} for input".format(
                  args.ctm_in))
  
      symbol_table = set()
      oov_word = None
      if args.symbol_table != None:
          if args.oov == -1:
              print("get_ctm_edits.py: error: if you set the the --symbol-table option "
                    "you must also set the --oov option", file = sys.stderr)
          try:
              f = open(args.symbol_table, 'r', encoding='utf-8')
              for line in f.readlines():
                  [ word, integer ] = line.split()
                  if int(integer) == args.oov:
                      oov_word = word
                  symbol_table.add(word)
          except:
              sys.exit("get_ctm_edits.py: error opening symbol-table file {0} for "
                       "input (or bad file), exception is: {1}".format(args.symbol_table))
          f.close()
          if oov_word == None:
              sys.exit("get_ctm_edits.py: OOV word not found: check the values of "
                       "--symbol-table={0} and --oov={1}".format(args.symbol_table,
                                                                 args.oov))
  
  # This function takes two lists
  # edits_array = [ [ hyp_word1, ref_word1], [ hyp_word2, ref_word2 ], ... ]
  # ctm_array = [ [ start1, duration1, hyp_word1, confidence1 ], ... ]
  #
  # and pads them with new list elements so that the entries 'match up'.  What we
  # are aiming for is that for each i, ctm_array[i][2] == edits_array[i][0].  The
  # reasons why this is not automatically true are:
  #
  #  (1) There may be deletions in the hypothesis sequence, which would lead to
  #      pairs like [ '<eps>', ref_word ].
  #  (2) The ctm may have been written 'with silence', which will lead to
  #      ctm entries like [ 1, 7.8, 0.9, '<eps>' ] where the '<eps>' refers
  #      to the optional-silence from the lexicon.
  #
  # We introduce suitable entries in to edits_array and ctm_array as necessary
  # to make them 'match up'.  This function returns the pair (new_edits_array,
  # new_ctm_array).
  def PadArrays(edits_array, ctm_array):
      new_edits_array = []
      new_ctm_array = []
      edits_len = len(edits_array)
      ctm_len = len(ctm_array)
      edits_pos = 0
      ctm_pos = 0
      # current_time is the end of the last ctm segment we processesed.
      current_time = ctm_array[0][0] if ctm_len > 0 else 0.0
      while edits_pos < edits_len or ctm_pos < ctm_len:
          if edits_pos < edits_len and ctm_pos < ctm_len and \
                  edits_array[edits_pos][0] == ctm_array[ctm_pos][2] and \
                  edits_array[edits_pos][0] != '<eps>':
              # This is the normal case, where there are 2 entries where
              # they hyp-words match up
              new_edits_array.append(edits_array[edits_pos])
              edits_pos += 1
              new_ctm_array.append(ctm_array[ctm_pos])
              current_time = ctm_array[ctm_pos][0] + ctm_array[ctm_pos][1]
              ctm_pos += 1
          elif edits_pos < edits_len and edits_array[edits_pos][0] == '<eps>':
              # There was a deletion.  Pad with an empty ctm segment with '<eps>' as
              # the word.
              new_edits_array.append(edits_array[edits_pos])
              edits_pos += 1
              duration = 0.0
              confidence = 1.0
              new_ctm_array.append([ current_time, duration, '<eps>', confidence])
          elif ctm_pos < ctm_len and ctm_array[ctm_pos][2] == '<eps>':
              # There was silence in the ctm, and either we're reached the end of the
              # edits sequence, or the hyp word was not '<eps>':
  
              new_edits_array.append(['<eps>', '<eps>'])
              new_ctm_array.append(ctm_array[ctm_pos])
              current_time = ctm_array[ctm_pos][0] + ctm_array[ctm_pos][1]
              ctm_pos += 1
          else:
              raise Exception("Could not align edits_array = {0} and ctm_array = {1}; "
                              "edits-position = {2}, ctm-position = {3}, "
                              "pending-edit={4}, pending-ctm-entry={5}".format(
                      edits_array, ctm_array, edits_pos, ctm_pos,
                      edits_array[edits_pos] if edits_pos < edits_len else None,
                      ctm_array[ctm_pos] if ctm_pos < ctm_len else None))
      assert len(new_edits_array) == len(new_ctm_array)
      return (new_edits_array, new_ctm_array)
  
  
  # This function returns the appropriate edit-type to output in the ctm-edits
  # file.  The ref_word and hyp_word and duration are the values we'll print in
  # the ctm-edits file.
  def GetEditType(hyp_word, ref_word, duration):
      global oov_word
      if hyp_word == ref_word and hyp_word !='<eps>':
          return 'cor'
      elif hyp_word != '<eps>' and ref_word == '<eps>':
          return 'ins'
      elif hyp_word == '<eps>' and ref_word != '<eps>' and duration == 0.0:
          return 'del'
      elif hyp_word == oov_word and \
           len(symbol_table) != 0 and not ref_word in symbol_table:
          return 'cor'   # this special case is treated as correct.
      elif hyp_word == '<eps>' == ref_word and duration > 0.0:
          # silence in hypothesis; we don't match this up with any reference word.
          return 'sil'
      else:
          # The following assertion is because, based on how PadArrays
          # works, we shouldn't hit this case.
          assert hyp_word != '<eps>' and ref_word != '<eps>'
          return 'sub'
  
  # this prints a number with a certain number of digits after
  # the point, while removing trailing zeros.
  def FloatToString(f):
      num_digits = 6 # we want to print 6 digits after the zero
      g = f
      while abs(g) > 1.0:
          g *= 0.1
          num_digits += 1
      format_str = '%.{0}g'.format(num_digits)
      return format_str % f
  
  
  def OutputCtm(utterance_id, edits_array, ctm_array):
      global ctm_edits_out
      # note: this function expects the padded entries created by PadARrays.
      assert len(edits_array) == len(ctm_array)
      channel = '1'  # this is hardcoded at both input and output, since this CTM
                     # doesn't really represent recordings, only utterances.
      for i in range(len(edits_array)):
          ( hyp_word, ref_word ) = edits_array[i]
          ( start_time, duration, hyp_word2, confidence ) = ctm_array[i]
          if not hyp_word == hyp_word2:
              print("Error producing output CTM for edit = {0} and ctm = {1}".format(
                      edits_array[i], ctm_array[i]), file = sys.stderr)
              sys.exit(1)
          assert hyp_word == hyp_word2
          edit_type = GetEditType(hyp_word, ref_word, duration)
          print(utterance_id, channel, FloatToString(start_time),
                FloatToString(duration), hyp_word, confidence, ref_word,
                edit_type, file = ctm_edits_out)
  
  
  def ProcessOneUtterance(utterance_id, edits_line, ctm_lines):
      try:
          # Remove the utterance-id from the beginning of the edits line
          edits_fields = edits_line[len(utterance_id) + 1:]
  
          # e.g. if edits_fields is now 'i i ; see be ; my my ', edits_array will become
          #  [ ['i', 'i'], ['see', 'be'], ['my', 'my'] ]
          fields_split = edits_fields.split()
          first_fields, second_fields = fields_split[0::3], fields_split[1::3]
          if (
              len(first_fields) != len(second_fields) or
              (len(fields_split) >= 3 and set(fields_split[2::3]) != {';'})
          ):
              sys.exit("get_ctm_edits.py: could not make sense of edits line: " + edits_line)
  
          edits_array = list(zip(first_fields, second_fields))
  
          # ctm_array will now become something like [ ['1', '1.010', '0.240', 'little ' ], ... ]
          ctm_array = [ x.split() for x in ctm_lines ]
          ctm_array = []
          for line in ctm_lines:
              try:
                  # Strip off the utterance-id and split the remaining fields
                  # which should be: channel==1, start, dur, word, [confidence]
                  a = line[len(utterance_id) + 1:].split()
                  if len(a) == 4:
                      a.append(1.0)  # confidence defaults to 1.0.
                  [ channel, start, dur, word, confidence ] = a
                  if channel != '1':
                      raise Exception("Channel should be 1, got: " + channel)
                  ctm_array.append([ float(start), float(dur), word, float(confidence) ])
              except Exception as e:
                  sys.exit("get_ctm_edits.py: error procesing ctm line {0} "
                           "... exception is: {1} {2}".format(line, type(e), str(e)))
          # ctm_array will now be something like [ [ 1.010, 0.240, 'little ', 1.0 ], ... ]
  
          # The following call pads the edits and ctm arrays with appropriate
          # entries so that they have the same length and the elements 'match up'.
          (edits_array, ctm_array) = PadArrays(edits_array, ctm_array)
      except Exception as e:
          sys.exit("get_ctm_edits.py: error processing utterance {0}, error was: {1}".format(
                  utterance_id, str(e)))
      OutputCtm(utterance_id, edits_array, ctm_array)
  
  def ProcessData():
      num_utterances_processed = 0
  
      pending_ctm_line = ctm_in.readline()
  
      while True:
          this_edits_line = edits_in.readline()
          if this_edits_line == '':
              if pending_ctm_line != '':
                  sys.exit("get_ctm_edits.py: edits_in input {0} ended before "
                           "ctm input was ended.  We processed {1} "
                           "utterances.".format(args.edits_in, num_utterances_processed))
              break
          a = this_edits_line.split()
          if len(a) == 0:
              sys.exit("get_ctm_edits.py: edits_input {0} had an empty line".format(
                      args.edits_in))
          utterance_id = a[0]
          utterance_id_len = len(utterance_id)
          this_utterance_ctm_lines = []
          while len(pending_ctm_line.strip()) > 0 and pending_ctm_line.split()[0] == utterance_id:
              this_utterance_ctm_lines.append(pending_ctm_line)
              pending_ctm_line = ctm_in.readline()
          ProcessOneUtterance(utterance_id, this_edits_line,
                              this_utterance_ctm_lines)
          num_utterances_processed += 1
      print("get_ctm_edits.py: processed {0} utterances".format(
              num_utterances_processed), file=sys.stderr)
  
  
  OpenFiles()
  ProcessData()