local Module = torch.class('nn.Module')

function Module:__init()
   self.gradInput = torch.Tensor()
   self.output = torch.Tensor()
   self._type = self.output:type()
end

function Module:parameters()
   if self.weight and self.bias then
      return {self.weight, self.bias}, {self.gradWeight, self.gradBias}
   elseif self.weight then
      return {self.weight}, {self.gradWeight}
   elseif self.bias then
      return {self.bias}, {self.gradBias}
   else
      return
   end
end

function Module:updateOutput(input)
   return self.output
end

function Module:forward(input)
   return self:updateOutput(input)
end

function Module:backward(input, gradOutput, scale)
   scale = scale or 1
   self:updateGradInput(input, gradOutput)
   self:accGradParameters(input, gradOutput, scale)
   return self.gradInput
end

function Module:backwardUpdate(input, gradOutput, lr)
   self:updateGradInput(input, gradOutput)
   self:accUpdateGradParameters(input, gradOutput, lr)
   return self.gradInput
end

function Module:updateGradInput(input, gradOutput)
   return self.gradInput
end

function Module:accGradParameters(input, gradOutput, scale)
end

function Module:accUpdateGradParameters(input, gradOutput, lr)
   if self.shared then
      self:sharedAccUpdateGradParameters(input, gradOutput, lr)
   else
      self:defaultAccUpdateGradParameters(input, gradOutput, lr)
   end
end

function Module:defaultAccUpdateGradParameters(input, gradOutput, lr)
   local gradWeight = self.gradWeight
   local gradBias = self.gradBias
   self.gradWeight = self.weight
   self.gradBias = self.bias
   self:accGradParameters(input, gradOutput, -lr)
   self.gradWeight = gradWeight
   self.gradBias = gradBias
end

function Module:sharedAccUpdateGradParameters(input, gradOutput, lr)
   if self:parameters() then
      self:zeroGradParameters()
      self:accGradParameters(input, gradOutput, 1)
      self:updateParameters(lr)
   end
end

function Module:zeroGradParameters()
   local _,gradParams = self:parameters()
   if gradParams then
      for i=1,#gradParams do
         gradParams[i]:zero()
      end
   end
end

function Module:updateParameters(learningRate)
   local params, gradParams = self:parameters()
   if params then
      for i=1,#params do
         params[i]:add(-learningRate, gradParams[i])
      end
   end
end

function Module:training()
   self.train = true
end

function Module:evaluate()
   self.train = false
end

function Module:share(mlp, ...)
   local arg = {...}
   for i,v in ipairs(arg) do
      if self[v] ~= nil then
         self[v]:set(mlp[v])
         self.shared = true
         mlp.shared = true
      end
   end
   return self
end

local function sharedWrite(...)
   local arg = {...}
   local shared = {}
   for i,v in ipairs(arg) do
       shared[v] = true
   end
   return function(self, file)
      local object = {}
      for k, v in pairs(self) do
         if shared[k] then
            assert(torch.isTensor(v), 'Shared parameters have to be Tensors')
            object[k] = v.new()
         else
            object[k] = v
         end
      end
      file:writeObject(object)
   end
end

function Module:clone(...)
   local oldWrite = nn.Module.write
   nn.Module.write = sharedWrite(...)

   local f = torch.MemoryFile("rw"):binary()
   f:writeObject(self)
   f:seek(1)
   local clone = f:readObject()
   f:close()

   nn.Module.write = oldWrite

   if select('#',...) > 0 then
      clone:share(self,...)
   end
   return clone
end

function Module:type(type, tensorCache)
   if not type then
      return self._type
   end

   tensorCache = tensorCache or {}

   -- find all tensors and convert them
   for key,param in pairs(self) do
      self[key] = nn.utils.recursiveType(param, type, tensorCache)
   end

   self._type = type
   return self
end

function Module:float(...)
   return self:type('torch.FloatTensor',...)
