require 'torch'
require 'nn'
require 'sys'

profiler = {}

function profiler.decor(model, functionsToDecorate)

   local functionsToDecorate = functionsToDecorate or
      {
         'updateOutput',
         'backward'
      }

   for _, name in pairs(functionsToDecorate) do
      model.timings = 0

      local functionTable = model

      if not rawget(functionTable, name) then
         functionTable = getmetatable(model)
      end

      if functionTable[name] and not (functionTable.orig and functionTable.orig[name]) then
         print('Profiler decoring ' .. functionTable.__typename .. '.' .. name)
         functionTable.orig = functionTable.orig or {}
         functionTable.orig[name] = functionTable[name]
         functionTable[name] = function(self, ...)
            local startTime = sys.clock()
            local result = { self.orig[name](self, unpack({...})) }
            local endTime = sys.clock()
            self.timings = self.timings + endTime - startTime
            return unpack(result)
         end
      end

   end

   if torch.isTypeOf(model, nn.Container) then
      for _, m in ipairs(model.modules) do
         profiler.decor(m, functionsToDecorate)
      end
   end

end

function profiler.print(model, nbSamples)
   print('----------------------------------------------------------------------')
   print(model)
   if nbSamples then
      print(string.format('acc_time %.02fs (%.1ems/sample)', model.timings, 1000 * model.timings / nbSamples))
   else
      print(string.format('acc_time %.02fs', model.timings))
   end

   if torch.isTypeOf(model, nn.Container) then
      for _, m in ipairs(model.modules) do
         profiler.print(m, nbSamples)
      end
   end
end

return profiler