import tensorflow as tf

class EpsilonLRP(GradientBasedMethod):
    eps = None

    def __init__(self, T, X, session, keras_learning_phase, epsilon=1e-4, Y_shape=None):
        assert epsilon > 0.0, 'LRP epsilon must be greater than zero'
        global eps
        eps = epsilon
        super(EpsilonLRP, self).__init__(T, X, session, keras_learning_phase, Y_shape)

    def get_symbolic_attribution(self):
        return [g * x for g, x in zip(
            tf.gradients(ys=self.T, xs=self.X),
            self.X if self.has_multiple_inputs else [self.X])]

    @classmethod
    def nonlinearity_grad_override(cls, op, grad):
        output = op.outputs[0]
        input = op.inputs[0]
        return grad * output / (input + eps *
                                tf.compat.v1.where(input >= 0, tf.ones_like(input), -1 * tf.ones_like(input)))


