# Copyright 2016 Vijayaditya Peddinti.
  #           2016 Vimal Manohar
  #           2017 Johns Hopkins University (author: Daniel Povey)
  # Apache 2.0
  
  """ This module contains several utility functions and classes that are
  commonly used in many kaldi python scripts.
  """
  
  from __future__ import print_function
  from __future__ import division
  import argparse
  import logging
  import math
  import os
  import subprocess
  import sys
  import threading
  
  try:
      import thread as thread_module
  except:
      import _thread as thread_module
  
  logger = logging.getLogger(__name__)
  logger.addHandler(logging.NullHandler())
  
  
  def send_mail(message, subject, email_id):
      try:
          subprocess.Popen(
              'echo "{message}" | mail -s "{subject}" {email}'.format(
                  message=message,
                  subject=subject,
                  email=email_id), shell=True)
      except Exception as e:
          logger.info("Unable to send mail due to error:
   {error}".format(
                          error=str(e)))
          pass
  
  
  def str_to_bool(value):
      if value == "true":
          return True
      elif value == "false":
          return False
      else:
          raise ValueError
  
  
  class StrToBoolAction(argparse.Action):
      """ A custom action to convert bools from shell format i.e., true/false
          to python format i.e., True/False """
  
      def __call__(self, parser, namespace, values, option_string=None):
          try:
              setattr(namespace, self.dest, str_to_bool(values))
          except ValueError:
              raise Exception(
                  "Unknown value {0} for --{1}".format(values, self.dest))
  
  
  class NullstrToNoneAction(argparse.Action):
      """ A custom action to convert empty strings passed by shell to None in
      python. This is necessary as shell scripts print null strings when a
      variable is not specified. We could use the more apt None in python. """
  
      def __call__(self, parser, namespace, values, option_string=None):
          if values.strip() == "":
              setattr(namespace, self.dest, None)
          else:
              setattr(namespace, self.dest, values)
  
  
  class smart_open(object):
      """
      This class is designed to be used with the "with" construct in python
      to open files. It is similar to the python open() function, but
      treats the input "-" specially to return either sys.stdout or sys.stdin
      depending on whether the mode is "w" or "r".
  
      e.g.: with smart_open(filename, 'w') as fh:
              print ("foo", file=fh)
      """
      def __init__(self, filename, mode="r"):
          self.filename = filename
          self.mode = mode
          assert self.mode == "w" or self.mode == "r"
  
      def __enter__(self):
          if self.filename == "-" and self.mode == "w":
              self.file_handle = sys.stdout
          elif self.filename == "-" and self.mode == "r":
              self.file_handle = sys.stdin
          else:
              self.file_handle = open(self.filename, self.mode)
          return self.file_handle
  
      def __exit__(self, *args):
          if self.filename != "-":
              self.file_handle.close()
  
  
  class smart_open(object):
      """
      This class is designed to be used with the "with" construct in python
      to open files. It is similar to the python open() function, but
      treats the input "-" specially to return either sys.stdout or sys.stdin
      depending on whether the mode is "w" or "r".
  
      e.g.: with smart_open(filename, 'w') as fh:
              print ("foo", file=fh)
      """
      def __init__(self, filename, mode="r"):
          self.filename = filename
          self.mode = mode
          assert self.mode == "w" or self.mode == "r"
  
      def __enter__(self):
          if self.filename == "-" and self.mode == "w":
              self.file_handle = sys.stdout
          elif self.filename == "-" and self.mode == "r":
              self.file_handle = sys.stdin
          else:
              self.file_handle = open(self.filename, self.mode)
          return self.file_handle
  
      def __exit__(self, *args):
          if self.filename != "-":
              self.file_handle.close()
  
  
  def check_if_cuda_compiled():
      p = subprocess.Popen("cuda-compiled")
      p.communicate()
      if p.returncode == 1:
          return False
      else:
          return True
  
  
  def execute_command(command):
      """ Runs a kaldi job in the foreground and waits for it to complete; raises an
          exception if its return status is nonzero.  The command is executed in
          'shell' mode so 'command' can involve things like pipes.  Often,
          'command' will start with 'run.pl' or 'queue.pl'.  The stdout and stderr
          are merged with the calling process's stdout and stderr so they will
          appear on the screen.
  
          See also: get_command_stdout, background_command
      """
      p = subprocess.Popen(command, shell=True)
      p.communicate()
      if p.returncode is not 0:
          raise Exception("Command exited with status {0}: {1}".format(
                  p.returncode, command))
  
  
  def get_command_stdout(command, require_zero_status = True):
      """ Executes a command and returns its stdout output as a string.  The
          command is executed with shell=True, so it may contain pipes and
          other shell constructs.
  
          If require_zero_stats is True, this function will raise an exception if
          the command has nonzero exit status.  If False, it just prints a warning
          if the exit status is nonzero.
  
          See also: execute_command, background_command
      """
      p = subprocess.Popen(command, shell=True,
                           stdout=subprocess.PIPE)
  
      stdout = p.communicate()[0]
      if p.returncode is not 0:
          output = "Command exited with status {0}: {1}".format(
              p.returncode, command)
          if require_zero_status:
              raise Exception(output)
          else:
              logger.warning(output)
      return stdout if type(stdout) is str else stdout.decode()
  
  
  
  
  def wait_for_background_commands():
      """ This waits for all threads to exit.  You will often want to
          run this at the end of programs that have launched background
          threads, so that the program will wait for its child processes
          to terminate before it dies."""
      for t in threading.enumerate():
          if not t == threading.current_thread():
              t.join()
  
  def background_command(command, require_zero_status = False):
      """Executes a command in a separate thread, like running with '&' in the shell.
         If you want the program to die if the command eventually returns with
         nonzero status, then set require_zero_status to True.  'command' will be
         executed in 'shell' mode, so it's OK for it to contain pipes and other
         shell constructs.
  
         This function returns the Thread object created, just in case you want
         to wait for that specific command to finish.  For example, you could do:
               thread = background_command('foo | bar')
               # do something else while waiting for it to finish
               thread.join()
  
         See also:
           - wait_for_background_commands(), which can be used
             at the end of the program to wait for all these commands to terminate.
           - execute_command() and get_command_stdout(), which allow you to
             execute commands in the foreground.
  
      """
  
      p = subprocess.Popen(command, shell=True)
      thread = threading.Thread(target=background_command_waiter,
                                args=(command, p, require_zero_status))
      thread.daemon=True  # make sure it exits if main thread is terminated
                          # abnormally.
      thread.start()
      return thread
  
  
  def background_command_waiter(command, popen_object, require_zero_status):
      """ This is the function that is called from background_command, in
          a separate thread."""
  
      popen_object.communicate()
      if popen_object.returncode is not 0:
          str = "Command exited with status {0}: {1}".format(
              popen_object.returncode, command)
          if require_zero_status:
              logger.error(str)
              # thread.interrupt_main() sends a KeyboardInterrupt to the main
              # thread, which will generally terminate the program.
              thread_module.interrupt_main()
          else:
              logger.warning(str)
  
  
  def get_number_of_leaves_from_tree(alidir):
      stdout = get_command_stdout(
          "tree-info {0}/tree 2>/dev/null | grep num-pdfs".format(alidir))
      parts = stdout.split()
      assert(parts[0] == "num-pdfs")
      num_leaves = int(parts[1])
      if num_leaves == 0:
          raise Exception("Number of leaves is 0")
      return num_leaves
  
  
  def get_number_of_leaves_from_model(dir):
      stdout = get_command_stdout(
          "am-info {0}/final.mdl 2>/dev/null | grep -w pdfs".format(dir))
      parts = stdout.split()
      # number of pdfs 7115
      assert(' '.join(parts[0:3]) == "number of pdfs")
      num_leaves = int(parts[3])
      if num_leaves == 0:
          raise Exception("Number of leaves is 0")
      return num_leaves
  
  
  def get_number_of_jobs(alidir):
      try:
          num_jobs = int(open('{0}/num_jobs'.format(alidir)).readline().strip())
      except (IOError, ValueError) as e:
          logger.error("Exception while reading the "
                       "number of alignment jobs: ", exc_info=True)
          raise SystemExit(1)
      return num_jobs
  
  
  def get_ivector_dim(ivector_dir=None):
      if ivector_dir is None:
          return 0
      stdout_val = get_command_stdout(
          "feat-to-dim --print-args=false "
          "scp:{dir}/ivector_online.scp -".format(dir=ivector_dir))
      ivector_dim = int(stdout_val)
      return ivector_dim
  
  def get_ivector_extractor_id(ivector_dir=None):
      if ivector_dir is None:
          return None
      stdout_val = get_command_stdout(
          "steps/nnet2/get_ivector_id.sh {dir}".format(dir=ivector_dir))
  
      if (stdout_val.strip() == "") or (stdout_val is None):
          return None
  
      return stdout_val.strip()
  
  def get_feat_dim(feat_dir):
      if feat_dir is None:
          return 0
      stdout_val = get_command_stdout(
          "feat-to-dim --print-args=false "
          "scp:{data}/feats.scp -".format(data=feat_dir))
      feat_dim = int(stdout_val)
      return feat_dim
  
  
  def get_feat_dim_from_scp(feat_scp):
      stdout_val = get_command_stdout(
          "feat-to-dim --print-args=false "
          "scp:{feat_scp} -".format(feat_scp=feat_scp))
      feat_dim = int(stdout_val)
      return feat_dim
  
  
  def read_kaldi_matrix(matrix_file):
      """This function reads a kaldi matrix stored in text format from
      'matrix_file' and stores it as a list of rows, where each row is a list.
      """
      try:
          lines = [x.split() for x in open(matrix_file).readlines()]
          first_field = lines[0][0]
          last_field = lines[-1][-1]
          lines[0] = lines[0][1:]
          lines[-1] = lines[-1][:-1]
          if not (first_field == "[" and last_field == "]"):
              raise Exception(
                  "Kaldi matrix file has incorrect format, "
                  "only text format matrix files can be read by this script")
          for i in range(len(lines)):
              lines[i] = [int(float(x)) for x in lines[i]]
          return lines
      except IOError:
          raise Exception("Error while reading the kaldi matrix file "
                          "{0}".format(matrix_file))
  
  
  def write_kaldi_matrix(output_file, matrix):
      """This function writes the matrix stored as a list of lists
      into 'output_file' in kaldi matrix text format.
      """
      with open(output_file, 'w') as f:
          f.write("[ ")
          num_rows = len(matrix)
          if num_rows == 0:
              raise Exception("Matrix is empty")
          num_cols = len(matrix[0])
  
          for row_index in range(len(matrix)):
              if num_cols != len(matrix[row_index]):
                  raise Exception("All the rows of a matrix are expected to "
                                  "have the same length")
              f.write(" ".join([str(x) for x in matrix[row_index]]))
              if row_index != num_rows - 1:
                  f.write("
  ")
          f.write(" ]")
  
  
  def write_matrix_ascii(file_or_fd, mat, key=None):
      """This function writes the matrix 'mat' stored as a list of lists
      in kaldi matrix text format.
      The destination can be a file or an opened file descriptor.
      If key is provided, then matrix is written to an archive with the 'key'
      as the index field.
      """
      try:
          fd = open(file_or_fd, 'w')
      except TypeError:
          # 'file_or_fd' is opened file descriptor,
          fd = file_or_fd
  
      try:
          if key is not None:
              print ("{0} [".format(key),
                     file=fd)  # ark-files have keys (utterance-id)
          else:
              print (" [", file=fd)
  
          num_cols = 0
          for i, row in enumerate(mat):
              line = ' '.join(["{0:f}".format(x) for x in row])
              if i == 0:
                  num_cols = len(row)
              elif len(row) != num_cols:
                  raise Exception("All the rows of a matrix are expected to "
                                  "have the same length")
  
              if i == len(mat) - 1:
                  line += " ]"
              print (line, file=fd)
      finally:
          if fd is not file_or_fd : fd.close()
  
  
  def read_matrix_ascii(file_or_fd):
      """This function reads a matrix in kaldi matrix text format
      and stores it as a list of lists.
      The input can be a file or an opened file descriptor.
      """
      try:
          fd = open(file_or_fd, 'r')
          fname = file_or_fd
      except TypeError:
          # 'file_or_fd' is opened file descriptor,
          fd = file_or_fd
          fname = file_or_fd.name
  
      first = fd.read(2)
      if first != ' [':
          logger.error(
              "Kaldi matrix file %s has incorrect format, "
              "only text format matrix files can be read by this script",
              fname)
          raise RuntimeError
  
      rows = []
      while True:
          line = fd.readline()
          if len(line) == 0:
              logger.error("Kaldi matrix file %s has incorrect format; "
                           "got EOF before end of matrix", fname)
          if len(line.strip()) == 0 : continue # skip empty line
          arr = line.strip().split()
          if arr[-1] != ']':
              rows.append([float(x) for x in arr])  # not last line
          else:
              rows.append([float(x) for x in arr[:-1]])  # lastline
              return rows
      if fd is not file_or_fd:
          fd.close()
  
  
  def read_key(fd):
    """ [str] = read_key(fd)
     Read the utterance-key from the opened ark/stream descriptor 'fd'.
    """
    str_ = ''
    while True:
      char = fd.read(1)
      if char == '':
          break
      if char == ' ':
          break
      str_ += char
    str_ = str_.strip()
    if str_ == '':
        return None   # end of file,
    return str_
  
  
  def read_mat_ark(file_or_fd):
      """This function reads a kaldi matrix archive in text format
      and yields a dictionary output indexed by the key (utterance-id).
      The input can be a file or an opened file descriptor.
  
      Example usage:
      mat_dict = { key: mat for key, mat in read_mat_ark(file) }
      """
      try:
          fd = open(file_or_fd, 'r')
          fname = file_or_fd
      except TypeError:
          # 'file_or_fd' is opened file descriptor,
          fd = file_or_fd
          fname = file_or_fd.name
  
      try:
          key = read_key(fd)
          while key:
            mat = read_matrix_ascii(fd)
            yield key, mat
            key = read_key(fd)
      finally:
          if fd is not file_or_fd:
              fd.close()
  
  
  def force_symlink(file1, file2):
      import errno
      try:
          os.symlink(file1, file2)
      except OSError as e:
          if e.errno == errno.EEXIST:
              os.remove(file2)
              os.symlink(file1, file2)
  
  
  def compute_lifter_coeffs(lifter, dim):
      coeffs = [0] * dim
      for i in range(0, dim):
          coeffs[i] = 1.0 + 0.5 * lifter * math.sin(math.pi * i / float(lifter))
  
      return coeffs
  
  
  def compute_idct_matrix(K, N, cepstral_lifter=0):
      matrix = [[0] * K for i in range(N)]
      # normalizer for X_0
      normalizer = math.sqrt(1.0 / float(N))
      for j in range(0, N):
          matrix[j][0] = normalizer
      # normalizer for other elements
      normalizer = math.sqrt(2.0 / float(N))
      for k in range(1, K):
          for n in range(0, N):
              matrix[n][
                  k] = normalizer * math.cos(math.pi / float(N) * (n + 0.5) * k)
  
      if cepstral_lifter != 0:
          lifter_coeffs = compute_lifter_coeffs(cepstral_lifter, K)
          for k in range(0, K):
              for n in range(0, N):
                  matrix[n][k] = float(matrix[n][k]) / lifter_coeffs[k]
  
      return matrix
  
  
  def write_idct_matrix(feat_dim, cepstral_lifter, file_path):
      # generate the IDCT matrix and write to the file
      idct_matrix = compute_idct_matrix(feat_dim, feat_dim, cepstral_lifter)
      # append a zero column to the matrix, this is the bias of the fixed affine
      # component
      for k in range(0, feat_dim):
          idct_matrix[k].append(0)
      write_kaldi_matrix(file_path, idct_matrix)