require 'torch';

require 'nn';

require 'optim';

cudaFlag = true

if cudaFlag then

    require 'cutorch';

    require 'cunn';

end

c = os.clock()

t = os.time()

-- parameters

local learningRate = 0.05

local learningRateDecay = 0.0005

local weightdecay = 0.0000

local momentum = 0.9

local maxIteration = 5

local p = 0.25

local batchSize = 200

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

dofile("data.lua")

local classes, classList, imagePaths = getImagePaths(folder)

dofile('/home/andrew/mitosis/models/model2.lua')

net = model2()

if cudaFlag then

    net = net:cuda()

end

dofile('/home/andrew/mitosis/code/autoencoder.lua')

autoencoder = convnet2autoencoder(net)

--[

autoencoder:insert(nn.Dropout(p),1)

if cudaFlag then

    autoencoder = autoencoder:cuda()

end

criterion = nn.MSECriterion()

if cudaFlag then

    criterion = criterion:cuda()

end

-- compute size of each batch

batchSizes, numBatches = getBatchSizes(classes, classList, batchSize)

-- shuffle the images

classList = shuffleImages(classList, classes)

-- train

print("# StochasticGradient: training")

autoencoder:training()

subNet = nn.Sequential()

subNet:insert(autoencoder:get(1),1)

-- count the number of convolutional layers

numConvLayers = 0

for i = 1, net:size() do

    if torch.typename(net:get(i)) == 'nn.SpatialConvolution' then

        numConvLayers = numConvLayers + 1

    end

end

errors = {}

for j = 1, numConvLayers do

    c0 = os.clock()

    t0 = os.time()

--[

    subNet:insert(autoencoder:get(3*(j-1)+2),3*(j-1)+2)

    subNet:insert(autoencoder:get(3*(j-1)+3),3*(j-1)+3)

    subNet:insert(autoencoder:get(3*(j-1)+4),3*(j-1)+4)

    subNet:insert(autoencoder:get(autoencoder:size()-3*(j-1)-2),3*(j-1)+5)

    subNet:insert(autoencoder:get(autoencoder:size()-3*(j-1)-1),3*(j-1)+6)

    subNet:insert(autoencoder:get(autoencoder:size()-3*(j-1)),3*(j-1)+7)

    params, gradParams = subNet:getParameters()

    optimState = {}

    optimState.learningRate = learningRate

    optimState.learningRateDecay = learningRateDecay

    optimState.weightDecay = weightDecay

    optimState.momentum = momentum

--]]

    for epoch = 1, maxIteration do

        c1 = os.clock()

        t1 = os.time()

        local currentError = 0

        local sampleSum = {}

        for i = 1, #classes do

            sampleSum[i] = 0

        end

        for i = 1, numBatches do

            c2 = os.clock()

            t2 = os.time()

            -- split classList into batches

            local sampleList = {}

            for j=1,#classes do

                sampleList[j] = classList[j][{{sampleSum[j] + 1, sampleSum[j] + batchSizes[j][i]}}]

                sampleSum[j] = sampleSum[j] + batchSizes[j][i]

            end

            local dataset = getSample(classes, sampleList, imagePaths)

            if cudaFlag then

                dataset.data = dataset.data:cuda()

            end

            dataset.label = dataset.data

            local input = dataset.data

            local target = dataset.label

            function feval(params)

                gradParams:zero()

                local outputs = subNet:forward(input)

                local loss = criterion:forward(outputs, target)

                local dloss_doutputs = criterion:backward(outputs, target)

                subNet:backward(input, dloss_doutputs)

                return loss, gradParams

            end

             _, fs = optim.sgd(feval, params, optimState)

            print('Layer = ' .. j .. ' of ' .. numConvLayers)

            print('Epoch = ' .. epoch .. ' of ' .. maxIteration)

            print('Batch = ' .. i .. ' of ' .. numBatches)

            for k=1,#errors do

                print('Final Error for Layer ' .. k .. ' = ' .. errors[k])

            end

            print('Error = ' .. fs[1])

            print('CPU batch time = ' .. os.clock()-c2 .. ' seconds')

            print('Actual batch time (rounded) = ' .. os.time()-t2 .. ' seconds')

            if epochClock then

                print('CPU epoch time = ' .. epochClock .. ' seconds')

                print('Actual epoch time (rounded) = ' .. epochTime .. ' seconds')

            end

            if layerClock then

                print('CPU layer time = ' .. layerClock .. ' seconds')

                print('Actual layer time (rounded) = ' ..layerTime .. ' seconds')

            end

            print('Total CPU time so far = ' .. os.clock()-c .. ' seconds')

            print('Total actual time so far (rounded) = ' .. os.time()-t .. ' seconds')

            print('')

        end

        epochClock = os.clock()-c1

        epochTime = os.time()-t1

    end

    errors[j] = fs[1]

    layerClock = os.clock()-c0

    layerTime = os.time()-t0

end

--autoencoder = torch.load('/home/andrew/mitosis/data/nets/pretrain.t7')

-- get CNN from autoencoder

net = autoencoder2convnet(autoencoder,net)

torch.save('/home/andrew/mitosis/data/nets/model2-pretrained-greedylayerwise2.t7',net)

totalClock = os.clock()-c

totalTime = os.time()-t

print('Total CPU time = ' .. totalClock .. ' seconds')

print('Total actual time (rounded) ' .. totalTime .. ' seconds')