import torch
from .Module import Module
from .utils import clear


class SpatialCrossMapLRN(Module):

    def __init__(self, size, alpha=1e-4, beta=0.75, k=1):
        super(SpatialCrossMapLRN, self).__init__()

        self.size = size
        self.alpha = alpha
        self.beta = beta
        self.k = k
        self.scale = None
        self.paddedRatio = None
        self.accumRatio = None

    def updateOutput(self, input):
        assert input.dim() == 4

        if self.scale is None:
            self.scale = input.new()
        if input.type() == 'torch.cuda.FloatTensor':
            self._backend.SpatialCrossMapLRN_updateOutput(
                self._backend.library_state,
                input,
                self.output,
                self.scale,
                self.size,
                self.alpha,
                self.beta,
                self.k
            )
        else:
            batchSize = input.size(0)
            channels = input.size(1)
            inputHeight = input.size(2)
            inputWidth = input.size(3)

            self.output.resize_as_(input)
            self.scale.resize_as_(input)

            # use output storage as temporary buffer
            inputSquare = self.output
            torch.pow(input, 2, out=inputSquare)

            prePad = int((self.size - 1) / 2 + 1)
            prePadCrop = channels if prePad > channels else prePad

            scaleFirst = self.scale.select(1, 0)
            scaleFirst.zero_()
            # compute first feature map normalization
            for c in range(prePadCrop):
                scaleFirst.add_(inputSquare.select(1, c))

            # reuse computations for next feature maps normalization
            # by adding the next feature map and removing the previous
            for c in range(1, channels):
                scalePrevious = self.scale.select(1, c - 1)
                scaleCurrent = self.scale.select(1, c)
                scaleCurrent.copy_(scalePrevious)
                if c < channels - prePad + 1:
                    squareNext = inputSquare.select(1, c + prePad - 1)
                    scaleCurrent.add_(1, squareNext)

                if c > prePad:
                    squarePrevious = inputSquare.select(1, c - prePad)
                    scaleCurrent.add_(-1, squarePrevious)

            self.scale.mul_(self.alpha / self.size).add_(self.k)

            torch.pow(self.scale, -self.beta, out=self.output)
            self.output.mul_(input)

        return self.output

    def updateGradInput(self, input, gradOutput):
        assert input.dim() == 4

        if input.type() == 'torch.cuda.FloatTensor':
            self._backend.SpatialCrossMapLRN_updateGradInput(
                self._backend.library_state,
                input,
                gradOutput,
                self.gradInput,
                self.scale,
                self.output,
                self.size,
                self.alpha,
                self.beta,
                self.k
            )
        else:
            batchSize = input.size(0)
            channels = input.size(1)
            inputHeight = input.size(2)
            inputWidth = input.size(3)

            if self.paddedRatio is None:
                self.paddedRatio = input.new()
            if self.accumRatio is None:
                self.accumRatio = input.new()
            self.paddedRatio.resize_(channels + self.size - 1, inputHeight, inputWidth)
            self.accumRatio.resize_(inputHeight, inputWidth)

            cacheRatioValue = 2 * self.alpha * self.beta / self.size
            inversePrePad = int(self.size - (self.size - 1) / 2)

            self.gradInput.resize_as_(input)
            torch.pow(self.scale, -self.beta, out=self.gradInput).mul_(gradOutput)

            self.paddedRatio.zero_()
            paddedRatioCenter = self.paddedRatio.narrow(0, inversePrePad, channels)
            for n in range(batchSize):
                torch.mul(gradOutput[n], self.output[n], out=paddedRatioCenter)
                paddedRatioCenter.div_(self.scale[n])
                torch.sum(self.paddedRatio.narrow(0, 0, self.size - 1), 0, keepdim=False, out=self.accumRatio)
                for c in range(channels):
                    self.accumRatio.add_(self.paddedRatio[c + self.size - 1])
                    self.gradInput[n][c].addcmul_(-cacheRatioValue, input[n][c], self.accumRatio)
                    self.accumRatio.add_(-1, self.paddedRatio[c])

        return self.gradInput

    def clearState(self):
        clear(self, 'scale', 'paddedRatio', 'accumRatio')
        return super(SpatialCrossMapLRN, self).clearState()
