using MuZero
using Base.Threads: @spawn

include("game.jl")
include("inference_tests.jl")
cd(@__DIR__)
config = getconfig(Tac)
config.batch_size = 32
config.root_dirichlet_alpha=0.3
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
config.training_steps = 200001

p_tests = getpolicytests(Tac)
p_tests.period = 250

learner = Learner(
    Tac; 
    dir="output", 
    config=config, 
    tests=[p_tests]
)

@spawn selfplay!(learner)
@spawn selfplay!(learner)
@spawn selfplay!(learner)
length(learner.replay_buffer)

@show @time fit!(learner; steps=3)
@spawn fit!(learner)

learner.stats.losses |> length
learner |> gettraindf |> display
learner |> gettraindf |> mean1000 |> display
learner |> gettraindfall |> mean1000 |> display


learner |> show_pv_tests
learner |> show_reward_tests

#@learnertest 4854 |> svg
#@learnersample 33

selfplay!(learner; limit=10)

learner |> show_game

# still a fence post problem here
s = learner |> randtrainingsample
learner |> randtrainingsample |> svg

learner |> randtrainingbatch

using Base.Threads: @spawn
for i = 1:num_actors(learner)
    @spawn selfplay!(learner)
end

length(learner.replay_buffer)

@show get_steps(learner), learner.stats.losses, learner.stats.accuracies, learner.times

@time fit!(learner, steps=1)

learner |> gettraindfall |> display

@show learner.stats.losses, learner.stats.accuracies, learner.times

get_steps(learner), get_accuracy(learner)
