require 'torch'
require 'image'
require 'paths'
require 'pl'
require 'cudnn'
NN_UTILS = require 'utils.nn_utils'
DATASET = require 'dataset_rgb'

OPT = lapp[[
    --save          (default "logs")                          Directory in which the networks are stored.
    --network       (default "adversarial.net")               Filename of the network to use.
    --neighbours                                              Whether to search for nearest neighbours of generated images in the dataset (takes long)
    --writeto       (default "samples")                       Directory to save the images to
    --seed          (default 1)                               Random number seed to use.
    --gpu           (default 0)                               GPU to run on
    --runs          (default 1)                               How often to sample and save images
    --noiseDim      (default 100)                             Noise vector size.
    --batchSize     (default 16)                              Sizes of batches.
    --dataset       (default "NONE")                          Directory that contains *.jpg images
]]

if OPT.gpu < 0 then
    print("[ERROR] Sample script currently only runs on GPU, set --gpu=x where x is between 0 and 3.")
    exit()
end

-- Start GPU mode
print("Starting gpu support...")
require 'cutorch'
require 'cunn'
torch.setdefaulttensortype('torch.FloatTensor')
cutorch.setDevice(OPT.gpu + 1)

-- initialize seeds
math.randomseed(OPT.seed)
torch.manualSeed(OPT.seed)
cutorch.manualSeed(OPT.seed)

-- Initialize dataset
DATASET.setFileExtension("jpg")

-- Main function that runs the sampling
function main()
    -- Load all models
    local G, D, height, width, dataset = loadModels()

    -- Image dimensions
    IMG_DIMENSIONS = {3, height, width}
    NOISE_DIM = {1, height, width}

    DATASET.setHeight(height)
    DATASET.setWidth(width)
    if OPT.dataset ~= "NONE" or dataset == nil then
        DATASET.setDirs({OPT.dataset})
    else
        DATASET.setDirs({dataset})
    end

    print("Sampling...")
    for run=1,OPT.runs do
        -- save 64 randomly selected images from the training set
        local imagesTrainList = DATASET.loadRandomImages(64)
        -- dont use nn_utils.toImageTensor here, because the metatable of imagesTrainList was changed
        local imagesTrain = torch.Tensor(#imagesTrainList, imagesTrainList[1].grayscale:size(1), imagesTrainList[1].grayscale:size(2), imagesTrainList[1].grayscale:size(3))
        for i=1,#imagesTrainList do
            imagesTrain[i] = imagesTrainList[i].grayscale
        end
        image.save(paths.concat(OPT.writeto, string.format('trainset_s1_%04d_base.jpg', run)), toGrid(imagesTrainList, nil, 8))

        -- sample 64 colorizations from G
        local noise = torch.Tensor(64, NOISE_DIM[1], NOISE_DIM[2], NOISE_DIM[3])
        noise:uniform(0, 1)
        local imagesGenerated = G:forward({noise, imagesTrain})

        -- validate image dimensions
        if imagesGenerated[1]:size(1) ~= IMG_DIMENSIONS[1] or imagesGenerated[1]:size(2) ~= IMG_DIMENSIONS[2] or imagesGenerated[1]:size(3) ~= IMG_DIMENSIONS[3] then
            print("[WARNING] dimension mismatch between images generated by base G and command line parameters")
            print("Dimension G:", images[1]:size())
            print("Settings:", IMG_DIMENSIONS)
        end

        -- save big images of those 1024 random images
        image.save(paths.concat(OPT.writeto, string.format('random64_%04d.jpg', run)), toGrid(imagesTrainList, imagesGenerated, 12))

        xlua.progress(run, OPT.runs)
    end

    print("Finished.")
end

-- Converts images to one image grid with set amount of rows.
-- @param images Tensor of images
-- @param nrow Number of rows.
-- @return Tensor
function toGrid(imagesOriginal, imagesGenerated, nrow)
    local N = 2 * imagesOriginal:size()
    if imagesGenerated ~= nil then N = N + imagesGenerated:size(1) end
    local images = torch.Tensor(N, 3, imagesOriginal[1].color:size(2), imagesOriginal[1].color:size(3))
    local idx = 1
    for i=1,#imagesOriginal do
        images[{{idx}, {1}, {}, {}}] = imagesOriginal[i].grayscale
        images[{{idx}, {2}, {}, {}}] = imagesOriginal[i].grayscale
        images[{{idx}, {3}, {}, {}}] = imagesOriginal[i].grayscale
        images[idx+1] = imagesOriginal[i].color
        if imagesGenerated ~= nil then
            images[idx+2] = imagesGenerated[i]
            idx = idx + 1
        end
        idx = idx + 2
    end

    return image.toDisplayTensor{input=images, nrow=nrow}
end

-- Loads all necessary models/networks and returns them.
-- @returns G, D, height, width, dataset directory
function loadModels()
    local file = torch.load(paths.concat(OPT.save, OPT.network))

    local G = file.G
    local D = file.D
    local opt_loaded = file.opt
    G:evaluate()
    D:evaluate()

    return G, D, opt_loaded.height, opt_loaded.width, opt_loaded.dataset
end

main()
