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

cudaFlag = true

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

--torch.setnumthreads(4)

-- parameters
local batchSize = 100
local learningRate = 0.01
local weightDecay = 0.000
local maxIteration = 10

local folder  = '/home/andrew/mitosis/data/mitosis-train-old/'

dofile("data.lua")
local classes, classList, imagePaths = getImagePaths(folder)

-- split dataset into training and test sets
local trainClassList = {}
local testClassList = {}
--[
trainClassList[1] = classList[1][{{1,math.ceil(classList[1]:size(1)/2)}}]
trainClassList[2] = classList[2][{{1,math.ceil(classList[2]:size(1)/2)}}]
testClassList[1] = classList[1][{{math.ceil(classList[1]:size(1)/2)+1,classList[1]:size(1)}}]
testClassList[2] = classList[2][{{math.ceil(classList[2]:size(1)/2)+1,classList[2]:size(1)}}]
--]]
--[[
trainClassList[1] = classList[1][{{1,math.ceil(classList[1]:size(1)/20)}}]
trainClassList[2] = classList[2][{{1,math.ceil(classList[2]:size(1)/20)}}]
testClassList[1] = classList[1][{{math.ceil(classList[1]:size(1)/20)+1,math.ceil(classList[1]:size(1)/10)}}]
testClassList[2] = classList[2][{{math.ceil(classList[2]:size(1)/20)+1,math.ceil(classList[1]:size(1)/10)}}]
--]]
--trainClassList = classList

local classRatio = trainClassList[2]:size(1)/trainClassList[1]:size(1)

-- define the model
dofile("/home/andrew/mitosis/models/model.lua")

-- load the pre-trained model
torch.load('/home/andrew/mitosis/data/nets/model-pretrained.t7')

-- train the network
dofile("train.lua")
train(net, criterion, classes, trainClassList, imagePaths, batchSize, learningRate, weightDecay, maxIteration, classRatio, false)

-- save the model
torch.save('/home/andrew/mitosis/data/nets/testNet.t7', net)
--net = torch.load('/home/andrew/mitosis/nets/testNet.t7')

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