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


class Min(Module):

    def __init__(self, dimension=0):
        super(Min, self).__init__()
        self.dimension = dimension
        self._output = None
        self._indices = None

    def _getPositiveDimension(self, input):
        dimension = self.dimension
        if dimension < 0:
            dimension = input.dim() + dimension

        return dimension

    def _lazyInit(self):
        if self._output is None:
            self._output = self.output.new()
        if self._indices is None:
            self._indices = \
                (torch.cuda.LongTensor() if self.output.type() == 'torch.cuda.FloatTensor'
                 else torch.LongTensor())

    def updateOutput(self, input):
        self._lazyInit()
        dimension = self._getPositiveDimension(input)
        torch.min(input, dimension, out=(self._output, self._indices), keepdim=True)
        if input.dim() > 1:
            self.output.set_(self._output.select(dimension, 0))
        else:
            self.output.set_(self._output)

        return self.output

    def updateGradInput(self, input, gradOutput):
        self._lazyInit()
        dimension = self._getPositiveDimension(input)
        if input.dim() > 1:
            gradOutputView = addSingletondimension(gradOutput, dimension)
        else:
            gradOutputView = gradOutput

        self.gradInput.resize_as_(input).zero_().scatter_(dimension, self._indices, gradOutputView)
        return self.gradInput

    def type(self, type, tensorCache=None):
        # torch.min expects a LongTensor as indices, whereas cutorch.max expects a CudaTensor.
        if type == 'torch.cuda.FloatTensor':
            indices, self._indices = self._indices, None
            super(Min, self).type(type, tensorCache)
            self._indices = indices.type('torch.cuda.LongTensor') if indices is not None else None
        else:
            # self._indices must be a LongTensor. Setting it to nil temporarily avoids
            # unnecessary memory allocations.
            indices, self._indices = self._indices, None
            super(Min, self).type(type, tensorCache)
            self._indices = indices.long() if indices is not None else None

        return self

    def clearState(self):
        clear(self, '_indices', '_output')
        return super(Min, self).clearState()
