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


class MaskedSelect(Module):

    def __init__(self):
        super(MaskedSelect, self).__init__()
        self._maskIndices = torch.LongTensor()
        self._maskIndexBuffer = torch.LongTensor()
        self._maskIndexBufferCPU = torch.FloatTensor()
        self._gradBuffer = torch.Tensor()
        self._gradMask = torch.ByteTensor()

    def updateOutput(self, input):
        input, mask = input
        torch.masked_select(input, mask, out=self.output)
        return self.output

    def updateGradInput(self, input, gradOutput):
        input, mask = input
        if input.type() == 'torch.cuda.FloatTensor':
            torch.arange(0, mask.nelement(), out=self._maskIndexBufferCPU).resize_(mask.size())
            self._maskIndexBuffer.resize_(self._maskIndexBufferCPU.size()).copy_(self._maskIndexBufferCPU)
        else:
            torch.arange(0, mask.nelement(), out=self._maskIndexBuffer).resize_(mask.size())

        torch.masked_select(self._maskIndexBuffer, mask, out=self._maskIndices)
        self._gradBuffer.resize_(input.nelement()).zero_()
        self._gradBuffer.scatter_(0, self._maskIndices, gradOutput)
        self._gradBuffer.resize_(input.size())
        self.gradInput = [self._gradBuffer, self._gradMask.resize_(mask.size()).fill_(0)]
        return self.gradInput

    def type(self, type=None, tensorCache=None):
        if type is None:
            return self._type

        self._gradBuffer = self._gradBuffer.type(type)
        self.gradInput = self.gradInput.type(type)
        self.output = self.output.type(type)

        # These casts apply when switching between cuda/non-cuda types
        if type != 'torch.cuda.FloatTensor':
            self._maskIndexBuffer = self._maskIndexBuffer.long()
            self._maskIndices = self._maskIndices.long()
            self._gradMask = self._gradMask.byte()
        else:
            self._maskIndexBuffer = self._maskIndexBuffer.cuda()
            self._maskIndices = self._maskIndices.cuda()
            self._gradMask = self._gradMask.cuda()

        self._type = type
        return self

    def clearState(self):
        return clear(self, ['output',
                            'gradInput',
                            '_maskIndexBuffer',
                            '_maskIndexBufferCPU',
                            '_maskIndices',
                            '_gradBuffer',
                            '_gradMask'])