end

function Module:double(...)
   return self:type('torch.DoubleTensor',...)
end

function Module:cuda(...)
   return self:type('torch.CudaTensor',...)
end

function Module:reset()
end

function Module:write(file)
  -- Write all values in the object as a table.
  local object = {}
  for k, v in pairs(self) do
    object[k] = v
  end
  file:writeObject(object)
end

function Module:read(file)
  local object = file:readObject()
  for k, v in pairs(object) do
    self[k] = v
  end
end

-- This function is not easy to understand. It works as follows:
--
-- - gather all parameter tensors for this module (and children);
--   count all parameter values (floats)
-- - create one ginormous memory area (Storage object) with room for all
--   parameters
-- - remap each parameter tensor to point to an area within the ginormous
--   Storage, and copy it there
--
-- It has the effect of making all parameters point to the same memory area,
-- which is then returned.
--
-- The purpose is to allow operations over all parameters (such as momentum
-- updates and serialization), but it assumes that all parameters are of
-- the same type (and, in the case of CUDA, on the same device), which
-- is not always true. Use for_each() to iterate over this module and
-- children instead.
--
-- Module._flattenTensorBuffer can be used by other packages (e.g. cunn)
-- to specify the type of temporary buffers. For example, the temporary
-- buffers for CudaTensor could be FloatTensor, to avoid GPU memory usage.
--
-- TODO: This logically belongs to torch.Tensor, not nn.
Module._flattenTensorBuffer = {}
function Module.flatten(parameters)

   -- returns true if tensor occupies a contiguous region of memory (no holes)
   local function isCompact(tensor)
      local sortedStride, perm = torch.sort(
            torch.LongTensor(tensor:nDimension()):set(tensor:stride()), 1, true)
      local sortedSize = torch.LongTensor(tensor:nDimension()):set(
            tensor:size()):index(1, perm)
      local nRealDim = torch.clamp(sortedStride, 0, 1):sum()
      sortedStride = sortedStride:narrow(1, 1, nRealDim):clone()
      sortedSize   = sortedSize:narrow(1, 1, nRealDim):clone()
      local t = tensor.new():set(tensor:storage(), 1,
                                 sortedSize:storage(),
                                 sortedStride:storage())
      return t:isContiguous()
   end

   if not parameters or #parameters == 0 then
      return torch.Tensor()
   end
   local Tensor = parameters[1].new
   local TmpTensor = Module._flattenTensorBuffer[torch.type(parameters[1])] or Tensor

   -- 1. construct the set of all unique storages referenced by parameter tensors
   local storages = {}
   local nParameters = 0
   local parameterMeta = {}
   for k = 1,#parameters do
      local param = parameters[k]
      local storage = parameters[k]:storage()
      local storageKey = torch.pointer(storage)

      if not storages[storageKey] then
         storages[storageKey] = {storage, nParameters}
         nParameters = nParameters + storage:size()
      end

      parameterMeta[k] = {storageOffset = param:storageOffset() +
                                          storages[storageKey][2],
                          size          = param:size(),
                          stride        = param:stride()}
   end

   -- 2. construct a single tensor that will hold all the parameters
   local flatParameters = TmpTensor(nParameters):zero()

   -- 3. determine if there are elements in the storage that none of the
   --    parameter tensors reference ('holes')
   local tensorsCompact = true
   for k = 1,#parameters do
      local meta = parameterMeta[k]
      local tmp = TmpTensor():set(
         flatParameters:storage(), meta.storageOffset, meta.size, meta.stride)
      tmp:fill(1)
      tensorsCompact = tensorsCompact and isCompact(tmp)
   end

   local maskParameters  = flatParameters:byte():clone()
   local compactOffsets  = flatParameters:long():cumsum(1)
   local nUsedParameters = compactOffsets[-1]

   -- 4. copy storages into the flattened parameter tensor
   for _, storageAndOffset in pairs(storages) do
      local storage, offset = table.unpack(storageAndOffset)
      flatParameters[{{offset+1,offset+storage:size()}}]:copy(Tensor():set(storage))
   end

   -- 5. allow garbage collection
   storages = nil
   for k = 1,#parameters do
       parameters[k]:set(Tensor())
   end

   -- 6. compact the flattened parameters if there were holes
   if nUsedParameters ~= nParameters then
      assert(tensorsCompact,
         "Cannot gather tensors that are not compact")

      flatParameters = TmpTensor(nUsedParameters):copy(
            flatParameters:maskedSelect(maskParameters))
      for k = 1,#parameters do
        parameterMeta[k].storageOffset =
              compactOffsets[parameterMeta[k].storageOffset]
      end
   end

   if TmpTensor ~= Tensor then
      flatParameters = Tensor(flatParameters:nElement()):copy(flatParameters)
   end

   -- 7. fix up the parameter tensors to point at the flattened parameters
   for k = 1,#parameters do
      parameters[k]:set(flatParameters:storage(),
          parameterMeta[k].storageOffset,
          parameterMeta[k].size,
          parameterMeta[k].stride)
   end

   return flatParameters
