class SoftMax(Module):
    def __init__(self):
         super().__init__()

    def updateOutput(self, input):
        exps = np.exp(input - input.max(axis=1, keepdims=True))
        self.output = exps / np.sum(exps, axis=1, keepdims=True)
        return self.output

    def updateGradInput(self, input, gradOutput):
        self.gradInput = self.output * (1-self.output) * gradOutput \
                       - self.output * np.sum(self.output * gradOutput, axis=1, keepdims=True) \
                       + self.output * self.output * gradOutput
        return self.gradInput

    def __repr__(self):
        return "SoftMax"

class Crossentropy(Criterion):
    def __init__(self):
        super().__init__()

    def updateOutput(self, input, target):
        eps = 1e-15
        input_clamp = np.clip(input, eps, 1 - eps)
        self.output = - np.sum(np.log(input) * target) / input.shape[0]
        return self.output

    def updateGradInput(self, input, target):
        input_clamp = np.maximum(1e-15, np.minimum(input, 1 - 1e-15) )
        self.gradInput =  - (target / input) / input.shape[0]
        return self.gradInput

    def __repr__(self):
        return "ClassNLLCriterion"