import tensorflow as tf
from keras import layers
from keras import backend as K


class LayerNormalizerZeroOne(layers.Layer):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    # def build(self, input_shape):
    #     output_dim = input_shape[-1]
    #     self.kernel = self.add_weight(
    #         shape=(output_dim * 2, output_dim),
    #         initializer=self.initializer,
    #         name="kernel",
    #         trainable=True,
    #     )

    def call(self, inputs, **kwargs):
        # print(inputs.shape, 'inputs shape')
        # print(inputs, 'inputs')
        # print('-----------------------------')

        input_max = K.max(inputs, axis=[1, 2], keepdims=True)

        # print(input_max, 'input_max')
        # print(input_max.shape, 'input_max shape')
        # print('-----------------------------')

        outputs = tf.math.divide_no_nan(inputs, input_max)
        # outputs = 0 if max == 0 else inputs / input_max

        # print(outputs, 'outputs')

        return outputs


# def test_layer():
    # import tensorflow as tf
    # from bsmu.retinal_fundus.models.layers import LayerNormalizerZeroOne
    #
    # tf.enable_eager_execution()
    #
    # # x = tf.zeros(shape=(3, 2, 2, 1))
    # x = tf.random.normal(shape=(3, 2, 2, 1))
    #
    # # v = tf.Variable(x)
    # # v[0, 0, 0, 0].assign(0)
    # # v[0, 0, 1, 0].assign(0)
    # # v[0, 1, 0, 0].assign(0)
    # # v[0, 1, 1, 0].assign(0)
    #
    # y = LayerNormalizerZeroOne()(x)
