text: # Copyright 2020 Alibaba Group Holding Limited. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # =============================================================================
  """Classes used to construct EgoTensor encoders."""
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  try:
    # https://www.tensorflow.org/guide/migrate
    import tensorflow.compat.v1 as tf
    tf.disable_v2_behavior()
  except ImportError:
    import tensorflow as tf
  
  from graphlearn.python.model.base_encoder import BaseGraphEncoder
  from graphlearn.python.model.tf.utils.offsets_to_segment_ids import offsets_to_segment_ids
  
  EMB_PT_SIZE = 128 * 1024
  
  
  class EgoGraphEncoder(BaseGraphEncoder):
    """Encoder of densor EgoTensor.
  
    Args:
      feature_encoders: A list of feature encoders. Length: len(nbr_num_list) + 1
      conv_layers: A list of convolutional layers. Length: len(nbr_num_list).
      nbr_num_list: A list contains the neighbors num to be sampled of each hop.
      dropout: Dropout rate.
    """
  
    def __init__(self,
                 feature_encoders,
                 conv_layers,
                 nbr_num_list=None,
                 dropout=0.0):
      self._feature_encoders = feature_encoders
      self._conv_layers = conv_layers
      self._nbr_num_list = nbr_num_list
      self._depth = len(conv_layers)
      self._dropout = dropout
  
      assert self._depth + 1 == len(feature_encoders)
      if nbr_num_list is not None:
        assert self._depth == len(nbr_num_list)
  
  
    def _organize_samples(self, ego_tensor):
      """Organizes all necessary nodes' attributes into a list.
      Args:
        ego_tensor: an EgoTensor instance.
      """
      self._samples = []
      self._samples.append([ego_tensor.src.continuous_attrs,
                            ego_tensor.src.categorical_attrs])
      for neigh in ego_tensor.hops:
        #TODO(baole): support edges.
        self._samples.append([neigh.nodes.continuous_attrs,
                              neigh.nodes.categorical_attrs])
  
    def _forward(self, inputs):
      """ A common routine to aggregate information on graph layer by layer.
  
      For example, given a graph inputs below:
      1
      | \
      2  3
      |   \
      4    5
      if we want to update node 1's embedding based on its neighbors,
      we need two convolution layers:
      conv layer 1: update node 2 with node 4 and itself: h1[2] <= (h0[2],h0[4]),
                    update node 3 with node 5 and itself: h1[3] <= (h0[3],h0[5]),
                    update node 1 with node 2,3 and itself: h1[1] <= (h0[1],h0[2],h0[3])
      conv layer 2: update node 1 with updated node 2,3 and itself:
                    h2[1] <= (h1[1],h1[2],h1[3])
  
      Args:
        inputs: A list of input embeddings
      Returns:
        updated node embedding
      """
      hiddens = inputs
      for layer_idx in range(len(self._conv_layers)):
        tmp_vecs = []
        for hop in range(self._depth - layer_idx):
          src_vecs = hiddens[hop]
          neigh_vecs = tf.reshape(hiddens[hop + 1],
                                  [-1, self._nbr_num_list[hop],
                                   src_vecs.get_shape().as_list()[-1]])
          h = self._conv_layers[layer_idx].forward(src_vecs, neigh_vecs)
          if self._dropout != 0.0:
            tmp_vecs.append(tf.nn.dropout(h, 1 - self._dropout))
          else:
            tmp_vecs.append(h)
        hiddens = tmp_vecs
  
      return hiddens[0]
  
    def encode(self, ego_tensor):
      self._organize_samples(ego_tensor)
      # encode features.
      hiddens = []
      for idx in range(self._depth + 1):
        hiddens.append(self._feature_encoders[idx].encode(self._samples[idx]))
  
      return self._forward(hiddens)
  
  
  class SparseEgoGraphEncoder(EgoGraphEncoder):
    """Encoder for sparse EgoTensor.
  
    Sparse graph is generated by full neighbor sampler, which returns a batch of
    nodes with different numbers of neighbors.
  
    Args:
      feature_encoders: A list of feature encoders with length (nbr_num_list + 1)
      conv_layers: A list of convolutional layer with length nbr_num_list.
      dropout: Dropout rate.
    """
  
    def __init__(self,
                 feature_encoders,
                 conv_layers,
                 dropout=0.0):
      super(SparseEgoGraphEncoder, self).__init__(feature_encoders,
                                                  conv_layers,
                                                  dropout=dropout)
  
    def _organize_samples(self, ego_tensor):
      self._samples = []
      self._segment_ids = []
      self._samples.append([ego_tensor.src.continuous_attrs,
                            ego_tensor.src.categorical_attrs])
      for neigh in ego_tensor.hops:
        #TODO(baole): support edges.
        self._samples.append([neigh.nodes.continuous_attrs,
                              neigh.nodes.categorical_attrs])
        self._segment_ids.append(offsets_to_segment_ids(neigh.nodes.offsets))
  
    def _forward(self, inputs):
      hiddens = inputs
  
      for layer_idx in range(len(self._conv_layers)):
        tmp_vecs = []
        for hop in range(self._depth - layer_idx):
          src_vecs = hiddens[hop]
          neigh_vecs = hiddens[hop + 1]
          segment_ids = self._segment_ids[hop]
          h = self._conv_layers[layer_idx].forward(src_vecs,
                                                   neigh_vecs,
                                                   segment_ids)
          if self._dropout != 0.0:
            tmp_vecs.append(tf.nn.dropout(h, 1 - self._dropout))
          else:
            tmp_vecs.append(h)
        hiddens = tmp_vecs
  
      return hiddens[0]
  
  
  class LookupEncoder(BaseGraphEncoder):
    """Embedding Lookup Encoder.
  
    Args:
      name: The name of encoder.
      num: Total number of nodes or edges to be encoded.
      dim: Embedding dimension.
      init: Tensorflow initializer object.
      str2hash: Set it to True if need string to hash conversion.
      ps_hosts: ps_hosts used for TF distributed training.
      use_edge: Set it to True if use edge.
    """
    def __init__(self,
                 num,
                 dim,
                 init=None,
                 str2hash=True,
                 ps_hosts=None,
                 use_edge=False,
                 name=''):
      self._num = num
      self._dim = dim
      self._str2hash = str2hash
      self._name = name
      self._use_edge = use_edge
      self._initializer = init
  
      ps_num = 1
      if ps_hosts:
        ps_num = len(ps_hosts.split(","))
  
      emb_partitioner = \
        tf.min_max_variable_partitioner(max_partitions=ps_num,
                                        min_slice_size=EMB_PT_SIZE)
  
      with tf.variable_scope(self._name + 'lookup_embedding',
                             reuse=tf.AUTO_REUSE,
                             partitioner=emb_partitioner):
        self._emb_table = tf.get_variable("emb_lookup_table",
                                          [self._num, self._dim],
                                          initializer=self._initializer,
                                          partitioner=emb_partitioner)
  
      with tf.variable_scope(self._name + 'lookup_bias', reuse=tf.AUTO_REUSE,
                             partitioner=emb_partitioner):
        self._bias_table = tf.get_variable("bias_lookup_table",
                                           [self._num],
                                           partitioner=emb_partitioner,
                                           initializer=tf.zeros_initializer(),
                                           trainable=False)
    def encode(self, ego_tensor):
      if not self._use_edge:
        ids = ego_tensor.src.ids
      else:
        # TODO: replace this hack with edge ids.
        ids = tf.cast(ego_tensor.src.continuous_attrs, dtype=tf.int64)
  
      if self._str2hash:
        index = tf.as_string(ids)
        ids = tf.string_to_hash_bucket_fast(index,
                                            self._num,
                                            name=self._name +
                                            'str_to_hash_bucket_op')
  
      emb = tf.nn.embedding_lookup(self._emb_table,
                                   ids,
                                   name=self._name + 'ids_embedding_lookup_op')
      emb = tf.reshape(emb, [-1, self._dim])
  
      return emb
  
    @property
    def num(self):
      return self._num
  
    @property
    def emb_table(self):
      return self._emb_table
  
    @property
    def bias_table(self):
      return self._bias_table
  