
# -*- coding: utf-8 -*-
"""
Created on Wed Oct 13 14:47:13 2021

@author: huzongxiang
"""

import tensorflow as tf
from tensorflow.keras import layers


class PartitionPadding(layers.Layer):
    def __init__(self, batch_size, **kwargs):
        super().__init__(**kwargs)
        self.batch_size = batch_size


    def call(self, inputs):
        features, graph_indices = inputs

        # Obtain subgraphs
        features = tf.dynamic_partition(
            features, graph_indices, self.batch_size
        )

        # Pad and stack subgraphs
        num_features = [tf.shape(f)[0] for f in features]
        max_num = tf.reduce_max(num_features)
        features_padded = tf.stack(
            [
                tf.pad(f, [(0, max_num - n), (0, 0)])
                for f, n in zip(features, num_features)
            ],
            axis=0,
        )
        
        # Remove empty subgraphs (usually for last batch)
        nonempty_examples = tf.where(tf.reduce_sum(features_padded, (1, 2)) != 0)
        nonempty_examples = tf.squeeze(nonempty_examples, axis=-1)

        features_batch =  tf.gather(features_padded, nonempty_examples, axis=0)

        return features_batch
    
    
    def get_config(self):
        config = super().get_config()
        config.update({"batch": self.batch_size})
        return config
    

class PartitionPaddingPair(layers.Layer):
    def __init__(self, batch_size, **kwargs):
        super().__init__(**kwargs)
        self.batch_size = batch_size


    def call(self, inputs):
        features, graph_indices = inputs

        # Obtain subgraphs
        features = tf.dynamic_partition(
            features, graph_indices, self.batch_size
        )

        # Pad and stack subgraphs
        num_features = [tf.shape(f)[0] for f in features]
        max_num = tf.reduce_max(num_features)
        features_padded = tf.stack(
            [
                tf.pad(f, [(0, max_num - n), (0, 0)])
                for f, n in zip(features, num_features)
            ],
            axis=0,
        )
       
        
        # Remove empty subgraphs (usually for last batch)
        nonempty_examples = tf.unique(graph_indices)[0]

        features_batch =  tf.gather(features_padded, nonempty_examples, axis=0)

        return features_batch
        

    def get_config(self):
        config = super().get_config()
        config.update({"batch_size": self.batch_size})
        return config
[((23, 19, 25, 9), 'tensorflow.dynamic_partition', 'tf.dynamic_partition', ({(24, 12, 24, 20): 'features', (24, 22, 24, 35): 'graph_indices', (24, 37, 24, 52): 'self.batch_size'}, {}), '(features, graph_indices, self.batch_size)', True, 'import tensorflow as tf\n'), ((29, 18, 29, 45), 'tensorflow.reduce_max', 'tf.reduce_max', ({(29, 32, 29, 44): 'num_features'}, {}), '(num_features)', True, 'import tensorflow as tf\n'), ((40, 28, 40, 66), 'tensorflow.squeeze', 'tf.squeeze', (), '', True, 'import tensorflow as tf\n'), ((42, 26, 42, 79), 'tensorflow.gather', 'tf.gather', (), '', True, 'import tensorflow as tf\n'), ((63, 19, 65, 9), 'tensorflow.dynamic_partition', 'tf.dynamic_partition', ({(64, 12, 64, 20): 'features', (64, 22, 64, 35): 'graph_indices', (64, 37, 64, 52): 'self.batch_size'}, {}), '(features, graph_indices, self.batch_size)', True, 'import tensorflow as tf\n'), ((69, 18, 69, 45), 'tensorflow.reduce_max', 'tf.reduce_max', ({(69, 32, 69, 44): 'num_features'}, {}), '(num_features)', True, 'import tensorflow as tf\n'), ((82, 26, 82, 79), 'tensorflow.gather', 'tf.gather', (), '', True, 'import tensorflow as tf\n'), ((80, 28, 80, 52), 'tensorflow.unique', 'tf.unique', ({(80, 38, 80, 51): 'graph_indices'}, {}), '(graph_indices)', True, 'import tensorflow as tf\n'), ((28, 24, 28, 35), 'tensorflow.shape', 'tf.shape', ({(28, 33, 28, 34): 'f'}, {}), '(f)', True, 'import tensorflow as tf\n'), ((32, 16, 32, 53), 'tensorflow.pad', 'tf.pad', ({(32, 23, 32, 24): 'f', (32, 26, 32, 52): '[(0, max_num - n), (0, 0)]'}, {}), '(f, [(0, max_num - n), (0, 0)])', True, 'import tensorflow as tf\n'), ((39, 37, 39, 75), 'tensorflow.reduce_sum', 'tf.reduce_sum', ({(39, 51, 39, 66): 'features_padded', (39, 68, 39, 74): '(1, 2)'}, {}), '(features_padded, (1, 2))', True, 'import tensorflow as tf\n'), ((68, 24, 68, 35), 'tensorflow.shape', 'tf.shape', ({(68, 33, 68, 34): 'f'}, {}), '(f)', True, 'import tensorflow as tf\n'), ((72, 16, 72, 53), 'tensorflow.pad', 'tf.pad', ({(72, 23, 72, 24): 'f', (72, 26, 72, 52): '[(0, max_num - n), (0, 0)]'}, {}), '(f, [(0, max_num - n), (0, 0)])', True, 'import tensorflow as tf\n')]