text: # Copyright (C) 2016, 2017 University of Vienna
  # All rights reserved.
  # BSD license.
  # Author: Ali Baharev <ali.baharev@gmail.com>
  from __future__ import print_function, division
  from functools import partial
  from itertools import count
  
  from six import iteritems, itervalues
  from nx import OrderedDiGraph as nx_OrderedDiGraph, dfs_preorder_nodes
  
  from nl_interpreter import get_problem, interpret, Range, get_J_rowwise
  from py3compat import imap, izip
  
  #-------------------------------------------------------------------------------
  # Hopefully, we can remove this hack one day when nx.OrderedDigraph is ready.
  # Note that the code here won't have to be changed; we will be able to use
  # nx.OrderedDiGraph directly. Currently, we have to override the subgraph()
  # method. Networkx at:
  # ali@X230:~/nx/networkx/networkx$ git log -n 1
  # commit e43a7b08bd5ab2640b5a9c3350ed5355bdb82c65
  # Date:   Fri Jun 23 17:57:33 2017 +0200
  
  class OrderedDiGraph(nx_OrderedDiGraph):
      
      #@profile
      def subgraph(self, nbunch):
          bunch = self.nbunch_iter(nbunch)
          # create new graph and copy subgraph into it
          H = self.__class__()
          # copy node and attribute dictionaries
          for n in bunch:
              H.node[n]=self.node[n]
          # namespace shortcuts for speed
          H_succ=H.succ
          H_pred=H.pred
          self_succ=self.succ
          self_pred=self.pred
          # add nodes
          for n in H:
              H_succ[n]=H.adjlist_inner_dict_factory()
              H_pred[n]=H.adjlist_inner_dict_factory()
          # add successors
          for u in H_succ:
              Hnbrs=H_succ[u]
              for v,datadict in self_succ[u].items():
                  if v in H_succ:
                      Hnbrs[v]=datadict
          # add predecessors
          for u in H_pred:
              Hnbrs=H_pred[u]
              for v,datadict in self_pred[u].items():
                  if v in H_pred:
                      Hnbrs[v]=datadict
          H.graph=self.graph
          return H
  
  #-------------------------------------------------------------------------------
  
  class ntype:
      # these have AMPL names and AMPL indices; var and con have bounds
      var    = 'var'
      defvar = 'defvar'
      con    = 'con'
      # ---
      num    = 'num'
      # An operator not in the .nl file, but corresponding to the linear part
      lin_comb = 'lin_comb'
      # operators
      plus     = 'plus'
      minus    = 'minus'
      multiply = 'mult'  
      divide   = 'div'   
      negate   = 'neg'
      sumlist  = 'sum'   
      exp      = 'exp'   
      log      = 'log'   
      power    = 'pow'   
  
  
  OPERATOR = {
      'plus':  ntype.plus,
      'minus': ntype.minus,
      'mult':  ntype.multiply,
      'div':   ntype.divide,
      'neg':   ntype.negate,
      'sum':   ntype.sumlist,
      'exp':   ntype.exp,
      'log':   ntype.log,
      'pow':   ntype.power,
      # non-AMPL operator
      'lin_comb': ntype.lin_comb, 
  }
  
  OPS = set(itervalues(OPERATOR))
  
  #-------------------------------------------------------------------------------
  
  def main():    
      #probname = 'tunnelDiodesSum'
      probname = 'mss10TornScaled'
      g, _problem = create_dag(probname)
      print(g, g.number_of_nodes(), g.number_of_edges())
      for n, d in iteritems(g.node):
          print(n, d)
      #
      gen_ops = (n for n, d in iteritems(g.node) if d['kind'] in OPS)
      for op in gen_ops:
          print(op, '=', g.node[op]['kind'], ' '.join(g.predecessors(op)))
  
  def create_dag(problem_name):
      p = get_problem(problem_name)
      g = OrderedDiGraph()
      g.graph['name'] = problem_name
      add_all_vars(g, p.col_names)
      cons, def_vars = p.segments.cons, p.segments.def_vars
      C_names, V_names = p.row_names, p.def_var_names
      t_start = p.nl_header.n_vars + len(p.def_var_names) 
      t_counter, n_counter = count(start=t_start), count()
      for kind, index in p.segments.eval_order:
          if kind == 'C':
              add_constraint(g, index, cons, t_counter, n_counter, C_names)
          elif kind == 'V':
              add_defined_variable(g, index, def_vars, t_counter, n_counter, V_names)
      #
      Jrows = get_J_rowwise(p)
      check_sparsity_pattern(g, Jrows, p.nl_header.n_vars, p.nl_header.n_nzeros, 
                                                           p.segments.con_jacobian)
      return g, p
  
  def add_constraint(g, index, cons, t_counter, n_counter, C_names):
      nl_part, lin_part, con_range = cons[index]
      nl_last, lin_last = add(g, nl_part, lin_part, t_counter, n_counter)
      rng = add_range(g, con_range, n_counter)
      last = sum_up(g, t_counter, nl_last, lin_last, rng)
      add_last(g, last, 'C%d' % index, ntype.con, C_names[index])
  
  def add_defined_variable(g, index, def_vars, t_counter, n_counter, V_names):
      nl_part, lin_part = def_vars[index]
      nl_last, lin_last = add(g, nl_part, lin_part, t_counter, n_counter)
      last = sum_up(g, t_counter, nl_last, lin_last)
      add_last(g, last, 'v%d' % index, ntype.defvar,  V_names[index])
  
  def add_last(g, last, node_id, kind, name):
      add_node(g, node_id, kind, name=name)
      add_edge(g, last, node_id)    
  
  def sum_up(g, t_counter, nl_last, lin_last, rng=None):
      not_none = tuple(n for n in (nl_last, lin_last, rng) if n)
      if len(not_none) == 1:
          return not_none[0]
      return add_with_children(g, 't%d' % next(t_counter), ntype.sumlist, not_none)
  
  def add_range(g, con_range, n_counter):
      kind = con_range[0]
      assert kind == Range.eq, kind
      (value,) = con_range[1]
      if value == '0':
          return
      # con = rhs - range -> rhs + (-range); flipping the sign:
      value = '-' + value if value[0] != '-' else value[1:]
      return add_num_node(g, n_counter, value)
  
  def add(g, nl_part, lin_part,  t_counter, n_counter):
      nl_last  = add_nonlinear(g, nl_part, t_counter, n_counter)
      lin_last = add_linear(g, lin_part, t_counter)
      assert nl_last or lin_last
      return nl_last, lin_last 
  
  def add_linear(g, lin_part, t_counter):
      indices, coefficients = lin_part
      nzero_coeffs = tuple(c for c in coefficients if c != '0')
      if not nzero_coeffs:
          return
      children = ['v%d' % i for i, c in izip(indices, coefficients) if c != '0']
      return add_with_children(g, 't%d' % next(t_counter), ntype.lin_comb, 
                                                   children, coeffs=nzero_coeffs)
  
  def add_nonlinear(g, nl_part, t_counter, n_counter):
      body = list(interpret(nl_part, t_counter))
      if body == ['n0']:
          return    
      for lhs, op_name, args in body[:-1]:
          args = add_numbers(g, args, n_counter)
          add_with_children(g, lhs, OPERATOR[op_name], args)
      return add_if_num(g, n_counter, body[-1]) # the last node, value of nl_part    
  
  def add_numbers(g, args, n_counter):
      # replace n-1.23 with n42 where 42 is a unique id, else just return arg
      return tuple(imap(partial(add_if_num, g, n_counter), args))
  
  def add_if_num(g, n_counter, arg):
      return add_num_node(g, n_counter, arg[1:]) if arg[0] == 'n' else arg
  
  def add_num_node(g, n_counter, value):
      return add_node(g, 'n%d' % next(n_counter), ntype.num, value=value)
  
  def add_with_children(g, node_id, kind, children, **kwargs):
      # Adds the node_id, the children must already be in g
      add_node(g, node_id, kind, **kwargs)
      for child in children:
          add_edge(g, child, node_id)
      return node_id
  
  def add_all_vars(g, col_names):
      for i, name in enumerate(col_names):
          add_node(g, 'v%d' % i, ntype.var, name=name)
  
  #-------------------------------------------------------------------------------
  # Only these functions add nodes / edges to g through g.add_*
  
  def add_node(g, node_id, kind, **kwargs):
      assert node_id not in g, node_id
      g.add_node(node_id, kind=kind, **kwargs)
      return node_id
  
  def add_edge(g, src, dst):
      assert src in g, src
      assert dst in g, dst
      g.add_edge(src, dst)
  
  #-------------------------------------------------------------------------------
  
  def check_sparsity_pattern(g, Jrows, n_vars, n_zeros, k_segment):
      # Cross-check with the k_segment, compute column lengths from rowwise
      col_counts  = get_col_counts(Jrows, n_vars)
      col_lengths = get_intermediate_col_lengths(col_counts)
      assert col_lengths[-1] == n_zeros
      assert k_segment == col_lengths[:-1]
      # Cross-check against DFS from C_i, reaching the base variables
      g.reverse(copy=False)
      for i, cols in enumerate(Jrows):
          check_base_vars(g, n_vars, i, cols)
      g.reverse(copy=False)
  
  def check_base_vars(g, n_vars, i, cols):
      # Used to be  g.node[n]['index'] instead of int(n[1:])
      vrs = [int(n[1:]) for n in dfs_preorder_nodes(g, source='C%d' % i) if n[0] == 'v']
      indices = sorted(index for index in vrs if index < n_vars)
      assert indices == cols, (indices, cols)
  
  def get_col_counts(Jrows, n_vars):
      col_counts = [0]*n_vars
      for cols in Jrows:
          for j in cols:
              col_counts[j] += 1
      return col_counts
  
  def get_intermediate_col_lengths(col_count):
      lengths, accum = [], 0
      for count in col_count:
          accum += count
          lengths.append(accum)
      return lengths
  
  
  if __name__ == '__main__':
      main()
  