
# Making initiaizations for the environment and packages

println("Loading environment...")
sleep(1.5)
using Pkg

try
    Pkg.activate("fc2finalproject/")
    Pkg.status()
catch
    error("Environment was either not found or failed loading")
end

println("Loading packages...")

begin
    using Flux
    using Flux: @epochs
    using Plots, Statistics
    using ProgressMeter, Alert
    using BSON: @save
    using BSON: @load
    using DelimitedFiles
    gr()
end

println("Packages loaded successfully")
println("Loading auxiliary functions...")
sleep(1.5)
function tuple2vec(test::Tuple)
    return [test[1] , test[2]]
end


function generatetraindata(tuplemesh)
    row, col = size(tuplemesh)
    data = Matrix{Float64}(undef, 2, row*col)
    
    current = 0 # initialize
    for k in 1:col, m in 1:row
        current += 1
        data[:,current] = tuple2vec( tuplemesh[m,k] )

    end
    return data
end

println("Auxiliary functions loaded successfully")
println("Loading data...")
sleep(1.5)

intensities = readdlm("data.csv",',')

ydims, xdims = size(intensities)

sleep(1.5)
println("Your picture is of size $xdims px by $ydims px.")

sleep(1.5)
println("Generating training data...")

z_orthogonal_complement = [ (x,y) for y in 1:xdims, x in 1:ydims ]
Xtrain = generatetraindata(z_orthogonal_complement)
Ytrain = Float32.( Matrix( vec(intensities)' ) )
data = Flux.Data.DataLoader((Xtrain, Ytrain), batchsize=100,shuffle=true);

println("Data loaded successfully")
sleep(1.5)

goal = heatmap( unique(Xtrain[1,:]), unique(Xtrain[2,:]),
reverse(reshape(Ytrain, ydims,xdims)), c=:grays, clims = (0,1));
savefig("goal.pdf")



println("Finding last save state model.bson")
search = findall(x-> x.== "model.bson", readdir() )

if length(search) < 1

    println("Could not find a .bson file with a model")
    sleep(1.5)
    println("Creating model from scrach")
    sleep(1.5)

    function NeuralNetwork()
        return Chain(
                Dense(2, 200,relu),
                Dense(200, 100,relu),
                Dense(100, 50,relu),
                Dense(50, 25,relu),
                Dense(25, 10,relu),
                Dense(10,1,x->σ.(x))
                )
    end

    global model = NeuralNetwork()
    println("New model created successfully")
    sleep(1.5)
    
    im0 = heatmap( unique(Xtrain[1,:]), unique(Xtrain[2,:]),
    reverse(reshape(model(Xtrain), ydims,xdims)), c=:grays, clims = (0,1));
    savefig("History/initialguess.pdf")
else
    sleep(1.5)
    @load "model.bson"  model
    println("model.bson loaded successfully")
    sleep(1.5)

end

opt = Descent(0.05)
loss(x, y) = sum(Flux.Losses.binarycrossentropy(model(x), y))


save_state = 0; #initialize
total_iterations = 0; # initialize
ps = Flux.params(model);
epochs = 100000;

@showprogress 1 "Training..." for i in 1:epochs
    global save_state
    global total_iterations
    Flux.train!(loss, ps, data, opt)

    if (i % 5000) == 0
        # println("Saving... $i iterations have been completed")
        save_state += 1
        total_iterations += 5000
        @save "model.bson" model
        @save "Checkpoints/model_c$(save_state).bson" model
        image = heatmap( unique(Xtrain[1,:]), unique(Xtrain[2,:]),
        reverse(reshape(model(Xtrain), ydims,xdims)), c=:grays, clims = (0,1) );
        savefig("History/result$(total_iterations).pdf") 
    end

end






