-- module for memory management

-- reuseMem is used for reusing output tensor for storing gradInput and optimizing memory allocation
-- use :reuseMem() on the module to allow the feature
-- then apply setReuse after initialization
-- only applies if output and gradinput are of the same type
function nn.Module:reuseMem(name)
    self.reuse = true
    return self
end

function nn.Module:setReuse()
    if self.reuse then
        assert(type(self.output) == type(self.gradInput), "invalid use of reuseMem:")
        self.gradInput = self.output
    end
    return self
end

-- usePrealloc is based on the same principle but use pre-allocated memory at the beginning of the process that can be shared
-- between different objects
-- use to prellocate gradInput, or output - useful for intermediate calculations working on large input
preallocTable = nil

function preallocateMemory(switch)
    if switch == 1 then
        preallocTable = {}
        print('Switching on memory preallocation')
    end
end

function preallocateTensor(name, D)
    if #D > 1 then
        local T = {}
        for i = 1, #D do
            table.insert(T, preallocateTensor(name, {D[i]}))
        end
        return T
    else
        D = D[1]
    end
    local t = torch.zeros(torch.LongStorage(D))
    if opt.gpuid >= 0 then
      t = t:cuda()
    end
    return t
end

-- enable reuseMemory - if preallocation disable, then switched back to reuseMem checking for 'reuse' in name
function nn.Module:usePrealloc(preallocName, inputDim, outputDim)
    if preallocTable == nil then
        if string.find(preallocName, "reuse") then
            self:reuseMem()
        end
        return self
    end
    self.prealloc = preallocName
    self.name = preallocName
    self.preallocInputDim = inputDim
    self.preallocOutputDim = outputDim
    return self
end

function nn.Module:setPrealloc()
    if self.prealloc and (self.preallocInputDim ~= nil or self.preallocOutputDim ~= nil) then
        if preallocTable[self.prealloc] == nil then
            preallocTable[self.prealloc] = {
            }
            if self.preallocInputDim ~= nil then
                preallocTable[self.prealloc].GI = preallocateTensor(self.prealloc, self.preallocInputDim)
            end
            if self.preallocOutputDim ~= nil then
                preallocTable[self.prealloc].O = preallocateTensor(self.prealloc, self.preallocOutputDim)
            end
        end
        local memmap = preallocTable[self.prealloc]
        if memmap["GI"] ~= nil then
            assert(type(self.gradInput) == type(memmap.GI), "invalid use of usePrealloc ["..self.prealloc.."]/GI: "..type(self.gradInput).."/"..type(memmap.GI))
            self.gradInput = memmap["GI"]
        end
        if memmap["O"] ~= nil then
            assert(type(self.output) == type(memmap.O), "invalid use of usePrealloc ["..self.prealloc.."]/O:"..type(self.output).."/"..type(memmap.O))
            self.output = memmap["O"]
        end
    end
    return self
end
