#! /usr/bin/env python
  
  # Copyright 2016    Vimal Manohar
  # Apache 2.0.
  
  """This module aligns a hypothesis (CTM or text) with a reference to
  find the best matching sub-sequence in the reference for the hypothesis
  using Smith-Waterman like alignment.
  
  e.g.: align_ctm_ref.py --hyp-format=CTM --ref=data/train/text --hyp=foo/ctm
          --output=foo/ctm_edits
  """
  
  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__)
  handler = logging.StreamHandler()
  formatter = logging.Formatter("%(asctime)s [%(pathname)s:%(lineno)s - "
                                "%(funcName)s - %(levelname)s ] %(message)s")
  handler.setFormatter(formatter)
  logger.setLevel(logging.DEBUG)
  
  verbose_level = 0
  
  
  def get_args():
      parser = argparse.ArgumentParser(description="""
      This module aligns a hypothesis (CTM or text) with a reference to find the
      best matching sub-sequence in the reference for the hypothesis using
      Smith-Waterman like alignment.
  
      e.g.: align_ctm_ref.py --align-full-hyp=false --hyp-format=CTM
      --reco2file-and-channel=data/foo/reco2file_and_channel --ref=data/train/text
      --hyp=foo/ctm --output=foo/ctm_edits
      """)
  
      parser.add_argument("--hyp-format", type=str, choices=["Text", "CTM"],
                          default="CTM",
                          help="Format used for the hypothesis")
      parser.add_argument("--reco2file-and-channel", type=argparse.FileType('r'),
                          help="""reco2file_and_channel file.
                          This will be used to match references that are usually
                          indexed by the recording-id with the CTM lines that have
                          file and channel. This option is typically not
                          required.""")
      parser.add_argument("--eps-symbol", type=str, default="-",
                          help="Symbol used to contain alignment "
                          "to empty symbol")
      parser.add_argument("--oov-word", type=str, default=None,
                          action=common_lib.NullstrToNoneAction,
                          help="Symbol of OOV word in hypothesis")
      parser.add_argument("--symbol-table", type=argparse.FileType('r'),
                          help="""Symbol table for words in vocabulary. Used
                          to determine if a word is a OOV or not""")
  
      parser.add_argument("--correct-score", type=int, default=1,
                          help="Score for correct matches")
      parser.add_argument("--substitution-penalty", type=int, default=1,
                          help="Penalty for substitution errors")
      parser.add_argument("--deletion-penalty", type=int, default=1,
                          help="Penalty for deletion errors")
      parser.add_argument("--insertion-penalty", type=int, default=1,
                          help="Penalty for insertion errors")
  
      parser.add_argument("--align-full-hyp", type=str,
                          action=common_lib.StrToBoolAction,
                          choices=["true", "false"], default=True,
                          help="""Align full hypothesis i.e. trackback from
                          the end to get the alignment. This is different
                          from the normal Smith-Waterman alignment, where the
                          traceback will be from the maximum score.""")
  
      parser.add_argument("--debug-only", type=str, default="false",
                          choices=["true", "false"],
                          help="Run test functions only")
      parser.add_argument("--verbose", type=int, default=0,
                          choices=[0, 1, 2, 3],
                          help="Use larger value for more verbose logging.")
  
      parser.add_argument("--ref", dest='ref_in_file',
                          type=argparse.FileType('r'), required=True,
                          help="Reference text file")
      parser.add_argument("--hyp", dest='hyp_in_file', required=True,
                          type=argparse.FileType('r'),
                          help="Hypothesis text or CTM file")
      parser.add_argument("--output", dest='alignment_out_file', required=True,
                          type=argparse.FileType('w'),
                          help="""File to write output alignment.
                          If hyp-format=CTM, then the output is in the form of
                          CTM, but with two additional columns of Edit-type and
                          Reference-word matched to the hypothesis.""")
  
      args = parser.parse_args()
  
      args.debug_only = bool(args.debug_only == "true")
  
      global verbose_level
      verbose_level = args.verbose
      if args.verbose > 2:
          handler.setLevel(logging.DEBUG)
      else:
          handler.setLevel(logging.INFO)
      logger.addHandler(handler)
  
      return args
  
  
  def read_text(text_file):
      """Reads a kaldi-format text file and yield elements of a dictionary
          { utterane_id : transcript (as a list of words) }
  
      The first-column of the text file is the utterance-id, which will be
      used as the key to index the dictionary elements.
      The remaining columns of the file are text of the transcript and they are
      returned as a list of words.
      """
      for line in text_file:
          parts = line.strip().split()
          if len(parts) < 1:
              raise RuntimeError(
                  "Did not get enough columns; line {0} in {1}"
                  "".format(line, text_file.name))
          elif len(parts) == 1:
              logger.warn("Empty transcript for utterance %s in %s",
                          parts[0], text_file.name)
              yield parts[0], []
          else:
              yield parts[0], parts[1:]
      text_file.close()
  
  
  def read_ctm(ctm_file, file_and_channel2reco=None):
      """Reads a CTM file and yields elements of a dictionary
          { utterance-id : CTM for the utterance },
      where CTM for the utterance is stored as a list of lines
      from a CTM correponding to the utterance.
  
      Note: *_reco in the variables usually correspond to utterances rather
      than recordings.
      """
      prev_reco = ""
      ctm_lines = []
      for line in ctm_file:
          try:
              parts = line.strip().split()
              parts[2] = float(parts[2])
              parts[3] = float(parts[3])
  
              if len(parts) == 5:
                  parts.append(1.0)   # confidence defaults to 1.0.
  
              if len(parts) != 6:
                  raise ValueError("CTM must have 6 fields.")
  
              if file_and_channel2reco is None:
                  reco = parts[0]
                  if parts[1] != '1':
                      raise ValueError("Channel should be 1, "
                                       "got {0}".format(parts[1]))
              else:
                  reco = file_and_channel2reco[(parts[0], parts[1])]
              if prev_reco != "" and reco != prev_reco:
                  # New recording
                  yield prev_reco, ctm_lines
                  ctm_lines = []
              ctm_lines.append(parts[2:])
              prev_reco = reco
          except Exception:
              logger.error("Error in processing CTM line {0}".format(line))
              raise
      if prev_reco != "" and len(ctm_lines) > 0:
          yield prev_reco, ctm_lines
      ctm_file.close()
  
  
  def smith_waterman_alignment(ref, hyp, similarity_score_function,
                               del_score, ins_score,
                               eps_symbol="<eps>", align_full_hyp=True):
      """Does Smith-Waterman alignment of reference sequence and hypothesis
      sequence.
      This is a special case of the Smith-Waterman alignment that assumes that
      the deletion and insertion costs are linear with number of incorrect words.
  
      If align_full_hyp is True, then the traceback of the alignment
      is started at the end of the hypothesis. This is when we want the
      reference that aligns with the full hypothesis.
      This differs from the normal Smith-Waterman alignment, where the traceback
      is from the highest score in the alignment score matrix. This
      can be obtained by setting align_full_hyp as False. This gets only the
      sub-sequence of the hypothesis that best matches with a
      sub-sequence of the reference.
  
      Returns a list of tuples where each tuple has the format:
          (ref_word, hyp_word, ref_word_from_index, hyp_word_from_index,
           ref_word_to_index, hyp_word_to_index)
      """
      output = []
  
      ref_len = len(ref)
      hyp_len = len(hyp)
  
      bp = [[] for x in range(ref_len+1)]
  
      # Score matrix of size (ref_len + 1) x (hyp_len + 1)
      # The index m, n in this matrix corresponds to the score
      # of the best matching sub-sequence pair between reference and hypothesis
      # ending with the reference word ref[m-1] and hypothesis word hyp[n-1].
      # If align_full_hyp is True, then the hypothesis sub-sequence is from
      # the 0th word i.e. hyp[0].
      H = [[] for x in range(ref_len+1)]
  
      for ref_index in range(ref_len+1):
          if align_full_hyp:
              H[ref_index] = [-(hyp_len+2) for x in range(hyp_len+1)]
              H[ref_index][0] = 0
          else:
              H[ref_index] = [0 for x in range(hyp_len+1)]
          bp[ref_index] = [(0, 0) for x in range(hyp_len+1)]
  
          if align_full_hyp and ref_index == 0:
              for hyp_index in range(1, hyp_len+1):
                  H[0][hyp_index] = H[0][hyp_index-1] + ins_score
                  bp[ref_index][hyp_index] = (ref_index, hyp_index-1)
                  logger.debug(
                      "({0},{1}) -> ({2},{3}): {4}"
                      "".format(ref_index, hyp_index-1, ref_index, hyp_index,
                                H[ref_index][hyp_index]))
  
      max_score = -float("inf")
      max_score_element = (0, 0)
  
      for ref_index in range(1, ref_len+1):     # Reference
          for hyp_index in range(1, hyp_len+1):     # Hypothesis
              sub_or_ok = (H[ref_index-1][hyp_index-1]
                           + similarity_score_function(ref[ref_index-1],
                                                       hyp[hyp_index-1]))
  
              if ((not align_full_hyp and sub_or_ok > 0)
                      or (align_full_hyp
                          and sub_or_ok >= H[ref_index][hyp_index])):
                  H[ref_index][hyp_index] = sub_or_ok
                  bp[ref_index][hyp_index] = (ref_index-1, hyp_index-1)
                  logger.debug(
                      "({0},{1}) -> ({2},{3}): {4} ({5},{6})"
                      "".format(ref_index-1, hyp_index-1, ref_index, hyp_index,
                                H[ref_index][hyp_index],
                                ref[ref_index-1], hyp[hyp_index-1]))
  
              if H[ref_index-1][hyp_index] + del_score > H[ref_index][hyp_index]:
                  H[ref_index][hyp_index] = H[ref_index-1][hyp_index] + del_score
                  bp[ref_index][hyp_index] = (ref_index-1, hyp_index)
                  logger.debug(
                      "({0},{1}) -> ({2},{3}): {4}"
                      "".format(ref_index-1, hyp_index, ref_index, hyp_index,
                                H[ref_index][hyp_index]))
  
              if H[ref_index][hyp_index-1] + ins_score > H[ref_index][hyp_index]:
                  H[ref_index][hyp_index] = H[ref_index][hyp_index-1] + ins_score
                  bp[ref_index][hyp_index] = (ref_index, hyp_index-1)
                  logger.debug(
                      "({0},{1}) -> ({2},{3}): {4}"
                      "".format(ref_index, hyp_index-1, ref_index, hyp_index,
                                H[ref_index][hyp_index]))
  
              #if hyp_index == hyp_len and H[ref_index][hyp_index] >= max_score:
              if ((not align_full_hyp or hyp_index == hyp_len)
                      and H[ref_index][hyp_index] >= max_score):
                  max_score = H[ref_index][hyp_index]
                  max_score_element = (ref_index, hyp_index)
  
      ref_index, hyp_index = max_score_element
      score = max_score
      logger.debug("Alignment score: %s for (%d, %d)",
                   score, ref_index, hyp_index)
  
      while ((not align_full_hyp and score >= 0)
             or (align_full_hyp and hyp_index > 0)):
          try:
              prev_ref_index, prev_hyp_index = bp[ref_index][hyp_index]
  
              if ((prev_ref_index, prev_hyp_index) == (ref_index, hyp_index)
                      or (prev_ref_index, prev_hyp_index) == (0, 0)):
                  ref_index, hyp_index = (prev_ref_index, prev_hyp_index)
                  score = H[ref_index][hyp_index]
                  break
  
              if (ref_index == prev_ref_index + 1
                      and hyp_index == prev_hyp_index + 1):
                  # Substitution or correct
                  output.append(
                      (ref[ref_index-1] if ref_index > 0 else eps_symbol,
                       hyp[hyp_index-1] if hyp_index > 0 else eps_symbol,
                       prev_ref_index, prev_hyp_index, ref_index, hyp_index))
              elif (prev_hyp_index == hyp_index):
                  # Deletion
                  assert prev_ref_index == ref_index - 1
                  output.append(
                      (ref[ref_index-1] if ref_index > 0 else eps_symbol,
                       eps_symbol,
                       prev_ref_index, prev_hyp_index, ref_index, hyp_index))
              elif (prev_ref_index == ref_index):
                  # Insertion
                  assert prev_hyp_index == hyp_index - 1
                  output.append(
                      (eps_symbol,
                       hyp[hyp_index-1] if hyp_index > 0 else eps_symbol,
                       prev_ref_index, prev_hyp_index, ref_index, hyp_index))
              else:
                  raise RuntimeError
  
              ref_index, hyp_index = (prev_ref_index, prev_hyp_index)
              score = H[ref_index][hyp_index]
          except Exception:
              logger.error("Unexpected entry (%d,%d) -> (%d,%d), %s, %s",
                           prev_ref_index, prev_hyp_index, ref_index, hyp_index,
                           ref[prev_ref_index], hyp[prev_hyp_index])
              raise RuntimeError("Unexpected result: Bug in code!!")
  
      assert (align_full_hyp or score == 0)
  
      output.reverse()
  
      if verbose_level > 2:
          for ref_index in range(ref_len+1):
              for hyp_index in range(hyp_len+1):
                  print ("{0} ".format(H[ref_index][hyp_index]), end='',
                         file=sys.stderr)
              print ("", file=sys.stderr)
  
      logger.debug("Aligned output:")
      logger.debug("  -  ".join(["({0},{1})".format(x[4], x[5])
                                 for x in output]))
      logger.debug("REF: ")
      logger.debug("    ".join(str(x[0]) for x in output))
      logger.debug("HYP:")
      logger.debug("    ".join(str(x[1]) for x in output))
  
      return (output, max_score)
  
  
  def print_alignment(recording, alignment, out_file_handle):
      out_text = [recording]
      for line in alignment:
          try:
              out_text.append(line[1])
          except Exception:
              logger.error("Something wrong with alignment. "
                           "Invalid line {0}".format(line))
              raise
      print (" ".join(out_text), file=out_file_handle)
  
  
  def get_edit_type(hyp_word, ref_word, duration=-1, eps_symbol='<eps>',
                    oov_word=None, symbol_table=None):
      if hyp_word == ref_word and hyp_word != eps_symbol:
          return 'cor'
      if hyp_word != eps_symbol and ref_word == eps_symbol:
          return 'ins'
      if hyp_word == eps_symbol and ref_word != eps_symbol and duration == 0.0:
          return 'del'
      if (hyp_word == oov_word and symbol_table is not None
              and len(symbol_table) > 0 and ref_word not in symbol_table):
          return 'cor'    # this special case is treated as correct
      if hyp_word == eps_symbol and ref_word == eps_symbol and duration > 0.0:
          # silence in hypothesis; we don't match this up with any reference
          # word.
          return 'sil'
      # The following assertion is because, based on how get_ctm_edits()
      # works, we shouldn't hit this case.
      assert hyp_word != eps_symbol and ref_word != eps_symbol
      return 'sub'
  
  
  def get_ctm_edits(alignment_output, ctm_array, eps_symbol="<eps>",
                    oov_word=None, symbol_table=None):
      """
      This function takes two lists
          alignment_output = The output of smith_waterman_alignment() which is a
              list of tuples (ref_word, hyp_word, ref_word_from_index,
              hyp_word_from_index, ref_word_to_index, hyp_word_to_index)
          ctm_array = [ [ start1, duration1, hyp_word1, confidence1 ], ... ]
      and pads them with new list elements so that the entries 'match up'.
  
      Returns CTM edits lines, which are CTM lines appended with reference word
      and edit type.
  
      What we are aiming for is that for each i, ctm_array[i][2] ==
      alignment_output[i][1].  The reasons why this is not automatically true
      are:
  
       (1) There may be insertions in the hypothesis sequence that are not
           aligned with any reference words in the beginning of the
           alignment_output.
       (2) There may be deletions in the end of the alignment_output that
           do not correspond to any additional hypothesis CTM lines.
  
      We introduce suitable entries in to alignment_output and ctm_array as
      necessary to make them 'match up'.
      """
      ctm_edits = []
      ali_len = len(alignment_output)
      ctm_len = len(ctm_array)
      ali_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
  
      for (ref_word, hyp_word, ref_prev_i, hyp_prev_i,
           ref_i, hyp_i) in alignment_output:
          try:
              ctm_pos = hyp_prev_i
              # This is true because we cannot have errors at the end because
              # that will decrease the smith-waterman alignment score.
              assert ctm_pos < ctm_len
              assert len(ctm_array[ctm_pos]) == 4
  
              if hyp_prev_i == hyp_i:
                  assert hyp_word == eps_symbol
                  # These are deletions as there are no CTM entries
                  # corresponding to these alignments.
                  edit_type = get_edit_type(
                      hyp_word=eps_symbol, ref_word=ref_word,
                      duration=0.0, eps_symbol=eps_symbol,
                      oov_word=oov_word, symbol_table=symbol_table)
                  ctm_line = [current_time, 0.0, eps_symbol, 1.0,
                              ref_word, edit_type]
                  ctm_edits.append(ctm_line)
              else:
                  assert hyp_i == hyp_prev_i + 1
                  assert hyp_word == ctm_array[ctm_pos][2]
                  # This is the normal case, where there are 2 entries where
                  # they hyp-words match up.
                  ctm_line = list(ctm_array[ctm_pos])
                  if hyp_word == eps_symbol and ref_word != eps_symbol:
                      # This is a silence in hypothesis aligned with a reference
                      # word. We split this into two ctm edit lines where the
                      # first one is a deletion of duration 0 and the second
                      # one is a silence of duration given by the ctm line.
                      edit_type = get_edit_type(
                          hyp_word=eps_symbol, ref_word=ref_word,
                          duration=0.0, eps_symbol=eps_symbol,
                          oov_word=oov_word, symbol_table=symbol_table)
                      assert edit_type == 'del'
                      ctm_edits.append([current_time, 0.0, eps_symbol, 1.0,
                                        ref_word, edit_type])
  
                      edit_type = get_edit_type(
                          hyp_word=eps_symbol, ref_word=eps_symbol,
                          duration=ctm_line[1], eps_symbol=eps_symbol,
                          oov_word=oov_word, symbol_table=symbol_table)
                      assert edit_type == 'sil'
                      ctm_line.extend([eps_symbol, edit_type])
                      ctm_edits.append(ctm_line)
                  else:
                      edit_type = get_edit_type(
                          hyp_word=hyp_word, ref_word=ref_word,
                          duration=ctm_line[1], eps_symbol=eps_symbol,
                          oov_word=oov_word, symbol_table=symbol_table)
                      ctm_line.extend([ref_word, edit_type])
                      ctm_edits.append(ctm_line)
                  current_time = (ctm_array[ctm_pos][0]
                                  + ctm_array[ctm_pos][1])
          except Exception:
              logger.error("Could not get ctm edits for "
                           "edits@{edits_pos} = {0}, ctm@{ctm_pos} = {1}".format(
                              ("NONE" if ali_pos >= ali_len
                               else alignment_output[ali_pos]),
                              ("NONE" if ctm_pos >= ctm_len
                               else ctm_array[ctm_pos]),
                              edits_pos=ali_pos, ctm_pos=ctm_pos))
              logger.error("alignment = {0}".format(alignment_output))
              raise
      return ctm_edits
  
  
  def ctm_line_to_string(ctm_line):
      if len(ctm_line) != 8:
          raise RuntimeError("len(ctm_line) expected to be {0}. "
                             "Invalid line {1}".format(8, ctm_line))
  
      return " ".join([str(x) for x in ctm_line])
  
  
  def test_alignment(align_full_hyp):
      hyp = "GCCAT"
      ref = "AGCACACA"
  
      verbose = 3
      logger.info("REF: %s", ref)
      logger.info("HYP: %s", hyp)
  
      output, score = smith_waterman_alignment(
          ref, hyp, similarity_score_function=lambda x, y: 2 if (x == y) else -1,
          del_score=-1, ins_score=-1, eps_symbol="-", align_full_hyp=align_full_hyp)
  
      print_alignment("Alignment", output, out_file_handle=sys.stderr)
  
  
  def run(args):
      if args.debug_only:
          test_alignment(args.align_full_hyp)
          raise SystemExit("Exiting since --debug-only was true")
  
      def similarity_score_function(x, y):
          if x == y:
              return args.correct_score
          return -args.substitution_penalty
  
      del_score = -args.deletion_penalty
      ins_score = -args.insertion_penalty
  
      reco2file_and_channel = {}
      file_and_channel2reco = {}
  
      if args.reco2file_and_channel is not None:
          for line in args.reco2file_and_channel:
              parts = line.strip().split()
  
              reco2file_and_channel[parts[0]] = (parts[1], parts[2])
              file_and_channel2reco[(parts[1], parts[2])] = parts[0]
          args.reco2file_and_channel.close()
      else:
          file_and_channel2reco = None
  
      symbol_table = {}
      if args.symbol_table is not None:
          for line in args.symbol_table:
              parts = line.strip().split()
              symbol_table[parts[0]] = int(parts[1])
          args.symbol_table.close()
  
      if args.hyp_format == "Text":
          hyp_lines = {key: value
                       for (key, value) in read_text(args.hyp_in_file)}
      else:
          hyp_lines = {key: value
                       for (key, value) in read_ctm(args.hyp_in_file,
                                                    file_and_channel2reco)}
  
      num_err = 0
      num_done = 0
      for reco, ref_text in read_text(args.ref_in_file):
          try:
              if reco not in hyp_lines:
                  num_err += 1
                  raise Warning("Could not find recording {0} "
                                "in hypothesis {1}".format(
                                    reco, args.hyp_in_file.name))
                  continue
  
              if args.hyp_format == "CTM":
                  hyp_array = [x[2] for x in hyp_lines[reco]]
              else:
                  hyp_array = hyp_lines[reco]
  
              if args.reco2file_and_channel is None:
                  reco2file_and_channel[reco] = (reco, "1")
  
              logger.debug("Running Smith-Waterman alignment for %s", reco)
  
              output, score = smith_waterman_alignment(
                  ref_text, hyp_array, eps_symbol=args.eps_symbol,
                  similarity_score_function=similarity_score_function,
                  del_score=del_score, ins_score=ins_score,
                  align_full_hyp=args.align_full_hyp)
  
              if args.hyp_format == "CTM":
                  ctm_edits = get_ctm_edits(output, hyp_lines[reco],
                                            eps_symbol=args.eps_symbol,
                                            oov_word=args.oov_word,
                                            symbol_table=symbol_table)
                  for line in ctm_edits:
                      ctm_line = list(reco2file_and_channel[reco])
                      ctm_line.extend(line)
                      print(ctm_line_to_string(ctm_line),
                            file=args.alignment_out_file)
              else:
                  print_alignment(
                      reco, output, out_file_handle=args.alignment_out_file)
              num_done += 1
          except:
              logger.error("Alignment failed for recording {0} "
                           "with ref = {1} and hyp = {2}".format(
                               reco, " ".join(ref_text),
                               " ".join(hyp_array)))
              raise
  
      logger.info("Processed %d recordings; failed with %d", num_done, num_err)
  
      if num_done == 0:
          raise RuntimeError("Processed 0 recordings.")
  
  
  def main():
      args = get_args()
  
      try:
          run(args)
      except Exception:
          logger.error("Failed to align ref and hypotheses; "
                       "got exception ", exc_info=True)
          raise SystemExit(1)
      finally:
          if args.reco2file_and_channel is not None:
              args.reco2file_and_channel.close()
          args.ref_in_file.close()
          args.hyp_in_file.close()
          args.alignment_out_file.close()
  
  
  if __name__ == '__main__':
      main()