import Pkg
Pkg.rm(["Grids", "GridGames"])
Pkg.develop(["Grids", "GridGames"])

using CUDA

# Make sure no errors or 
# warnings with CUDA
CUDA.versioninfo()

using Base.Threads: @spawn
using MuZero
using GridGames
cd(@__DIR__)

# You need at least 5 
# threads.
Threads.nthreads()

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

# Play a few sample games to make sure the
# game dynamics are correct. Need to copy and 
# paste "play" function into REPL, if using VSCode.
#
# play(Tac)

# Initialize MuZero configuration
include("muzero_config.jl")
config = getconfig(Tac)
config.num_unroll_steps = 5
config.num_actors = 1

# Initialize the policy tests
include("inference_tests.jl")
rprv_tests = makerrvptests(Tac)
rprv_tests.period = 1000
pv_tests = makepolicyvaluetests(Tac)
pv_tests.period = 1000
#valuefn_debug = getvaluefn(Tac)

# Configure Learner
config.num_actors
config.num_unroll_steps
learner = Learner(
    config;
    uuid="b234fc40-a07b-4418-9fbc-adb0278dd4e1",
    dir="~/models/MuZeroModels.jl/tictactoe/",
    tests=[pv_tests, rprv_tests], #
    schedule=tacschedule!,#cycleschedule(min_lr=1e-6, max_lr=9e-3),#
    #valuefn_debug=valuefn_debug
)

xyz = randrrvptest(learner)
xyz.state |> display
xyz.state_next |> display

getpolicyvaluetests(learner)[300] |> svg

l = Learner(Tac)
Learner(Tac; uuid="20B4993D-D1B8-41B7-ABB2-9BC8C70EA5F9")
Learner(Tac; dir="~/models/MuZeroModels.jl/tictactoe/")
Learner(
    Tac;
    uuid="8bcecce1-c870-4e34-8ae9-df6433d033de",
    dir="./models/MuZeroModels.jl/tictactoe/"
)

Learner(config)
Learner(config; uuid="20B4993D-D1B8-41B7-ABB2-9BC8C70EA5F9")
Learner(config; dir="~/models/MuZeroModels.jl/tictactoe/")
Learner(
    config;
    uuid="8bcecce1-c870-4e34-8ae9-df6433d033de",
    dir="./models/MuZeroModels.jl/tictactoe/"
)

Learner(Tac, config)
Learner(Tac, config; uuid="20B4993D-D1B8-41B7-ABB2-9BC8C70EA5F9")
Learner(Tac, config; dir="~/models/MuZeroModels.jl/tictactoe/")
Learner(
    Tac, 
    config;
    uuid="8bcecce1-c870-4e34-8ae9-df6433d033de",
    dir="./models/MuZeroModels.jl/tictactoe/"
)



getpvtestresult(learner, 300) |> svg
acc, fails, value_mse = runpvtests(learner)
reward_mse, value_mse, acc, fails = runrrvptests(learner)
getpolicy(learner, getstate(Tac([2,8,6]))) |> display
getvalue(learner, getstate(Tac([2,8,6])), 1) |> display

# Let's make sure our policy test cases look reasonable
randpolicyvaluetest(learner) |> svg
randpolicytest(learner) |> svg
randpolicytest(learner) |> svg
randpolicytest(learner) |> svg



showrrvptestresults(learner)



# The predictions from this untrained model should be random
# and mostly incorrect.
showpolicyvaluetestresults(learner)
randpolicyvaluetestresult(learner) |> svg

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

# Show a random self-played game
g = randgame(learner)
g |> display
g |> svg
randtrainingsample(learner, g) |> display

# Get a random batch of training samples
randtrainingbatch(learner) |> display
randtrainingbatch(learner)[1] |> svg

# If game play seems to be working properly, then start indefinite self play on multiple threads.
r1 = @spawn play!(learner)

# Let's make sure threading is working propery and the games are being recorded
length(learner.replay_buffers)
getavggamestep(learner)
getgamestepquality(learner)
get_steps(learner)
r2 = @spawn fit!(learner, steps=1100)

learner.is_playing = true
res3 = @spawn selfplay!(learner)
learner.is_playing = false


updatecallback!(
    learner,
    LearnerCallback(
        "SaveGamePlayStats", 
        savegameplaystats!,
        period=1000,
        event=:afterstep
    )
)


learner.is_training = false
updateschedule!(learner, tacschedule!)
updateschedule!(learner, newcycleschedule(learner, min_lr=1e-5, max_lr=3e-5, cycle_steps=10_000))

# 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.
@show @time fit!(learner, steps=3)

# Let's ensure the initial loss was stored, which means the model is training
showlosses(learner)

# View training progress
showresults(learner)

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

# View training progress. We should see steady decrease
# in the loss and a steady increase in the accuracy.
showresults(learner)

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

fit!(learner; steps=3)




# 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)










@show @time selfplay!(learner; limit=3)
@show @time selfplay!(learner; limit=3)

using BenchmarkTools
@btime selfplay!(learner; limit=3)

@btime selfplay!(learner; limit=10)

@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
stoptraining!(learner)
stopplaying!(learner)
clearreplaybuffer!(learner)
length(learner.replay_buffer)


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)


# STOPPED HERE


tor = "/home/ubuntu/models/MuZeroModels.jl/tictactoe/eb14fc40-a07b-4418-9fbc-adb0278dd4e1/replay_buffers/e4ffabbe-1638-4580-8cc6-b57c8e7c6dbd/buffer.bin"
tor2 = "/home/ubuntu/models/MuZeroModels.jl/tictactoe/eb14fc40-a07b-4418-9fbc-adb0278dd4e1/replay_buffers/d5281c62-455a-4bd4-a145-0ff86597d058/buffer.bin"
reanalyzebuffer(learner, tor2)



# DEBUGGING PV FAILS
getpolicyvaluetestresults(learner)[1935:1935]
fit!(learner, steps=100) ; learner.stats.policy_test_fails

pvdf = getdf(getpolicyvaluetestresults(learner))
pvdf[learner.stats.policy_test_fails, :]

getpolicyvaluetestresults(learner)[1935:1935]
pt1935 = getpolicyvaluetests(learner)[1935]
getpolicy(learner, pt1935.state)
getvalue(learner, pt1935.state, 5)
getvalue(learner, pt1935.state, 2)
getvalue(learner, pt1935.state, 3)
getvalue(learner, pt1935.state, 7)

pt1154 = getpolicyvaluetests(learner)[1154]
pt1154 |> svg
getpolicy(learner, pt2438.state)

pvdf[pvdf."P/F" .== "Fail", :]



using Plots
plot(getlosses(learner))


# LR finder



writable = learner.replay_buffers.d[learner.replay_buffers.writable_uuid]
g1 = writable[end]
g1.net_step
g1
g1 |> svg
g1s5 = get_state(g1, 5)
g1s5 |> typeof

g2 = writable[end-1]



using Profile
@time play_game(learner)

@profile play_game(learner)

open("/tmp/prof.txt", "w") do s
    Profile.print(IOContext(s))
end


learner.play.representation_network[1].weight
learner.play.representation_network[2]
learner.play.representation_network[2].μ
learner.play.representation_network[2].μ
learner.play.representation_network[2].active
learner.train.representation_network[2].active
learner.test.representation_network[2].active
