--[[
Copyright (c) 2016 Michael Wilber

This software is provided 'as-is', without any express or implied
warranty. In no event will the authors be held liable for any damages
arising from the use of this software.

Permission is granted to anyone to use this software for any purpose,
including commercial applications, and to alter it and redistribute it
freely, subject to the following restrictions:

1. The origin of this software must not be misrepresented; you must not
   claim that you wrote the original software. If you use this software
   in a product, an acknowledgement in the product documentation would be
   appreciated but is not required.
2. Altered source versions must be plainly marked as such, and must not be
   misrepresented as being the original software.
3. This notice may not be removed or altered from any source distribution.
--]]

require 'residual-layers'
require 'nn'
require 'cutorch'
require 'cunn'
require 'cudnn'
require 'nngraph'
require 'train-helpers'
display = require 'display'
workbook = (require'lab-workbook-for-trello'):newExperiment{}

opt = lapp[[
      --batchSize       (default 64)      Sub-batch size
      --iterSize        (default 1)       How many sub-batches in each batch
      --nThreads        (default 4)       Data loader threads
      --dataTrainRoot   (default /mnt/imagenet/train)   Data root folder
      --dataValRoot     (default /mnt/imagenet/val)   Data root folder
      --verbose         (default true)
      --loadSize        (default 256)     Size of image when loading
      --fineSize        (default 224)     Size of image crop
      --loadFrom        (default "")      Model to load
]]
print(opt)

-- create data loader
local DataLoader = paths.dofile('data/data.lua')
dataTrain = DataLoader.new(opt.nThreads, 'folder', {dataRoot = opt.dataTrainRoot,
                                                    fineSize = opt.fineSize,
                                                    loadSize = opt.loadSize,
                                                    batchSize = opt.batchSize,
                                                })
dataVal = DataLoader.new(opt.nThreads, 'folder', {dataRoot = opt.dataValRoot,
                                                  fineSize = opt.fineSize,
                                                  loadSize = opt.loadSize,
                                                  batchSize = opt.batchSize,
                                              })
print("Dataset size: ", dataTrain:size())


-- dataset_train = Dataset.MNIST("mnist.hdf5", 'train')
-- dataset_test = Dataset.MNIST("mnist.hdf5", 'test')

-- mean,std = dataset_train:preprocess()
-- dataset_test:preprocess(mean,std)

-- -- LENET here
-- -- model = nn.Sequential()
-- -- -- stage 1 : mean suppresion -> filter bank -> squashing -> max pooling
-- -- model:add(nn.SpatialConvolutionMM(1, 32, 5, 5, 1,1))
-- -- model:add(nn.ReLU())
-- -- model:add(nn.SpatialMaxPooling(2,2, 2,2))
-- -- -- stage 2 : mean suppresion -> filter bank -> squashing -> max pooling
-- -- model:add(nn.SpatialConvolutionMM(32, 64, 5, 5, 1,1))
-- -- model:add(nn.ReLU())
-- -- model:add(nn.SpatialMaxPooling(2, 2, 2,2))
-- -- -- stage 3 : standard 2-layer MLP:
-- -- model:add(nn.Reshape(64*4*4))
-- -- model:add(nn.Linear(64*4*4, 10))
-- -- -- model:add(nn.ReLU())
-- -- -- model:add(nn.Linear(200, 10))
-- -- model:add(nn.LogSoftMax())

-- Residual network.
-- Input: 3x224x224
if opt.loadFrom == "" then
    input = nn.Identity()()
    ------> 64, 112,112
    model = cudnn.SpatialConvolution(3, 64, 7,7, 2,2, 3,3)(input)
    --model = nn.SpatialBatchNormalization(64)(model)
    model = cudnn.ReLU(true)(model)
    model = cudnn.SpatialMaxPooling(3,3,  2,2,  1,1)(model)
    ------> 64, 56,56
    model = addResidualLayer2(model, 64)
    model = addResidualLayer2(model, 64)
    model = addResidualLayer2(model, 64)
    ------> 128, 28,28
    model = addResidualLayer2(model, 64, 128, 2)
    model = addResidualLayer2(model, 128)
    model = addResidualLayer2(model, 128)
    model = addResidualLayer2(model, 128)
    ------> 256, 14,14
    model = addResidualLayer2(model, 128, 256, 2)
    model = addResidualLayer2(model, 256)
    model = addResidualLayer2(model, 256)
    model = addResidualLayer2(model, 256)
    model = addResidualLayer2(model, 256)
    model = addResidualLayer2(model, 256)
    ------> 512, 7,7
    model = addResidualLayer2(model, 256, 512, 2)
    model = addResidualLayer2(model, 512)
    model = addResidualLayer2(model, 512)
    ------> 1000, 1,1
    model = cudnn.ReLU(true)(cudnn.SpatialConvolution(512, 1000, 7,7)(model))
    ------> 1000
    model = nn.Reshape(1000)(model)
    model = nn.LogSoftMax()(model)

    model = nn.gModule({input}, {model})
    model:cuda()
else
    print("Loading model from "..opt.loadFrom)
    cutorch.setDevice(1)
    model = torch.load(opt.loadFrom)
    print "Done"
end

loss = nn.ClassNLLCriterion()
loss:cuda()

