import torch
from .Module import Module
from .utils import clear, recursiveResizeAs


class MixtureTable(Module):

    def __init__(self, dim=1):
        super(MixtureTable, self).__init__()
        self.dim = dim
        self.size = torch.Size()
        self.size2 = torch.Size()
        self.batchSize = 0
        self.backwardSetup = False
        self.gradInput = []

        self._gaterView = None
        self._expert = None
        self._expertView = None
        self._sum = None
        self._expertView2 = None
        self._expert2 = None
        self.table = False

    def updateOutput(self, input):
        gaterInput, expertInputs = input

        # buffers
        if self._gaterView is None:
            self._gaterView = input[0].new()
        if self._expert is None:
            self._expert = input[0].new()
        if self._expertView is None:
            self._expertView = input[0].new()

        self.dimG = 1
        batchSize = gaterInput.size(0)

        if self.table or isinstance(expertInputs, list):
            self.table = True
            if gaterInput.size(self.dimG) != len(expertInputs):
                raise RuntimeError("Should be one gater output per expert")

            expertInput = expertInputs[0]
            if self.batchSize != batchSize:
                size = [1] * (expertInput.dim() + 1)
                if self.dimG > 0:
                    size[0] = gaterInput.size(0)
                size[self.dim] = gaterInput.size(self.dimG)
                self.size = torch.Size(size)
                self.output.resize_as_(expertInput)
                self.backwardSetup = False
                self.batchSize = batchSize

            self._gaterView = gaterInput.view(self.size)
            self.output.zero_()
            # multiply accumulate gater outputs by their commensurate expert
            for i, expertInput in enumerate(expertInputs):
                gate = self._gaterView.select(self.dim, i).expand_as(expertInput)
                self.output.addcmul_(expertInput, gate)
        else:
            if self.batchSize != batchSize:
                size = [1] * expertInputs.dim()
                if self.dimG > 0:
                    size[0] = gaterInput.size(0)
                size[self.dim] = gaterInput.size(self.dimG)
                self.size = torch.Size(size)
                self.output.resize_as_(expertInputs.select(self.dim, 0))
                self.batchSize = batchSize
                self.backwardSetup = False

            self._gaterView = gaterInput.view(self.size)
            torch.mul(self._gaterView.expand_as(expertInputs), expertInputs, out=self._expert)
            torch.sum(self._expert, self.dim, True, out=self.output)
            self.output.resize_as_(expertInputs.select(self.dim, 0))

        return self.output

    def updateGradInput(self, input, gradOutput):
        gaterInput, expertInputs = input
        recursiveResizeAs(self.gradInput, input)
        gaterGradInput, expertGradInputs = self.gradInput

        # buffers
        if self._sum is None:
            self._sum = input[0].new()
        if self._expertView2 is None:
            self._expertView2 = input[0].new()
        if self._expert2 is None:
            self._expert2 = input[0].new()

        if self.table:
            if not self.backwardSetup:
                for i, expertInput in enumerate(expertInputs):
                    expertGradInput = expertGradInputs[i] or expertInput.clone()
                    expertGradInput.resize_as_(expertInput)
                    expertGradInputs[i] = expertGradInput

                gaterGradInput.resize_as_(gaterInput)
                self.backwardSetup = True

            # like CMulTable, but with broadcasting
            for i, expertGradInput in enumerate(expertGradInputs):
                # gater updateGradInput
                torch.mul(gradOutput, expertInputs[i], out=self._expert)
                if self.dimG == 0:
                    self._expertView = self._expert.view(-1)
                else:
                    self._expertView = self._expert.view(gradOutput.size(0), -1)

                torch.sum(self._expertView, self.dimG, True, out=self._sum)
                if self.dimG == 0:
                    gaterGradInput[i] = self._sum.select(self.dimG, 0)
                else:
                    gaterGradInput.select(self.dimG, i).copy_(self._sum.select(self.dimG, 0))

                # expert updateGradInput
                gate = self._gaterView.select(self.dim, i).expand_as(expertGradInput)
                expertGradInput.mul_(gate, gradOutput)
        else:
            if not self.backwardSetup:
                size2 = list(expertInputs.size())
                size2[self.dim] = 1
                self.size2 = torch.Size(size2)
                gaterGradInput.resize_as_(gaterInput)
                self.backwardSetup = True

            # gater updateGradInput
            self._expertView = gradOutput.contiguous().view(torch.Size(self.size2))
            gradOutput = self._expertView.expand_as(expertInputs)
            torch.mul(gradOutput, expertInputs, out=self._expert)
            expert = self._expert.transpose(self.dim, self.dimG)
            if not expert.is_contiguous():
                self._expert2.resize_as_(expert)
                self._expert2.copy_(expert)
                expert = self._expert2
            if self.dimG == 0:
                self._expertView2 = expert.view(gaterInput.size(0), -1)
            else:
                self._expertView2 = expert.view(gaterInput.size(0), gaterInput.size(1), -1)

            torch.sum(self._expertView2, self.dimG + 1, True, out=gaterGradInput)
            gaterGradInput.resize_as_(gaterInput)

            # expert updateGradInput
            torch.mul(self._gaterView.expand_as(expertInputs), gradOutput, out=expertGradInputs)

        return self.gradInput

    def type(self, type, tensorCache=None):
        self._gaterView = None
        self._expert = None
        self._expertView = None
        self._sum = None
        self._expert2 = None
        self._expertView2 = None
        return super(MixtureTable, self).type(type, tensorCache)

    def clearState(self, ):
        clear(self, [
            '_gaterView',
            '_expert',
            '_expertView',
            '_sum',
            '_expert2',
            '_expertView2',
        ])
        return super(MixtureTable, self).clearState()
