# Copyright 2016    Johns Hopkins University (Dan Povey)
  #           2016    Vijayaditya Peddinti
  # Apache 2.0.
  
  """ This module contains the top level xconfig parsing functions.
  """
  
  from __future__ import print_function
  
  import logging
  import sys
  import libs.nnet3.xconfig.layers as xlayers
  import libs.nnet3.xconfig.utils as xutils
  
  import libs.common as common_lib
  
  
  # We have to modify this dictionary when adding new layers
  config_to_layer = {
          'input' : xlayers.XconfigInputLayer,
          'output' : xlayers.XconfigTrivialOutputLayer,
          'output-layer' : xlayers.XconfigOutputLayer,
          'relu-layer' : xlayers.XconfigBasicLayer,
          'relu-renorm-layer' : xlayers.XconfigBasicLayer,
          'relu-batchnorm-dropout-layer' : xlayers.XconfigBasicLayer,
          'relu-dropout-layer': xlayers.XconfigBasicLayer,
          'relu-batchnorm-layer' : xlayers.XconfigBasicLayer,
          'relu-batchnorm-so-layer' : xlayers.XconfigBasicLayer,
          'batchnorm-so-relu-layer' : xlayers.XconfigBasicLayer,
          'batchnorm-layer' : xlayers.XconfigBasicLayer,
          'sigmoid-layer' : xlayers.XconfigBasicLayer,
          'tanh-layer' : xlayers.XconfigBasicLayer,
          'fixed-affine-layer' : xlayers.XconfigFixedAffineLayer,
          'idct-layer' : xlayers.XconfigIdctLayer,
          'affine-layer' : xlayers.XconfigAffineLayer,
          'lstm-layer' : xlayers.XconfigLstmLayer,
          'lstmp-layer' : xlayers.XconfigLstmpLayer,
          'lstmp-batchnorm-layer' : xlayers.XconfigLstmpLayer,
          'fast-lstm-layer' : xlayers.XconfigFastLstmLayer,
          'fast-lstm-batchnorm-layer' : xlayers.XconfigFastLstmLayer,
          'fast-lstmp-layer' : xlayers.XconfigFastLstmpLayer,
          'fast-lstmp-batchnorm-layer' : xlayers.XconfigFastLstmpLayer,
          'lstmb-layer' : xlayers.XconfigLstmbLayer,
          'stats-layer': xlayers.XconfigStatsLayer,
          'relu-conv-layer': xlayers.XconfigConvLayer,
          'conv-layer': xlayers.XconfigConvLayer,
          'conv-relu-layer': xlayers.XconfigConvLayer,
          'conv-renorm-layer': xlayers.XconfigConvLayer,
          'relu-conv-renorm-layer': xlayers.XconfigConvLayer,
          'batchnorm-conv-layer': xlayers.XconfigConvLayer,
          'conv-relu-renorm-layer': xlayers.XconfigConvLayer,
          'batchnorm-conv-relu-layer': xlayers.XconfigConvLayer,
          'relu-batchnorm-conv-layer': xlayers.XconfigConvLayer,
          'relu-batchnorm-noconv-layer': xlayers.XconfigConvLayer,
          'relu-noconv-layer': xlayers.XconfigConvLayer,
          'conv-relu-batchnorm-layer': xlayers.XconfigConvLayer,
          'conv-relu-batchnorm-so-layer': xlayers.XconfigConvLayer,
          'conv-relu-batchnorm-dropout-layer': xlayers.XconfigConvLayer,
          'conv-relu-dropout-layer': xlayers.XconfigConvLayer,
          'res-block': xlayers.XconfigResBlock,
          'res2-block': xlayers.XconfigRes2Block,
          'channel-average-layer': xlayers.ChannelAverageLayer,
          'attention-renorm-layer': xlayers.XconfigAttentionLayer,
          'attention-relu-renorm-layer': xlayers.XconfigAttentionLayer,
          'attention-relu-batchnorm-layer': xlayers.XconfigAttentionLayer,
          'relu-renorm-attention-layer': xlayers.XconfigAttentionLayer,
          'gru-layer' : xlayers.XconfigGruLayer,
          'pgru-layer' : xlayers.XconfigPgruLayer,
          'opgru-layer' : xlayers.XconfigOpgruLayer,
          'norm-pgru-layer' : xlayers.XconfigNormPgruLayer,
          'norm-opgru-layer' : xlayers.XconfigNormOpgruLayer,
          'fast-gru-layer' : xlayers.XconfigFastGruLayer,
          'fast-pgru-layer' : xlayers.XconfigFastPgruLayer,
          'fast-norm-pgru-layer' : xlayers.XconfigFastNormPgruLayer,
          'fast-opgru-layer' : xlayers.XconfigFastOpgruLayer,
          'fast-norm-opgru-layer' : xlayers.XconfigFastNormOpgruLayer,
          'tdnnf-layer': xlayers.XconfigTdnnfLayer,
          'prefinal-layer': xlayers.XconfigPrefinalLayer,
          'spec-augment-layer': xlayers.XconfigSpecAugmentLayer,
          'renorm-component': xlayers.XconfigRenormComponent,
          'batchnorm-component': xlayers.XconfigBatchnormComponent,
          'no-op-component': xlayers.XconfigNoOpComponent,
          'linear-component': xlayers.XconfigLinearComponent,
          'affine-component': xlayers.XconfigAffineComponent,
          'scale-component':  xlayers.XconfigPerElementScaleComponent,
          'dim-range-component': xlayers.XconfigDimRangeComponent,
          'offset-component':  xlayers.XconfigPerElementOffsetComponent,
          'combine-feature-maps-layer': xlayers.XconfigCombineFeatureMapsLayer
  }
  
  # Turn a config line and a list of previous layers into
  # either an object representing that line of the config file; or None
  # if the line was empty after removing comments.
  # 'prev_layers' is a list of objects corresponding to preceding layers of the
  # config file.
  def xconfig_line_to_object(config_line, prev_layers = None):
      try:
          x  = xutils.parse_config_line(config_line)
          if x is None:
              return None
          (first_token, key_to_value) = x
          if not first_token in config_to_layer:
              raise RuntimeError("No such layer type '{0}'".format(first_token))
          return config_to_layer[first_token](first_token, key_to_value, prev_layers)
      except Exception:
          logging.error(
              "***Exception caught while parsing the following xconfig line:
  "
              "*** {0}".format(config_line))
          raise
  
  
  def get_model_component_info(model_filename):
      """
      This function reads existing model (*.raw or *.mdl) and returns array
      of XconfigExistingLayer one per {input,output}-node or component-node
      with same 'name' used in the raw model and 'dim' equal to 'output-dim'
      for component-node and 'dim' for {input,output}-node.
  
      e.g. layer in *.mdl -> corresponding 'XconfigExistingLayer' layer
           'input-node name=ivector dim=100' ->
           'existing name=ivector dim=100'
           'component-node name=tdnn1.affine ... input-dim=1000 '
           'output-dim=500' ->
           'existing name=tdnn1.affine dim=500'
      """
  
      all_layers = []
      try:
          f = open(model_filename, 'r')
      except Exception as e:
          sys.exit("{0}: error reading model file '{1}'".format(sys.argv[0],
                                                                model_filename,
                                                                repr(e)))
  
      # use nnet3-info to get component names in the model.
      out = common_lib.get_command_stdout("""nnet3-info {0} | grep '\-node' """
                                          """ """.format(model_filename))
  
      # out contains all {output, input, component}-nodes used in model_filename
      # It can parse lines in out like:
      # i.e. input-node name=input dim=40
      #   component-node name=tdnn1.affine component=tdnn1.affine input=lda
      #   input-dim=300 output-dim=512
      layer_names = []
      key_to_value = dict()
      for line in out.split("
  "):
          parts = line.split(" ")
          dim = -1
          for  field in parts:
              key_value = field.split("=")
              if len(key_value) == 2:
                  key = key_value[0]
                  value = key_value[1]
                  if key == "name":           # name=**
                      layer_name = value
                  elif key == "dim":          # for input-node
                      dim = int(value)
                  elif key == "output-dim":   # for component-node
                      dim = int(value)
  
          if layer_name is not None and layer_name not in layer_names:
              layer_names.append(layer_name)
              key_to_value['name'] = layer_name
              assert(dim != -1)
              key_to_value['dim'] = dim
              all_layers.append(xlayers.XconfigExistingLayer('existing', key_to_value, all_layers))
      if len(all_layers) == 0:
          raise RuntimeError("{0}: model filename '{1}' is empty.".format(
              sys.argv[0], model_filename))
      f.close()
      return all_layers
  
  
  # This function reads xconfig file and returns it as a list of layers
  # (usually we use the variable name 'all_layers' elsewhere for this).
  # It will die if the xconfig file is empty or if there was
  # some error parsing it.
  # 'existing_layers' contains some layers of type 'existing' (layers which are not really
  # layers but are actual component node names from an existing neural net model
  # and created using get_model_component_info function).
  # 'existing' layers can be used as input to component-nodes in layers of xconfig file.
  def read_xconfig_file(xconfig_filename, existing_layers=None):
      if existing_layers is None:
          existing_layers = []
      try:
          f = open(xconfig_filename, 'r')
      except Exception as e:
          sys.exit("{0}: error reading xconfig file '{1}'; error was {2}".format(
              sys.argv[0], xconfig_filename, repr(e)))
      all_layers = []
      while True:
          line = f.readline()
          if line == '':
              break
          # the next call will raise an easy-to-understand exception if
          # it fails.
          this_layer = xconfig_line_to_object(line, existing_layers)
          if this_layer is None:
              continue  # line was blank after removing comments.
          all_layers.append(this_layer)
          existing_layers.append(this_layer)
      if len(all_layers) == 0:
          raise RuntimeError("{0}: xconfig file '{1}' is empty".format(
              sys.argv[0], xconfig_filename))
      f.close()
      return all_layers