include("./utils.jl")
include("neural_net.jl")
using Flux.Tracker
using Flux: params, mse, crossentropy
using ProgressMeter
using CuArrays

function total_loss(y, ŷ, ps)
    (π, z), (p, v)  = y, ŷ
    crossentropy(p, π) + mse(z, v) + 0.0001f0*sum([sum(i.^2) for i in ps])
end

@inline function all_equal(x)
    e1 = x[1]
    @inbounds for i=2:30
        x[i] == e1 || return false
    end
    return true
end

# function retrain(curr_nn::NeuralNet, epoch::Int=1000)
#     curr_nn = mapleaves(cu, curr_nn)
#     total = 0.
#     losses = []
#     opt = ADAM()
#     anim, plt = get_animPlot()
#     @showprogress "training network: " for i=1:epoch
#         data = load_data()
#         x = data[1] |> gpu
#         π = data[2] |> gpu
#         z = data[3] |> gpu
#         y = (π, z)
#         ps, ŷ = params(curr_nn), curr_nn(x)
#         @sync begin
#             @async begin
#                 loss = total_loss(y, ŷ, ps).data
#                 total += loss
#                 push!(plt, [loss; total/i])
#                 push!(losses, [loss; total/i])
#                 frame(anim, plt)
#             end
#             @async begin gs = gradient(ps) do
#                 total_loss(y, ŷ, ps)
#                 end
#                 Tracker.update!(opt, ps, gs)
#             end
#         end
#         i>100 && all_equal(losses[end-30:end]) && break
#     end
#     save_retrain(anim, plt, losses)
#     save_model(curr_nn, true)
#     return curr_nn
# end