end

function Module:getParameters()
   -- get parameters
   local parameters,gradParameters = self:parameters()
   local p, g = Module.flatten(parameters), Module.flatten(gradParameters)
   assert(p:nElement() == g:nElement(),
      'check that you are sharing parameters and gradParameters')
   if parameters then
      for i=1,#parameters do
         assert(parameters[i]:storageOffset() == gradParameters[i]:storageOffset(),
            'misaligned parameter at ' .. tostring(i))
      end
   end
   return p, g
end

function Module:__call__(input, gradOutput)
   self:forward(input)
   if gradOutput then
      self:backward(input, gradOutput)
      return self.output, self.gradInput
   else
      return self.output
   end
end

-- Run a callback (called with the module as an argument) in preorder over this
-- module and its children.
--
function Module:apply(callback)
    callback(self)

    if self.modules then
        for _, module in ipairs(self.modules) do
            module:apply(callback)
        end
    end
end

function Module:findModules(typename, container)
  container = container or self
  local nodes = {}
  local containers = {}
  local mod_type = torch.typename(self)
  if mod_type == typename then
    nodes[#nodes+1] = self
    containers[#containers+1] = container
  end
  -- Recurse on nodes with 'modules'
  if (self.modules ~= nil) then
    if (torch.type(self.modules) == 'table') then
      for i = 1, #self.modules do
        local child = self.modules[i]
        local cur_nodes, cur_containers =
          child:findModules(typename, self)
        assert(#cur_nodes == #cur_containers,
          'Internal error: incorrect return length')  -- This shouldn't happen
        -- add the list items from our child to our list (ie return a
        -- flattened table of the return nodes).
        for j = 1, #cur_nodes do
          nodes[#nodes+1] = cur_nodes[j]
          containers[#containers+1] = cur_containers[j]
        end
      end
    end
  end
  return nodes, containers
end

-- returns a list of modules
function Module:listModules()
   local function tinsert(to, from)
      if torch.type(from) == 'table' then
         for i=1,#from do
            tinsert(to,from[i])
         end
      else
         table.insert(to,from)
      end
   end
   -- include self first
   local modules = {self}
   if self.modules then
      for i=1,#self.modules do
         local modulas = self.modules[i]:listModules()
         if modulas then
            tinsert(modules,modulas)
         end
      end
   end
   return modules
end

function Module:clearState()
   return nn.utils.clear(self, 'output', 'gradInput')
end

-- similar to apply, recursively goes over network and calls
-- a callback function which returns a new module replacing the old one
function nn.Module:replace(callback)
   local out = callback(self)
   if self.modules then
      for i, module in ipairs(self.modules) do
         self.modules[i] = module:replace(callback)
      end
   end
   return out
end