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


class SelectTable(Module):

    def __init__(self, index):
        super(SelectTable, self).__init__()
        self.index = index
        self.gradInput = []

    def updateOutput(self, input):
        # handle negative indices
        index = self.index if self.index >= 0 else input.size(self.dimension) + self.index
        assert len(input) > index
        self.output = input[index]
        return self.output

    def _zeroTableCopy(self, l1, l2):
        for i, v in enumerate(l2):
            if isinstance(v, list):
                if len(l1) > i:
                    l1[i] = self._zeroTableCopy(l1[i], l2[i])
                else:
                    l1.append(self._zeroTableCopy([], l2[i]))
            else:
                if i >= len(l1):
                    l1.append(v.new().resize_as_(v).zero_())
                else:
                    l1[i].resize_as_(v)
                    l1[i].zero_()
        del l1[len(l2):]
        return l1

    def updateGradInput(self, input, gradOutput):
        # make gradInput a zeroed copy of input
        self._zeroTableCopy(self.gradInput, input)
        # handle negative indices
        index = self.index if self.index >= 0 else input.size(self.dimension) + self.index
        # copy into gradInput[index] (necessary for variable sized inputs)
        assert self.gradInput[index] is not None
        recursiveCopy(self.gradInput[index], gradOutput)
        return self.gradInput

    def type(self, type, tensorCache=None):
        del self.gradInput[:]
        if isinstance(self.output, list):
            del self.output[:]
        return super(SelectTable, self).type(type, tensorCache)

    def __repr__(self):
        return super(SelectTable, self).__repr__() + '({})'.format(self.index)

    def clearState(self):
        clear(self, 'gradInput')
