from .affine import Affine
from ... import describe
from ... import check
from ...describe import Dimension, Synapses, Biases
from ...check import has_shape


class ReLu(Affine):
    @check.arg(1, has_shape(('nB', 'nI')))
    def predict(self, input__BI):
        output__BO = Affine.predict(self, input__BI)
        output__BO = self.ops.relu(output__BO, inplace=False)
        return output__BO

    @check.arg(1, has_shape(('nB', 'nI')))
    def begin_update(self, input__BI, drop=0.0):
        output__BO, finish_affine = Affine.begin_update(self, input__BI, drop=0.)
        output__BO = self.ops.relu(output__BO)
        @check.arg(0, has_shape(('nB', 'nO')))
        def finish_update(gradient, sgd=None):
            gradient = self.ops.backprop_relu(gradient, output__BO)
            return finish_affine(gradient, sgd)
        dropped, bp_dropout = self.ops.dropout(output__BO, drop, inplace=False)
        return dropped, bp_dropout(finish_update)
