#!/usr/bin/env luajit

require 'torch'
require 'nn'

require 'dagnn'

function checkGrad(model, criterion, input, target)
   local params, gradParams = model:getParameters()

   local epsilon = 1e-5

   local output = model:forward(input)
   local loss = criterion:forward(output, target)
   local gradOutput = criterion:backward(output, target)
   gradParams:zero()
   model:backward(input, gradOutput)
   local analyticalGradParam = gradParams:clone()

   for i = 1, params:size(1) do
      local x = params[i]

      params[i] = x - epsilon
      local output0 = model:forward(input)
      local loss0 = criterion:forward(output0, target)

      params[i] = x + epsilon
      local output1 = model:forward(input)
      local loss1 = criterion:forward(output1, target)

      params[i] = x

      local ana = analyticalGradParam[i]
      local num = (loss1 - loss0) / (2 * epsilon)
      local err = torch.abs(num - ana) / torch.abs(num)

      print(
         err .. ' checkGrad ' .. i
            .. ' analytical ' .. ana
            .. ' numerical ' .. num
      )
   end

end

function printTensorTable(t)
   if torch.type(t) == 'table' then
      for i, t in pairs(t) do
         print('-- ELEMENT [' .. i .. '] --')
         printTensorTable(t)
      end
   else
      print(tostring(t))
   end
end

-- torch.setnumthreads(params.nbThreads)
torch.setdefaulttensortype('torch.DoubleTensor')
torch.manualSeed(2)

--                     +--> c ----> e --+
--                    /            /     \
--                   /            /       \
--  input --> a --> b ---> d ----+         g --> output
--                          \             /
--                           \           /
--                            +--> f ---+

a = nn.Linear(10, 10)
b = nn.ReLU()
c = nn.Linear(10, 3)
d = nn.Linear(10, 3)
e = nn.CMulTable()
f = nn.Linear(3, 3)
g = nn.CAddTable()

----------------------------------------------------------------------

model = nn.DAG()

model:addEdge(a, b)
model:addEdge(b, c)
model:addEdge(b, d)
model:addEdge(c, e)
model:addEdge(d, e)
model:addEdge(d, f)
model:addEdge(e, g)
model:addEdge(f, g)

model:setInput(a)
model:setOutput(g)

input = torch.Tensor(3, 10):uniform()

print('******************************************************************')
print('** updateOutput **************************************************')
print('******************************************************************')

output = model:updateOutput(input):clone()

printTensorTable(output)

print('******************************************************************')
print('** updateGradInput ***********************************************')
print('******************************************************************')

gradInput = model:updateGradInput(input, output)

printTensorTable(gradInput)

print('******************************************************************')
print('** checkGrad *****************************************************')
print('******************************************************************')

output:uniform()

checkGrad(model, nn.MSECriterion(), input, output)