text: #!/usr/bin/env/python
  """
  Usage:
      chem_tensorflow_async.py [options]
  
  Options:
      -h --help                Show this screen.
      --config-file FILE       Hyperparameter configuration file path (in JSON format).
      --config CONFIG          Hyperparameter configuration dictionary (in JSON format).
      --log_dir DIR            Log dir name.
      --data_dir DIR           Data dir name.
      --restore FILE           File to restore weights from.
      --freeze-graph-model     Freeze weights of graph model components.
  """
  from typing import List, Tuple, Dict, Sequence, Any
  
  from docopt import docopt
  from collections import defaultdict
  import numpy as np
  import tensorflow as tf
  import sys, traceback
  import pdb
  
  from chem_tensorflow import ChemModel
  from utils import glorot_init, SMALL_NUMBER
  
  
  def bfs_visit(outgoing_edges: Dict[int, Sequence[int]], node_depths: Dict[int, int], v: int, depth: int):
      # Already seen, skip:
      if v in node_depths:
          return
      node_depths[v] = depth
      for (_, __, w) in outgoing_edges[v]:
          bfs_visit(outgoing_edges, node_depths, w, depth + 1)
  
  
  class AsyncGGNNChemModel(ChemModel):
      def __init__(self, args):
          super().__init__(args)
  
      @classmethod
      def default_params(cls):
          params = dict(super().default_params())
          params.update({
              'num_nodes': 100000,
              'use_edge_bias': False,
  
              'propagation_rounds': 4,  # Has to be an even number
              'propagation_substeps': 15,
  
              'graph_rnn_cell': 'GRU',  # GRU or RNN
              'graph_rnn_activation': 'tanh',  # tanh, ReLU
              'graph_state_dropout_keep_prob': 1.,
  
              'task_sample_ratios': {},
          })
          return params
  
      def prepare_specific_graph_model(self) -> None:
          h_dim = self.params['hidden_size']
          self.placeholders['initial_node_representation'] = tf.placeholder(tf.float32, [None, h_dim],
                                                                            name='node_features')
  
          # Initial nodes I_{r}: Node IDs that will have no incoming edges in round r.
          self.placeholders['initial_nodes'] = [tf.placeholder(tf.int32, [None], name="initial_nodes_round%i" % prop_round)
                                                for prop_round in range(self.params['propagation_rounds'])]
  
          # Sending nodes S_{r,s,e}: Source node ids of edges propagating in step s of round r.
          # Restrictions: If v in S_{r,s,e}, then v in R_{r,s'} for s' < s or v in I_{r}
          self.placeholders['sending_nodes'] = [[[tf.placeholder(tf.int32,
                                                                 [None],
                                                                 name="sending_nodes_round%i_step%i_edgetyp%i" % (prop_round, step, edge_typ))
                                                  for edge_typ in range(self.num_edge_types)]
                                                 for step in range(self.params['propagation_substeps'])]
                                                for prop_round in range(self.params['propagation_rounds'])]
  
          # Normalised edge target nodes T_{r,s}: Targets of edges propagating in step s of round r, normalised to a
          # continuous range starting from 0. This is used for aggregating messages from the sending nodes.
          self.placeholders['msg_targets'] = [[tf.placeholder(tf.int32,
                                                              [None],
                                                              name="msg_targets_nodes_round%i_step%i" % (prop_round, step))
                                               for step in range(self.params['propagation_substeps'])]
                                              for prop_round in range(self.params['propagation_rounds'])]
  
  
          # Receiving nodes R_{r,s}: Target node ids of aggregated messages in propagation step s of round r.
          # Restrictions: If v in R_{r,s}, v not in R_{r,s'} for all s' != s and v not in I_{r}
          self.placeholders['receiving_nodes'] = [[tf.placeholder(tf.int32,
                                                                  [None],
                                                                  name="receiving_nodes_round%i_step%i" % (prop_round, step))
                                                   for step in range(self.params['propagation_substeps'])]
                                                  for prop_round in range(self.params['propagation_rounds'])]
  
          # Number of receiving nodes N_{r,s}
          # Restrictions: N_{r,s} = len(R_{r,s})
          self.placeholders['receiving_node_num'] = [tf.placeholder(tf.int32,
                                                                    [self.params['propagation_substeps']],
                                                                    name="receiving_nodes_num_round%i" % (prop_round,))
                                                     for prop_round in range(self.params['propagation_rounds'])]
  
          self.placeholders['graph_nodes_list'] = tf.placeholder(tf.int64, [None, 2], name='graph_nodes_list')
          self.placeholders['graph_state_keep_prob'] = tf.placeholder(tf.float32, None, name='graph_state_keep_prob')
  
          activation_name = self.params['graph_rnn_activation'].lower()
          if activation_name == 'tanh':
              activation_fun = tf.nn.tanh
          elif activation_name == 'relu':
              activation_fun = tf.nn.relu
          else:
              raise Exception("Unknown activation function type '%s'." % activation_name)
  
          # Generate per-layer values for edge weights, biases and gated units. If we tie them, they are just copies:
          self.weights['edge_weights'] = [tf.Variable(glorot_init([h_dim, h_dim]), name='gnn_edge_weights_typ%i' % e_typ)
                                          for e_typ in range(self.num_edge_types)]
  
          if self.params['use_edge_bias']:
              self.weights['edge_biases'] = [tf.Variable(np.zeros([h_dim], dtype=np.float32), name='gnn_edge_biases_typ%i' % e_typ)
                                             for e_typ in range(self.num_edge_types)]
  
          cell_type = self.params['graph_rnn_cell'].lower()
          if cell_type == 'gru':
              cell = tf.nn.rnn_cell.GRUCell(h_dim, activation=activation_fun)
          elif cell_type == 'rnn':
              cell = tf.nn.rnn_cell.BasicRNNCell(h_dim, activation=activation_fun)
          else:
              raise Exception("Unknown RNN cell type '%s'." % cell_type)
          cell = tf.nn.rnn_cell.DropoutWrapper(cell,
                                               state_keep_prob=self.placeholders['graph_state_keep_prob'])
          self.weights['rnn_cells'] = cell
  
      def compute_final_node_representations(self) -> tf.Tensor:
          cur_node_states = self.placeholders['initial_node_representation']
  
          for prop_round in range(self.params['propagation_rounds']):
              with tf.variable_scope('prop_round%i' % (prop_round,)):
                  # ---- Declare and fill tensor arrays used in tf.while_loop:
                  sending_nodes_ta = tf.TensorArray(tf.int32,
                                                    infer_shape=False,
                                                    element_shape=[None],
                                                    size=self.params['propagation_substeps'] * self.num_edge_types,
                                                    name='sending_nodes')
                  msg_targets_ta = tf.TensorArray(tf.int32,
                                                  infer_shape=False,
                                                  element_shape=[None],
                                                  size=self.params['propagation_substeps'],
                                                  name='msg_targets')
                  receiving_nodes_ta = tf.TensorArray(tf.int32,
                                                      infer_shape=False,
                                                      element_shape=[None],
                                                      size=self.params['propagation_substeps'],
                                                      clear_after_read=False,
                                                      name='receiving_nodes')
                  receiving_node_num_ta = tf.TensorArray(tf.int32,
                                                         infer_shape=False,
                                                         element_shape=[],
                                                         size=self.params['propagation_substeps'],
                                                         name='receiving_nodes_num')
  
                  for step in range(self.params['propagation_substeps']):
                      for edge_typ in range(self.num_edge_types):
                          sending_nodes_ta = sending_nodes_ta.write(step * self.num_edge_types + edge_typ,
                                                                    self.placeholders['sending_nodes'][prop_round][step][edge_typ])
                      msg_targets_ta = msg_targets_ta.write(step, self.placeholders['msg_targets'][prop_round][step])
                      receiving_nodes_ta = receiving_nodes_ta.write(step, self.placeholders['receiving_nodes'][prop_round][step])
                  receiving_node_num_ta = receiving_node_num_ta.unstack(self.placeholders['receiving_node_num'][prop_round])
  
                  new_node_states_ta = tf.TensorArray(tf.float32,
                                                      infer_shape=False,
                                                      element_shape=[self.params['hidden_size']],
                                                      size=tf.shape(cur_node_states)[0],
                                                      clear_after_read=False,
                                                      name='new_node_states')
  
                  # ---- Actual propagation schedule implementation:
                  # Initialize the initial nodes with their state from last round:
                  new_node_states_ta = new_node_states_ta.scatter(self.placeholders['initial_nodes'][prop_round],
                                                                  tf.gather(cur_node_states, self.placeholders['initial_nodes'][prop_round]))
  
                  def do_substep(substep_id, new_node_states_ta):
                      # For each edge active in this substep, pull source state and transform:
                      sent_messages = []
                      for edge_typ in range(self.num_edge_types):
                          sending_states = new_node_states_ta.gather(sending_nodes_ta.read(substep_id * self.num_edge_types + edge_typ))
                          messages = tf.matmul(sending_states, self.weights['edge_weights'][edge_typ])
                          if self.params['use_edge_bias']:
                              messages += self.weights['edge_biases'][edge_typ]
                          sent_messages.append(messages)
  
                      # Stack all edge messages and aggregate as sum for each receiving node:
                      sent_messages = tf.concat(sent_messages, axis=0)
                      aggregated_received_messages = tf.unsorted_segment_sum(sent_messages,
                                                                             msg_targets_ta.read(substep_id),
                                                                             receiving_node_num_ta.read(substep_id))
  
                      # Collect old states for receiving nodes, and combine in RNN cell with incoming message
                      substep_receiving_nodes = receiving_nodes_ta.read(substep_id)
                      old_receiving_node_states = tf.gather(cur_node_states, substep_receiving_nodes)
                      aggregated_received_messages.set_shape([None, self.params['hidden_size']])
                      old_receiving_node_states.set_shape([None, self.params['hidden_size']])
                      substep_new_node_states = self.weights['rnn_cells'](aggregated_received_messages,
                                                                          old_receiving_node_states)[1]
  
                      # Write updated states back:
                      new_node_states_ta = new_node_states_ta.scatter(substep_receiving_nodes, substep_new_node_states)
                      return (substep_id + 1, new_node_states_ta)
  
                  def is_done(substep_id, new_node_states_ta_unused):
                      return tf.logical_and(substep_id < self.params['propagation_substeps'],
                                            tf.greater(tf.shape(receiving_nodes_ta.read(substep_id))[0], 0))
  
                  _, new_node_states_ta = tf.while_loop(cond=is_done,
                                                        body=do_substep,
                                                        loop_vars=[tf.constant(0), new_node_states_ta]
                                                       )
  
                  cur_node_states = new_node_states_ta.stack(name="state_stack_round%i" % (prop_round,))
  
          return cur_node_states
  
      def gated_regression(self, last_h, regression_gate, regression_transform):
          # last_h: [v x h]
          gate_input = tf.concat([last_h, self.placeholders['initial_node_representation']], axis=-1)  # [v x 2h]
          gated_outputs = tf.nn.sigmoid(regression_gate(gate_input)) * regression_transform(last_h)  # [v x 1]
  
          # Sum up all nodes per-graph
          num_nodes = tf.shape(gate_input, out_type=tf.int64)[0]
          graph_nodes = tf.SparseTensor(indices=self.placeholders['graph_nodes_list'],
                                        values=tf.ones_like(self.placeholders['graph_nodes_list'][:, 0],
                                                            dtype=tf.float32),
                                        dense_shape=[self.placeholders['num_graphs'], num_nodes])  # [g x v]
          return tf.squeeze(tf.sparse_tensor_dense_matmul(graph_nodes, gated_outputs), axis=[-1])  # [g]
  
      # ----- Data preprocessing and chunking into minibatches:
      def process_raw_graphs(self, raw_data: Sequence[Any], is_training_data: bool) -> Any:
          processed_graphs = []
          for d in raw_data:
              prop_schedules = self.__graph_to_propagation_schedules(d['graph'])
              processed_graphs.append({"init": d["node_features"],
                                       "prop_schedules": prop_schedules,
                                       "target_values": [d["targets"][task_id][0] for task_id in self.params['task_ids']]})
  
          if is_training_data:
              np.random.shuffle(processed_graphs)
              for task_id in self.params['task_ids']:
                  task_sample_ratio = self.params['task_sample_ratios'].get(str(task_id))
                  if task_sample_ratio is not None:
                      ex_to_sample = int(len(processed_graphs) * task_sample_ratio)
                      for ex_id in range(ex_to_sample, len(processed_graphs)):
                          processed_graphs[ex_id]['target_values'][task_id] = None
  
          return processed_graphs
  
      def __tensorise_edge_sequence(self, edges)\
              -> Tuple[np.ndarray, List[List[np.ndarray]], List[List[np.ndarray]], List[np.ndarray]]:
          sending_nodes = []  # type: List[List[np.ndarray]]
          msg_targets = []  # type: List[List[np.ndarray]]
          receiving_nodes = []  # type: List[np.ndarray]
          all_nodes = set()
          for step_edges in edges:
              msg_targets_uniq = set(w for (_, __, w) in step_edges)
              recv_nodes = list(sorted(msg_targets_uniq))
              recv_nodes_to_uniq_id = {v: i for (i, v) in enumerate(recv_nodes)}
  
              sending_nodes_in_step = []
              msg_targets_in_step = []
              for target_e_typ in range(self.num_edge_types):
                  sending_nodes_in_step.append(np.array([v for (v, e_typ, _) in step_edges if e_typ == target_e_typ], dtype=np.int32))
                  msg_targets_in_step.append(np.array([recv_nodes_to_uniq_id[w] for (_, e_typ, w) in step_edges if e_typ == target_e_typ], dtype=np.int32))
              msg_targets.append(msg_targets_in_step)
              sending_nodes.append(sending_nodes_in_step)
              receiving_nodes.append(np.array(recv_nodes, dtype=np.int32))
              all_nodes.update(v for (v, _, __) in step_edges)
              all_nodes.update(w for (_, __, w) in step_edges)
          
          all_updated_nodes = set()
          all_updated_nodes.update(v for step_receiving_nodes in receiving_nodes
                                     for v in step_receiving_nodes)
          initial_nodes = list(sorted(all_nodes - all_updated_nodes))
  
          #initialised_nodes = set()
          #initialised_nodes.update(initial_nodes)
          #for step in range(len(receiving_nodes)):
          #    sent_nodes = set()
          #    for edge_typ in range(self.num_edge_types):
          #        sent_nodes.update(sending_nodes[step][edge_typ])
          #    for v in sent_nodes:
          #        assert v in initialised_nodes
          #
          #    for v in receiving_nodes[step]:
          #        assert v not in initialised_nodes
          #    initialised_nodes.update(receiving_nodes[step])
  
          return (np.array(initial_nodes, dtype=np.int32), sending_nodes, msg_targets, receiving_nodes)
  
      def __graph_to_propagation_schedules(self, graph)\
              -> List[Tuple[np.ndarray, List[List[np.ndarray]], List[List[np.ndarray]], List[np.ndarray]]]:
          num_incoming_edges = defaultdict(lambda: 0)
          outgoing_edges = defaultdict(lambda: [])
          # Compute number of incoming edges per node, and build adjacency lists:
          for (v, typ, w) in graph:
              num_incoming_edges[v] += 1
              num_incoming_edges[w] += 1
              edge_bwd_typ = typ if self.params['tie_fwd_bkwd'] else self.num_edge_types + typ
              outgoing_edges[v].append((v, typ, w))
              outgoing_edges[w].append((w, edge_bwd_typ, v))
  
          # Sort them, pick node with lowest number of incoming edges:
          tensorised_prop_schedules = []
          for prop_round in range(int(self.params['propagation_rounds'] / 2)):
              dag_seed = min(num_incoming_edges.items(), key=lambda t: t[1])[prop_round]
              node_depths = {}
              bfs_visit(outgoing_edges, node_depths, dag_seed, 0)
  
              # Now split edges into forward/backward sets, by using their depths.
              # Intuitively, a node with depth h will get updated in step h.
              max_depth = max(node_depths.values())
              assert(max_depth <= self.params['propagation_substeps'])
              fwd_pass_edges = [[] for _ in range(max_depth)]
              bwd_pass_edges = [[] for _ in range(max_depth)]
              for (v, typ, w) in graph:
                  edge_bwd_type = typ if self.params['tie_fwd_bkwd'] else self.num_edge_types + typ
                  v_depth = node_depths[v]
                  w_depth = node_depths[w]
                  if v_depth < w_depth:  # "Forward": We are going up in depth:
                      fwd_pass_edges[w_depth - 1].append((v, typ, w))
                      bwd_pass_edges[-v_depth - 1].append((w, edge_bwd_type, v))
                  elif w_depth < v_depth:  # "Backward": We are going down in depth
                      fwd_pass_edges[v_depth - 1].append((w, edge_bwd_type, v))
                      bwd_pass_edges[-w_depth - 1].append((v, typ, w))
                  else:
                      # We ignore self-loops:
                      assert v == w
  
              tensorised_prop_schedules.append(self.__tensorise_edge_sequence(fwd_pass_edges))
              tensorised_prop_schedules.append(self.__tensorise_edge_sequence(bwd_pass_edges))
  
          return tensorised_prop_schedules
  
      def make_minibatch_iterator(self, data: Any, is_training: bool):
          """Create minibatches by flattening graphs into a single one with multiple disconnected components."""
          if is_training:
              np.random.shuffle(data)
          dropout_keep_prob = self.params['graph_state_dropout_keep_prob'] if is_training else 1.
  
          # Pack until we cannot fit more graphs in the batch
          num_graphs = 0
          while num_graphs < len(data):
              num_graphs_in_batch = 0
              batch_node_features = []
              batch_target_task_values = []
              batch_target_task_mask = []
              batch_graph_nodes_list = []
              node_offset = 0
  
              # Collect all indices; we'll strip out the batch dimension with a np.concatenate along that axis at the end:
              batch_initial_nodes = [[] for _ in range(self.params['propagation_rounds'])
                                    ]  # type: List[List[np.ndarray]]              # (prop_round, batch, None)
              batch_sending_nodes = [[[[] for _ in range(self.num_edge_types)]
                                      for _ in range(self.params['propagation_substeps'])]
                                     for _ in range(self.params['propagation_rounds'])
                                    ]  # type: List[List[List[List[np.ndarray]]]]  # (prop_round, step, edge_typ, batch, None)
              batch_msg_targets = [[[[] for _ in range(self.num_edge_types)]
                                    for _ in range(self.params['propagation_substeps'])]
                                   for _ in range(self.params['propagation_rounds'])
                                  ]  # type: List[List[List[List[np.ndarray]]]]    # (prop_round, step, edge_typ, batch, None)
              batch_receiving_nodes = [[[] for _ in range(self.params['propagation_substeps'])]
                                       for _ in range(self.params['propagation_rounds'])
                                      ]  # type: List[List[List[np.ndarray]]]      # (prop_round, step, batch, None)
              batch_receiving_node_num = [[0 for _ in range(self.params['propagation_substeps'])]
                                          for _ in range(self.params['propagation_rounds'])
                                         ]  # type: List[List[int]]                # (prop_round, step)
  
              msg_target_offsets = [[[0 for _ in range(self.num_edge_types)]
                                     for _ in range(self.params['propagation_substeps'])]
                                    for _ in range(self.params['propagation_rounds'])]
  
              while num_graphs < len(data) and node_offset + len(data[num_graphs]['init']) < self.params['num_nodes']:
                  cur_graph = data[num_graphs]
                  num_nodes_in_graph = len(cur_graph['init'])
                  padded_features = np.pad(cur_graph['init'],
                                           ((0, 0), (0, self.params['hidden_size'] - self.annotation_size)),
                                           'constant')
                  batch_node_features.extend(padded_features)
                  batch_graph_nodes_list.extend(
                      (num_graphs_in_batch, node_offset + i) for i in range(num_nodes_in_graph))
  
                  # Combine the different propagation schedules:
                  for prop_round in range(self.params['propagation_rounds']):
                      cur_prop_schedule = cur_graph['prop_schedules'][prop_round]
                      (graph_initial_nodes,
                       graph_sending_nodes,
                       graph_msg_targets,
                       graph_recv_nodes) = cur_prop_schedule
                      batch_initial_nodes[prop_round].append(graph_initial_nodes + node_offset)
                      for step in range(self.params['propagation_substeps']):
                          # Stop if we don't have that many steps:
                          if step >= len(graph_sending_nodes):
                              break
  
                          for e_typ in range(self.num_edge_types):
                              batch_sending_nodes[prop_round][step][e_typ].append(graph_sending_nodes[step][e_typ] + node_offset)
                              batch_msg_targets[prop_round][step][e_typ].append(graph_msg_targets[step][e_typ] + msg_target_offsets[prop_round][step][e_typ])
                              if len(graph_msg_targets[step][e_typ]) > 0:
                                  msg_target_offsets[prop_round][step][e_typ] += max(graph_msg_targets[step][e_typ]) + 1  # ... 0-based indexing!
                          batch_receiving_nodes[prop_round][step].append(graph_recv_nodes[step] + node_offset)
                          batch_receiving_node_num[prop_round][step] += len(graph_recv_nodes[step])
  
                  target_task_values = []
                  target_task_mask = []
                  for target_val in cur_graph['target_values']:
                      if target_val is None:  # This is one of the examples we didn't sample...
                          target_task_values.append(0.)
                          target_task_mask.append(0.)
                      else:
                          target_task_values.append(target_val)
                          target_task_mask.append(1.)
                  batch_target_task_values.append(target_task_values)
                  batch_target_task_mask.append(target_task_mask)
                  num_graphs += 1
                  num_graphs_in_batch += 1
                  node_offset += num_nodes_in_graph
  
              batch_feed_dict = {
                  self.placeholders['initial_node_representation']: np.array(batch_node_features),
                  self.placeholders['graph_nodes_list']: np.array(batch_graph_nodes_list, dtype=np.int32),
                  self.placeholders['target_values']: np.transpose(batch_target_task_values, axes=[1,0]),
                  self.placeholders['target_mask']: np.transpose(batch_target_task_mask, axes=[1, 0]),
                  self.placeholders['num_graphs']: num_graphs_in_batch,
                  self.placeholders['graph_state_keep_prob']: dropout_keep_prob,
              }
  
              for prop_round in range(self.params['propagation_rounds']):
                  batch_feed_dict[self.placeholders['initial_nodes'][prop_round]] = \
                      np.concatenate(batch_initial_nodes[prop_round], axis=0)
                  for step in range(self.params['propagation_substeps']):
                      msg_targets = []
                      for edge_typ in range(self.num_edge_types):
                          raw_senders = batch_sending_nodes[prop_round][step][edge_typ]
                          batch_feed_dict[self.placeholders['sending_nodes'][prop_round][step][edge_typ]] = \
                              np.concatenate(raw_senders, axis=0) if len(raw_senders) > 0 else np.empty(shape=(0,),
                                                                                                        dtype=np.int32)
                          raw_targets = batch_msg_targets[prop_round][step][edge_typ]
                          msg_targets.extend(np.concatenate(raw_targets, axis=0) if len(raw_targets) > 0 else np.empty(shape=(0,),
                                                                                                                       dtype=np.int32))
  
                      batch_feed_dict[self.placeholders['msg_targets'][prop_round][step]] = \
                          np.array(msg_targets, dtype=np.int32)
                      raw_recvs = batch_receiving_nodes[prop_round][step]
                      batch_feed_dict[self.placeholders['receiving_nodes'][prop_round][step]] = \
                          np.concatenate(raw_recvs, axis=0) if len(raw_recvs) > 0 else np.empty(shape=(0,),
                                                                                                dtype=np.int32)
                  batch_feed_dict[self.placeholders['receiving_node_num'][prop_round]] = \
                      np.array(batch_receiving_node_num[prop_round])
  
              #self.check_batch_invariants(batch_feed_dict)
              yield batch_feed_dict
  
  
      def check_batch_invariants(self, batch_feed_dict):
          for prop_round in range(self.params['propagation_rounds']):
              initialised_nodes = set()
              initialised_nodes.update(batch_feed_dict[self.placeholders['initial_nodes'][prop_round]])
              for step in range(self.params['propagation_substeps']):
                  sending_nodes = set()
                  for edge_typ in range(self.num_edge_types):
                      sending_nodes.update(batch_feed_dict[self.placeholders['sending_nodes'][prop_round][step][edge_typ]])
                  for v in sending_nodes:
                      assert v in initialised_nodes
  
                  recv_nodes = batch_feed_dict[self.placeholders['receiving_nodes'][prop_round][step]]
                  for v in recv_nodes:
                      assert v not in initialised_nodes
                  initialised_nodes.update(recv_nodes)
  
  def main():
      args = docopt(__doc__)
      try:
          model = AsyncGGNNChemModel(args)
          model.train()
      except:
          typ, value, tb = sys.exc_info()
          traceback.print_exc()
          pdb.post_mortem(tb)
  
  
  if __name__ == "__main__":
      main()
  