#!/usr/bin/env luajit

--[[

   Copyright (c) 2016 Idiap Research Institute, http://www.idiap.ch/
   Written by Francois Fleuret <francois.fleuret@idiap.ch>

   This file is free software: you can redistribute it and/or modify
   it under the terms of the GNU General Public License version 3 as
   published by the Free Software Foundation.

   It is distributed in the hope that it will be useful, but WITHOUT
   ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
   or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public
   License for more details.

   You should have received a copy of the GNU General Public License
   along with this file.  If not, see <http://www.gnu.org/licenses/>.

]]--

require 'torch'
require 'nn'

-- require 'cunn'

require 'dagnn'

torch.setdefaulttensortype('torch.DoubleTensor')
torch.manualSeed(1)

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

   local epsilon = epsilon or 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()

   local err = 0

   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)

      if num ~= ana then
         err = math.max(err, math.abs(num - ana) / math.max(epsilon, math.abs(num)))
      end
   end

   return err
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

--               +-- Linear(10, 10) --> ReLU --> d -->
--              /                               /
--             /                               /
--  --> a --> b -----------> c ---------------+
--                            \
--                             \
--                              +--------------- e -->

dag = nn.DAG()

a = nn.Linear(50, 10)
b = nn.ReLU()
c = nn.Linear(10, 15)
d = nn.CMulTable()
e = nn.Mul(-1)

dag:connect(a, b, c)
dag:connect(b, nn.Linear(10, 15), nn.ReLU(), d)
dag:connect(c, d)
dag:connect(c, e)

dag:setLabel(a, 'first module')

dag:setInput(a)
dag:setOutput({ d, e })

-- Check the output of the dot file. Generate a pdf with:
--
-- dot ./graph.dot -Lg -T pdf -o ./graph.pdf
--
print('Writing ./graph.dot')
dag:saveDot('./graph.dot')

-- Let's make a model where the dag is inside another nn.Container.
model = nn.Sequential()
   :add(nn.Linear(50, 50))
   :add(dag)
   :add(nn.CAddTable())

criterion = nn.MSECriterion()

if cunn then
   print("Using CUDA")
   model:cuda()
   criterion:cuda()
   torch.setdefaulttensortype('torch.CudaTensor')
   epsilon = 1e-3
end

local input = torch.Tensor(30, 50):uniform()
local output = model:updateOutput(input):clone()
output:uniform()

-- Check that DAG:accGradParameters and friends work okay
print('Gradient estimate error ' .. checkGrad(model, criterion, input, output, epsilon))

-- Check that we can save and reload the model
model:clearState()
torch.save('./test.t7', model)
local otherModel = torch.load('./test.t7')
print('Gradient estimate error ' .. checkGrad(otherModel, criterion, input, output, epsilon))

dag:print()