include("auto_encoder_task.jl")

# ::: NETWORK
networks = Dict(
    "500_300_150" => ([
        DenseLayer(INPUT_SIZE, 500),
        ActivationLayer(500, relu, Δrelu),
        DenseLayer(500, 300),
        ActivationLayer(300, relu, Δrelu),
        DenseLayer(300, 150),
        ActivationLayer(150, relu, Δrelu),
        DenseLayer(150, LATENT_SIZE),
        ActivationLayer(LATENT_SIZE, diminish, Δdiminish)
        ],[
        DenseLayer(LATENT_SIZE, 150),
        ActivationLayer(150, relu, Δrelu),
        DenseLayer(150, 300),
        ActivationLayer(300, relu, Δrelu),
        DenseLayer(300, 500),
        ActivationLayer(500, relu, Δrelu),
        DenseLayer(500, INPUT_SIZE),
        ActivationLayer(INPUT_SIZE, diminish, Δdiminish)
        ]),
   "600_400_200_100" => ([
        DenseLayer(INPUT_SIZE, 600),
        ActivationLayer(600, relu, Δrelu),
        DenseLayer(600, 400),
        ActivationLayer(400, relu, Δrelu),
        DenseLayer(400, 200),
        ActivationLayer(200, relu, Δrelu),
        DenseLayer(200, 100),
        ActivationLayer(100, relu, Δrelu),
        DenseLayer(100, LATENT_SIZE),
        ActivationLayer(LATENT_SIZE, diminish, Δdiminish)
        ],[
        DenseLayer(LATENT_SIZE, 100),
        ActivationLayer(100, relu, Δrelu),
        DenseLayer(100, 200),
        ActivationLayer(200, relu, Δrelu),
        DenseLayer(200, 400),
        ActivationLayer(400, relu, Δrelu),
        DenseLayer(400, 600),
        ActivationLayer(600, relu, Δrelu),
        DenseLayer(600, INPUT_SIZE),
        ActivationLayer(INPUT_SIZE, diminish, Δdiminish)
        ]),
    "500_500_250_250_125_125" => ([
        DenseLayer(INPUT_SIZE, 500),
        ActivationLayer(500, relu, Δrelu),
        DenseLayer(500, 500),
        ActivationLayer(500, relu, Δrelu),
        DenseLayer(500, 250),
        ActivationLayer(250, relu, Δrelu),
        DenseLayer(250, 250),
        ActivationLayer(250, relu, Δrelu),
        DenseLayer(250, 125),
        ActivationLayer(125, relu, Δrelu),
        DenseLayer(125, 125),
        ActivationLayer(125, relu, Δrelu),
        DenseLayer(125, LATENT_SIZE),
        ActivationLayer(LATENT_SIZE, diminish, Δdiminish)
       ],[
        DenseLayer(LATENT_SIZE, 125),
        ActivationLayer(125, relu, Δrelu),
        DenseLayer(125, 125),
        ActivationLayer(125, relu, Δrelu),
        DenseLayer(125, 250),
        ActivationLayer(250, relu, Δrelu),
        DenseLayer(250, 250),
        ActivationLayer(250, relu, Δrelu),
        DenseLayer(250, 500),
        ActivationLayer(500, relu, Δrelu),
        DenseLayer(500, 500),
        ActivationLayer(500, relu, Δrelu),
        DenseLayer(500, INPUT_SIZE),
        ActivationLayer(INPUT_SIZE, diminish, Δdiminish)
        ]),
    "500" => ([
        DenseLayer(INPUT_SIZE, 500),
        ActivationLayer(500, relu, Δrelu),
        DenseLayer(500, LATENT_SIZE),
        ActivationLayer(LATENT_SIZE, diminish, Δdiminish)
       ],[
        DenseLayer(LATENT_SIZE, 500),
        ActivationLayer(500, relu, Δrelu),
        DenseLayer(500, INPUT_SIZE),
        ActivationLayer(INPUT_SIZE, diminish, Δdiminish)
        ]))



# ::: TRAINING
runtime_plotting = false
plot_intervall = 20
begin
    tm1 = [NaN for _ in 1:TRAIN_EPISODES];
    tm2 = [NaN for _ in 1:TRAIN_EPISODES];
    tm3 = [NaN for _ in 1:TRAIN_EPISODES];
    tm4 = [NaN for _ in 1:TRAIN_EPISODES];


    println("---"); @time for e in 1:TRAIN_EPISODES
        t = time()

        z1 = forward!(networks["500_300_150"][1], train_data_x[e])
        z2 = forward!(networks["600_400_200_100"][1], train_data_x[e])
        z3 = forward!(networks["500_500_250_250_125_125"][1], train_data_x[e])
        z4 = forward!(networks["500"][1], train_data_x[e])

        y1 = forward!(networks["500_300_150"][2], z1)
        y2 = forward!(networks["600_400_200_100"][2], z2)
        y3 = forward!(networks["500_500_250_250_125_125"][2], z3)
        y4 = forward!(networks["500"][2], z4)

        dloss1 = -Δmse(train_data_x[e], y1)
        dloss2 = -Δmse(train_data_x[e], y2)
        dloss3 = -Δmse(train_data_x[e], y3)
        dloss4 = -Δmse(train_data_x[e], y4)


        dz1 = backward!(reverse(networks["500_300_150"][2]), dloss1, opt)
        _ = backward!(reverse(networks["500_300_150"][1]), dz1, opt)
        dz2 = backward!(reverse(networks["600_400_200_100"][2]), dloss2, opt)
        _ = backward!(reverse(networks["600_400_200_100"][1]), dz2, opt)
        dz3 = backward!(reverse(networks["500_500_250_250_125_125"][2]), dloss3, opt)
        _ = backward!(reverse(networks["500_500_250_250_125_125"][1]), dz3, opt)
        dz4 = backward!(reverse(networks["500"][2]), dloss4, opt)
        _ = backward!(reverse(networks["500"][1]), dz4, opt)


        tm1[e] = mse(train_data_x[e], y1)
        tm2[e] = mse(train_data_x[e], y2)
        tm3[e] = mse(train_data_x[e], y3)
        tm4[e] = mse(train_data_x[e], y4)


        # if runtime_plotting && e % plot_intervall == 0
        #     p = plot(FC_train_metric)
        #     display(p)
        # end

        # println("episode = $e : loss = $loss : dt = $(time() - t)")
    end
end


begin
    gr = 30
    plot([sum(tm1[i:i+gr])/gr for i in 1:TRAIN_EPISODES-gr], ylabel="MSE (lr=$LEARNING_RATE, batch size=$BATCH_SIZE)", xlabel="training iterations", label="FC [500_300_150]")
    plot!([sum(tm2[i:i+gr])/gr for i in 1:TRAIN_EPISODES-gr], xlabel="training iterations", label="FC [600_400_200_100]")
    plot!([sum(tm3[i:i+gr])/gr for i in 1:TRAIN_EPISODES-gr], xlabel="training iterations", label="FC [500_500_250_250_125_125]")
    plot!([sum(tm4[i:i+gr])/gr for i in 1:TRAIN_EPISODES-gr], xlabel="training iterations", label="FC [500]")
end






generate_images(10, "/Users/BBM2H16AHO/Documents/programming/DNN-Framework/auto_encoder/images/relu_sigmoid_test")
