require 'torch';

require 'nn';

require 'image';

cuda = false

require 'cutorch';

require 'cunn';

require 'sys';

require 'csvigo';

local dir = require 'pl.dir';

local c = os.clock()

local t = os.time()

local folder = '/home/andrew/mitosis/data/MITOS/testing/'

--local netPath1 = '/home/andrew/mitosis/data/nets/net.t7'

local netPath1 = '/home/andrew/mitosis/data/nets/dnn1_fullset_aug_20i_lr05_lrd0005_m09_mini200_aeptgl.t7'

local netPath2 = '/home/andrew/mitosis/data/nets/dnn2_fullset_aug_20i_lr05_lrd0005_m09_mini200_aeptgl.t7'

local threshold = 0.1

dofile("getImagePaths.lua")

local imagePaths = getImagePaths(folder)

if paths.dirp(folder .. 'results') == false then

    paths.mkdir(folder .. 'results');

end

dofile("scan.lua")

dofile("expand.lua")

-- define circular kernel

local d = 10

local kernel = torch.FloatTensor(2*d+1,2*d+1):fill(0.001)

for i=1,2*d+1 do

    for j=1,2*d+1 do

        if (i-d-1)^2 + (j-d-1)^2 >= d^2 then

            kernel[i][j] = 0

        end

    end

end

for k,imagePath in ipairs(imagePaths) do

    print(k)

    local c1 = os.clock()

    local t1 = os.time()

    --local windowWidth = 101

    --local windowHeight = 101

    local net1 = torch.load(netPath1)

    local net2 = torch.load(netPath2)

    net1 = expand(net1)

    net2 = expand(net2)

    if cuda then

        net1 = net1:cuda()

        net2 = net2:cuda()

    else

        net1 = net1:float()

        net2 = net2:float()

    end

    local img = image.load(imagePath, 3, 'float')

    if cuda then

        img = img:cuda()

    else

        img = img:float()

    end

    -- scan sixteen versions of the image

    -- four rotations, paired with their reflections, and two neural nets

    --]]

    --[

    local maps = {}

    for i=1,4 do

        for j=1,2 do

            local tmp = image.rotate(img, (i-1)*math.pi/2)

            if j == 2 then

                tmp = image.hflip(tmp)

                tmp = scan(tmp, net1)

                tmp = image.hflip(tmp)

            else

                tmp = scan(tmp, net1)

            end

            maps[(i-1)*2+j] = image.rotate(tmp, -(i-1)*math.pi/2)

        end

    end

    net1 = nil

--[

    for i=5,8 do    

        for j=1,2 do

            local tmp = image.rotate(img, (i-1)*math.pi/2)

            if j == 2 then

                tmp = image.hflip(tmp)

                tmp = scan(tmp, net2)

                tmp = image.hflip(tmp)

            else

                tmp = scan(tmp, net2)

            end

            maps[(i-1)*2+j] = image.rotate(tmp, -(i-1)*math.pi/2)

        end

    end

    net2 = nil

--]]

    -- take the mean of the sixteen maps

    --[

    local sum = maps[1]

    for i=2,#maps do

        sum = sum + maps[i]

    end

    local map = sum/#maps

    map = image.convolve(map, kernel, 'same')

    map = map/torch.max(map)                    -- normalize

    image.save('test.png',map)

    --map = image.load('test.png',1,'float')

    local results = {}

    local m = 1

    while true do

        local ind = map:eq(map:max()):nonzero()

        local row = ind[1][1]

        local col = ind[1][2]

        local val = map[row][col]

        results[m] = {row, col, val}

        m = m + 1

        for i=-2*d,2*d do

            for j=-2*d,2*d do

                if i^2 + j^2 < 4*d^2 and row+i>0 and col+j>0 and row+i<=map:size(1) and col+j<=map:size(2) then

                    map[row+i][col+j] = 0

                end

            end

        end

        if val < threshold then

            break

        end

    end

    local outfile = paths.concat(folder, 'results', paths.basename(imagePath, paths.extname(imagePath)) .. '.csv')

    csvigo.save(outfile, results)

    print(os.clock()-c1)

    print(os.time()-t1)

end

print(os.clock()-c)

print(os.time()-t)