using MuZero
import Zygote # @todo figure out how to remove

# User defined state and game implementations
include("state.jl")
include("game.jl")

# Configure Learner
learner = Learner(TicTacToe)

# Let's make sure our policy test cases look reasonable.
learner |> tests |> svg

# The predictions from this untrained model should be random
# and mostly incorrect.
learner |> test_results |> svg

#learner |> policy_tests |> first |> svg
#pt = learner |> policy_tests |> first
#policy_test_result = PolicyTestResult(pt, Flux.softmax(collect(rand(9))))

# Self-play 10 sample games to ensure proper game implementation.
selfplay!(learner; limit=10)

# Show a random self-played game
learner |> randgame |> svg

# Get a random training sample
learner |> randtrainingsample |> svg

# Get a random batch of training samples
learner |> randtrainingbatch |> show

# If game play seems to be working properly, then start indefinite self play on multiple threads.
using Base.Threads: @spawn
for i = 1:num_actors(learner)
    @spawn selfplay!(learner)
end

# Let's make sure threading is working propery and the games are being recorded
length(learner.replay_buffer)

# Train three steps to make sure there are no errors. This 1st step
# may take a several minutes if your machine has no usable
# Nvidia GPU.
@time fit!(learner, steps=3)

# Let's ensure the initial loss was stored, which means the model is training
learner |> losses |> svg

# Let's see how close our predictions are to their targets
learner |> rand_training_result |> svg
#learner |> training_results |> svg

# Visualize test set errors
learner |> test_results |> svg

# Train network for a few more steps
fit!(learner, steps=20)

# We should see steady decreases in the loss
learner |> losses |> svg

# We should see steady increase in the accuracy
learner |> accuracy |> svg

# Let's train indefinitely on a separate thread and periodically check the 
# progress.
res = @spawn fit!(learner)


# function losses_svg()
#     function Base.show(io::IO, m::MIME"image/svg+xml", g::GridSVG)
#         draw(g)
#         print(io, read(output_file(g), String))
#     end
# end

# # Visualize any game from batch using the encoded state key
# df = get_df(batch)
# state = TicTacToeState(df."State Key"[1])
# state |> show

# # Or, use the following convenience function to get the state
# # @todo Stopped here!
# get_state(batch, 1) |> show

# Learning rate finder
# @todo implement

# Open up "losses.png" and "accuracy.png" in separate windows
# to view training loss and testing accuracy graphs in real-time.

# Alternatively, print the "losses" and/or "accuracy" arrays
# in the REPL to view progress.
losses
accuracy

# If losses is empty and/or not increasing, see status.txt
# to check if there was an error during training.
`cat ./status.txt` |> run

# Also, you can check "fetch(result)" to see if there was an exception.
# If no exception prints (and the call to fetch hangs), then the training is still 
# starting, and you can interrupt the call to fetch
fetch(result)

# Export Flux model
MuZero.export(learner)