require 'torch';
require 'nn';
require 'optim';

--cudaFlag = true
cudaFlag = true

if cudaFlag then
	require 'cutorch';
	require 'cunn';
end

-- parameters
local batchSize = 200		-- batch size is approximate; actual batchSizes will vary by +/- N-1, where N is the number of classes
local learningRate = 0.05
local learningRateDecay = 0.0005
local weightDecay = 0.000
local momentum = 0.9
local maxIteration = 20
local augment = true

local trainFolder = '/home/andrew/mitosis/data/mitosis-train-large/'
local testFolder = '/home/andrew/mitosis/data/mitosis-test/'

dofile("data.lua")

local classes, trainClassList, trainImagePaths = getImagePaths(trainFolder)
local classes, testClassList, testImagePaths = getImagePaths(testFolder)


local classRatio = trainClassList[2]:size(1)/trainClassList[1]:size(1)
local weights = torch.Tensor(2)
weights[1] = classRatio
weights[2] = 1

--[
-- first network
-- define the model
dofile("/home/andrew/mitosis/models/model1.lua")
local net, criterion = model1(weights)
if cudaFlag then
	net = net:cuda()
	criterion = criterion:cuda()
end

-- train the network
dofile("train.lua")
--net = torch.load('/home/andrew/mitosis/data/nets/dnn1_halfset_aug_20i_lr001.t7')
--train(net, criterion, classes, trainClassList, imagePaths, batchSize, learningRate, maxIteration)
--torch.save('/home/andrew/mitosis/data/nets/dnn1_fullset_aug_30i_lr05_mini200.t7', net)

-- test the network
dofile("test.lua")
--test(classes, testClassList, imagePaths, batchSize, maxIteration)
--]]


-- second network
-- define the model
dofile("/home/andrew/mitosis/models/model2.lua")
local net, criterion = model2(weights)
if cudaFlag then
	net = net:cuda()
	criterion = criterion:cuda()
end

-- train the network
dofile("train.lua")
net = torch.load('/home/andrew/mitosis/data/nets/model2-pretrained-greedylayerwise2.t7')
train(net, criterion, classes, trainClassList, trainImagePaths, batchSize, learningRate, learningRateDecay, weightDecay, momentum, maxIteration, classRatio, augment, netFolder)
torch.save('/home/andrew/mitosis/data/nets/dnn2_fullset_aug_20i_lr05_lrd0005_m09_mini200_aeptgl.t7', net)

-- test the network
dofile("test.lua")
--net = torch.load('/home/andrew/mitosis/data/nets/dnn2_fullset_aug_20i_lr05_mini200_aept.t7')
test(net, classes, testClassList, testImagePaths, batchSize)