# Full Training Test (needs GPU for reasonable training time)using MuZero
using MuZero
import Zygote
using Grids
include("../../examples/tictactoe/state.jl")
include("../../examples/tictactoe/game.jl")
include("./old.jl")
include("./inference_tests.jl")
using Base.Threads: @spawn, @threads
config = make_tictactoe_config()
config.batch_size = 16
config.num_actors = 3
config.weight_decay = 1e-4
config.lr_init=0.002f0
config.save_interval = 10000
config.num_simulations = 100
config.num_unroll_steps = 1
losses, dynamics_stats = [], []

p_tests = getpolicytests(TicTacToe)
p_tests.period = 1000

policy_perfs = []; 
function calc_policy_perf(l)
    push!(policy_perfs, policy_perf(l.test))
end

test_cb = LearnerCallback(
    "test_cb",
    calc_policy_perf,
    period=1000
)

learner = Learner(
    TicTacToe; 
    config=config,
    dir="output",
    cb=[test_cb],
    tests=[p_tests]
)

#selfplay!(learner; limit=config.batch_size)
# @spawn run_selfplay!(config, storage, replay_buffer; limit=33)
# @spawn run_selfplay!(config, storage, replay_buffer; limit=33)
# @spawn run_selfplay!(config, storage, replay_buffer; limit=34)
# rand(1:length(replay_buffer))

for i = 1:(config.num_actors)
    @spawn selfplay!(learner)
end
learner.replay_buffer |> length

#cb = [(1000, 1, [calc_policy_perf]),]

#config.training_steps = 3
#train_network!(learner, losses=losses, callbacks=cb)

config.training_steps = 2000
@show @time fit!(learner; steps=3)
@spawn fit!(learner)
#@spawn train_network!(learner, losses=losses, callbacks=cb)
#@spawn train_network!(config, storage, replay_buffer, losses=losses, callbacks=cb)

@show 1
learner.stats.losses |> length
@show 2
