include("auto_encoder_task.jl")

# ::: NETWORK
networks = Dict(
    "full_sigmoid" => ([
        DenseLayer(INPUT_SIZE, 500),
        ActivationLayer(500, diminish, Δdiminish),
        DenseLayer(500, 300),
        ActivationLayer(300, diminish, Δdiminish),
        DenseLayer(300, 150),
        ActivationLayer(150, diminish, Δdiminish),
        DenseLayer(150, LATENT_SIZE),
        ActivationLayer(LATENT_SIZE, diminish, Δdiminish)
        ],[
        DenseLayer(LATENT_SIZE, 150),
        ActivationLayer(150, diminish, Δdiminish),
        DenseLayer(150, 300),
        ActivationLayer(300, diminish, Δdiminish),
        DenseLayer(300, 500),
        ActivationLayer(500, diminish, Δdiminish),
        DenseLayer(500, INPUT_SIZE),
        ActivationLayer(INPUT_SIZE, diminish, Δdiminish)
        ]),
   "partial_sigmoid1" => ([
        DenseLayer(INPUT_SIZE, 500),
        ActivationLayer(500, relu, Δrelu),
        DenseLayer(500, 300),
        ActivationLayer(300, diminish, Δdiminish),
        DenseLayer(300, 150),
        ActivationLayer(150, diminish, Δdiminish),
        DenseLayer(150, LATENT_SIZE),
        ActivationLayer(LATENT_SIZE, diminish, Δdiminish)
        ],[
        DenseLayer(LATENT_SIZE, 150),
        ActivationLayer(150, relu, Δrelu),
        DenseLayer(150, 300),
        ActivationLayer(300, diminish, Δdiminish),
        DenseLayer(300, 500),
        ActivationLayer(500, diminish, Δdiminish),
        DenseLayer(500, INPUT_SIZE),
        ActivationLayer(INPUT_SIZE, diminish, Δdiminish)
        ]),
    "partial_sigmoid2" => ([
        DenseLayer(INPUT_SIZE, 500),
        ActivationLayer(500, relu, Δrelu),
        DenseLayer(500, 300),
        ActivationLayer(300, relu, Δrelu),
        DenseLayer(300, 150),
        ActivationLayer(150, diminish, Δdiminish),
        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, diminish, Δdiminish),
        DenseLayer(500, INPUT_SIZE),
        ActivationLayer(INPUT_SIZE, diminish, Δdiminish)
        ]),
    "partial_sigmoid3" => ([
        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)
        ]))



# ::: TRAINING
runtime_plotting = false
plot_intervall = 20
begin
    full_sigmoid_tm = [NaN for _ in 1:TRAIN_EPISODES];
    partial_sigmoid_tm1 = [NaN for _ in 1:TRAIN_EPISODES];
    partial_sigmoid_tm2 = [NaN for _ in 1:TRAIN_EPISODES];
    partial_sigmoid_tm3 = [NaN for _ in 1:TRAIN_EPISODES];


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

        z1 = forward!(networks["full_sigmoid"][1], train_data_x[e])
        z2 = forward!(networks["partial_sigmoid1"][1], train_data_x[e])
        z3 = forward!(networks["partial_sigmoid2"][1], train_data_x[e])
        z4 = forward!(networks["partial_sigmoid3"][1], train_data_x[e])

        y1 = forward!(networks["full_sigmoid"][2], z1)
        y2 = forward!(networks["partial_sigmoid1"][2], z2)
        y3 = forward!(networks["partial_sigmoid2"][2], z3)
        y4 = forward!(networks["partial_sigmoid3"][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["full_sigmoid"][2]), dloss1, opt)
        _ = backward!(reverse(networks["full_sigmoid"][1]), dz1, opt)
        dz2 = backward!(reverse(networks["partial_sigmoid1"][2]), dloss2, opt)
        _ = backward!(reverse(networks["partial_sigmoid1"][1]), dz2, opt)
        dz3 = backward!(reverse(networks["partial_sigmoid2"][2]), dloss3, opt)
        _ = backward!(reverse(networks["partial_sigmoid2"][1]), dz3, opt)
        dz4 = backward!(reverse(networks["partial_sigmoid3"][2]), dloss4, opt)
        _ = backward!(reverse(networks["partial_sigmoid3"][1]), dz4, opt)


        full_sigmoid_tm[e] = mse(train_data_x[e], y1)
        partial_sigmoid_tm1[e] = mse(train_data_x[e], y2)
        partial_sigmoid_tm2[e] = mse(train_data_x[e], y3)
        partial_sigmoid_tm3[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(full_sigmoid_tm[i:i+gr])/gr for i in 1:TRAIN_EPISODES-gr], xlabel="training iterations", label="FC (full diminish): MSE (lr=$LEARNING_RATE, batch size=$BATCH_SIZE)")
    plot!([sum(partial_sigmoid_tm1[i:i+gr])/gr for i in 1:TRAIN_EPISODES-gr], xlabel="training iterations", label="FC (Relu [1] + diminish [3]): MSE (lr=$LEARNING_RATE, batch size=$BATCH_SIZE)")
    plot!([sum(partial_sigmoid_tm2[i:i+gr])/gr for i in 1:TRAIN_EPISODES-gr], xlabel="training iterations", label="FC (Relu [2] + diminish [2]): MSE (lr=$LEARNING_RATE, batch size=$BATCH_SIZE)")
    plot!([sum(partial_sigmoid_tm3[i:i+gr])/gr for i in 1:TRAIN_EPISODES-gr], xlabel="training iterations", label="FC (Relu [3] + diminish [1]): MSE (lr=$LEARNING_RATE, batch size=$BATCH_SIZE)")
end






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