-- Convert model to multi-GPU model :)
----[[
dmodel = nn.DataParallelTable(1)
dmodel:add(model, 1)
for i=2, 4 do
    cutorch.setDevice(i)
    dmodel:add(model:clone(), i)
end
cutorch.setDevice(1)
model = dmodel
--]]
-- required to fix call to SyncParameters
model:forward(torch.randn(8, 3, 224,224):cuda())

-- Dirty trick: make the first conv layer weights easier to modify
-- model.modules[2].weight:mul(0.5)

--[[
model:apply(function(m)
    -- Initialize weights
    local name = torch.type(m)
    if name:find('Convolution') then
        m.weight:normal(0.0, math.sqrt(2/(m.nInputPlane*m.kW*m.kH)))
        print(m.weight:std())
        m.bias:fill(0)
    elseif name:find('BatchNormalization') then
        if m.weight then m.weight:normal(1.0, 0.002) end
        if m.bias then m.bias:fill(0) end
    end
end)
--]]



sgdState = {
   --- For SGD with momentum ---
   ----[[
   -- My semi-working settings
   learningRate = 0.001,
   weightDecay    = 1e-4,
   -- Settings from their paper
   --learningRate = 0.1,
   --weightDecay    = 1e-4,

   momentum     = 0.9,
   dampening    = 0,
   nesterov     = true,
   --]]
   --- For rmsprop, which is very fiddly and I don't trust it at all ---
   --[[
   learningRate = 1e-5,
   alpha = 0.9,
   whichOptimMethod = 'rmsprop',
   --]]
   --- For adadelta, which sucks ---
   --[[
   rho              = 0.3,
   whichOptimMethod = 'adadelta',
   --]]
   --- For adagrad, which also sucks ---
   --[[
   learningRate = 3e-4,
   whichOptimMethod = 'adagrad',
   --]]
   --- For adam, which also sucks ---
   --[[
   learningRate = 0.005,
   whichOptimMethod = 'adam',
   --]]
   --- For the alternate implementation of NAG ---
   --[[
   learningRate = 0.01,
   weightDecay = 1e-6,
   momentum = 0.9,
   whichOptimMethod = 'nag',
   --]]
}


if opt.loadFrom ~= "" then
    print("Trying to load sgdState from "..string.gsub(opt.loadFrom, "model", "sgdState"))
    collectgarbage(); collectgarbage(); collectgarbage()
    sgdState = torch.load(""..string.gsub(opt.loadFrom, "model", "sgdState"))
    collectgarbage(); collectgarbage(); collectgarbage()
    print("Got", sgdState.nSampledImages,"images")
end

-- Actual Training! -----------------------------
weights, gradients = model:getParameters()
function forwardBackwardBatch(batch)
    -- After every batch, the different GPUs all have different gradients
    -- (because they saw different data), and only the first GPU's weights were
    -- actually updated.
    -- We have to do two changes:
    --   - Copy the new parameters from GPU #1 to the rest of them;
    --   - Zero the gradient parameters so we can accumulate them again.
   model:training()
   model:syncParameters() -- This copies parameters from first GPU to
   model:zeroGradParameters()

   --[[
   if sgdState.nSampledImages < 10000 then
       sgdState.learningRate = 0.001
   else
       sgdState.learningRate = 0.01
   end
   --]]

   local loss_val = 0
   local N = opt.iterSize
   local inputs, labels
   for i=1,N do
       inputs, labels = dataTrain:getBatch()
       inputs = inputs:cuda()
       labels = labels:cuda()
       collectgarbage(); collectgarbage();
       local y = model:forward(inputs)
       loss_val = loss_val + loss:forward(y, labels)
       local df_dw = loss:backward(y, labels)
       model:backward(inputs, df_dw)
       -- The above call will accumulate all GPUs' parameters onto GPU #1
   end
   loss_val = loss_val / N
   gradients:mul( 1.0 / N )

   if sgdState.nEvalCounter % 20 == 0 then
      display.image(model.modules[1].modules[2].weight, {win=24, title="First layer weights"})
   end

   return loss_val, gradients, inputs:size(1) * N
end


function evalModel()
   print("No evaluation...")
   -- if sgdState.epochCounter > 10 then os.exit(1) end
   local results = evaluateModel(model, dataVal)
   print(results)
   --table.insert(sgdState.accuracies, acc)
end

--[[
local results = evaluateModel(model, dataVal)
print(results)
--]]

--[[
require 'graph'
graph.dot(model.fg, 'MLP', '/tmp/MLP')
os.execute('convert /tmp/MLP.svg /tmp/MLP.png')
display.image(image.load('/tmp/MLP.png'), {title="Network Structure", win=23})
--]]

--[[
require 'ncdu-model-explore'
local y = model:forward(torch.randn(opt.batchSize, 3, 224,224):cuda())
local df_dw = loss:backward(y, torch.zeros(opt.batchSize):cuda())
model:backward(torch.randn(opt.batchSize,3,224,224):cuda(), df_dw)
exploreNcdu(model)
--]]


-- --[[
TrainingHelpers.trainForever(
   model.modules[1],
   forwardBackwardBatch,
   weights,
   sgdState,
   dataTrain:size(),
   evalModel,
   "snapshots/imagenet-residual-experiment2"
)
--]]