include("../network_functions.jl")

using Plots: plot, plot!, display
using MLDatasets: MNIST
using Random: seed!
using CuArrays
using Images
# set random seed
seed!(1301)
# disallow scalar operations
CuArrays.allowscalar(false)


# ::: DATA
train_x, _ = MNIST.traindata()
train_x = permutedims(Array{Float32}(train_x), [3,2,1])
train_x = reshape(train_x, (length(train_x[:, 1, 1]), length(train_x[1, :, 1]) * length(train_x[1, :, 1])))
get_MNIST_input_batch(x) = begin
    return CuArray(x[rand(1:length(x[:, 1])), :])
end

# ::: PARAMS
INPUT_SIZE = length(train_x[1, :, :])
LATENT_SIZE = 70
BATCH_SIZE = 1
LEARNING_RATE = 0.04
TRAIN_EPISODES = 10000

# optimizer
opt = SGD(LEARNING_RATE)

# training data
train_data_x = [get_MNIST_input_batch(train_x) for _ in 1:TRAIN_EPISODES]



# save test image for comparison
function generate_images(number_of_images, path)
    for n in 1:number_of_images
        all_test_image_ref = get_MNIST_batch(train_x, BATCH_SIZE)
        all_test_images =  all_test_image_ref |> encoder |> decoder |> Tracker.data

        random_index = rand(1:BATCH_SIZE)
        test_image = all_test_images[:, random_index]
        test_image_ref = all_test_image_ref[:, random_index]

        test_image_ref_r = reshape(test_image_ref, (28, 28))
        test_image_r = clamp.(reshape(test_image, (28, 28)), 0, 1)
        save("$path/$(n)_reference_.png", colorview(Gray, test_image_ref_r))
        save("$path/$(n)_reconstruction.png", colorview(Gray, test_image_r))
    end
end
