using Plots
using Random
using DataFrames
using UUIDs
using TOML
using Printf
using Dates
import CSV
import BSON
using Statistics
using GridGames
using Serialization
using Base.Threads: @spawn

struct LearnerCallback
    name::String
    event::Symbol
    fn::Function
    period::Int
    warmup::Int
    function LearnerCallback(name, fn; event=:afterstep, period=1, warmup=0)
        new(name, event, fn, period, warmup)
    end
end

tuple(x::LearnerCallback) = x.fn, x.period, x.warmup

mutable struct LearnerStats
    step_count::Counter
    losses::Dict{UInt64,Loss}
    accuracies::Dict{UInt64,PolicyAcc}
    reward_errors::Dict{UInt64,Float16}
    value_errors::Dict{UInt64,Float16}
    times::Dict{UInt64,Float16}
    test_times::Dict{UInt64,Float16}
    policy_test_fails::Union{Nothing,Vector{<:Integer}}
    rec_value_errors::Dict{UInt64,Float16}
    rec_accuracies::Dict{UInt64,PolicyAcc}
    rec_policy_test_fails::Union{Nothing,Vector{<:Integer}}
    games_played::Dict{UInt64,UInt64}
    games_reanalyzed::Dict{UInt64,UInt64}
    games_avg_step::Dict{UInt64,Float16}
end

mutable struct Learner <: AbstractLearner

    uuid::UUID
    dir::String
    training_started::Union{Nothing,DateTime}
    optimizer::Flux.Optimise.AbstractOptimiser
    config::MuZeroConfig
    train::Network
    test::Network
    play::Network
    training_batch::Union{Nothing,TrainingBatch}
    replay_buffers::ReplayBuffers
    callbacks::Vector{LearnerCallback}
    tests::Vector{TrainingTests}
    is_playing::Bool
    is_training::Bool
    play_only::Bool
    stats::LearnerStats
    
    #@todo - this should be a Function
    valuefn_debug::Union{Nothing,Dict} 

    function Learner(
        config, 
        train,
        test,
        play,
        replay_buffers;
        tests=[],
        training_batch=nothing,
        optimizer=ADAMW(config.lr_init, (0.9, 0.999), config.weight_decay),
        uuid=uuid4(),
        dir="learner",
        training_started=nothing,
        losses=Dict{UInt64,Float16}(), 
        accuracies=Dict{UInt64,Float16}(),
        reward_errors=Dict{UInt64,Float16}(),
        value_errors=Dict{UInt64,Float16}(),
        times=Dict{UInt64,Float16}(),
        test_times=Dict{UInt64,Float16}(),
        policy_test_fails=nothing,
        rec_value_errors=Dict{UInt64,Float16}(),
        rec_accuracies=Dict{UInt64,Float16}(),
        rec_policy_test_fails=nothing,
        callbacks=[],
        play_only=false,
        valuefn_debug=nothing,
        games_played=Dict{UInt64,UInt64}(),
        games_reanalyzed=Dict{UInt64,UInt64}(),
        games_avg_step=Dict{UInt64,Float16}(),
        )
        uuid_object = isnothing(uuid) ? uuid4() : UUID(uuid)
        is_playing = false
        is_training = false
        stats = LearnerStats(
            Counter(0),
            losses,
            accuracies,
            reward_errors,
            value_errors,
            times,
            test_times,
            policy_test_fails,
            rec_value_errors,
            rec_accuracies,
            rec_policy_test_fails,
            games_played,
            games_reanalyzed,
            games_avg_step,
        )
        new(
            uuid_object,
            dir,
            training_started,
            optimizer,
            config, 
            train,
            test,
            play,
            training_batch,
            replay_buffers, 
            callbacks, 
            tests,
            is_playing,
            is_training,
            play_only,
            stats,
            valuefn_debug,
        )
    end
end

function Learner(
    game_type::Type{<:AbstractEnv},
    config::MuZeroConfig;
    kwargs...
    )
    Learner(;
        game_type=game_type,
        config=config,
        kwargs...
    )
end

function Learner(
    game_type::Type{<:AbstractEnv};
    kwargs...
    )
    Learner(;
        game_type=game_type,
        config=getconfig(game_type),
        kwargs...
    )
end

function (x::LearnerCallback)(l::AbstractLearner)
    #println("\nLearnerCallback ($(get_steps(l))): $(x.name), $(x.event), $(x.period), $(x.warmup)")
    if get_steps(l) % x.period == 0 && get_steps(l) >= x.warmup
        #println("$(x.name) running...\n")
        x.fn(l)
    end
end

"Launch a mock training to find a good learning rate and return suggestions based on `suggest_funcs` as a named tuple"
# function lrfind(;
#     start_lr=1e-07, 
#     end_lr=10, 
#     num_it=100, 
#     stop_div=true, 
#     show_plot=true, 
#     #suggest_funcs=valley
#     )
#     n_epoch = num_it ÷ length(self.dls.train) + 1
#     cb=LRFinder(start_lr=start_lr, end_lr=end_lr, num_it=num_it, stop_div=stop_div)
#     with self.no_logging(): self.fit(n_epoch, cbs=cb)
#     if suggest_funcs is not None:
#         lrs, losses = tensor(self.recorder.lrs[num_it//10:-5]), tensor(self.recorder.losses[num_it//10:-5])
#         nan_idxs = torch.nonzero(torch.isnan(losses.view(-1)))
#         if len(nan_idxs) > 0:
#             drop_idx = min(nan_idxs)
#             lrs = lrs[:drop_idx]
#             losses = losses[:drop_idx]
#         _suggestions, nms = [], []
#         for func in tuplify(suggest_funcs):
#             nms.append(func.__name__ if not isinstance(func, partial) else func.func.__name__) # deal with partials
#             _suggestions.append(func(lrs, losses, num_it))

#         SuggestedLRs = collections.namedtuple('SuggestedLRs', nms)
#         lrs, pnts = [], []
#         for lr, pnt in _suggestions:
#             lrs.append(lr)
#             pnts.append(pnt)
#         if show_plot: self.recorder.plot_lr_find(suggestions=pnts, nms=nms)
#         return SuggestedLRs(*lrs)

#     elif show_plot: self.recorder.plot_lr_find()
# end

function needsreanalyze(l::AbstractLearner, stats)
    stats_path = "$(stats["buffer_dir"])/stats.toml"
    d = TOML.parsefile(stats_path)
    since_last_save = Dates.now() - DateTime(d["saved_at"])
    isabandoned = Dates.value(since_last_save) > 300_000
    if !d["at_capacity"] && !isabandoned
        return false
    end
    curr_step = l.play.step
    buffer_step_min = d["step_min"]
    locked_at = valornothing(d, "locked_at")
    step_diff = curr_step - buffer_step_min
    if step_diff < 1000 
        return false
    end
    if isnothing(locked_at)
        return true
    end
    locked_at_diff = Dates.now() - DateTime(locked_at)
    Dates.value(locked_at_diff) > 600_000
end

function getbufferstats(buffer_dir)
    d = TOML.parsefile(
        "$buffer_dir/stats.toml"
    )
    d["buffer_dir"] = buffer_dir
    d
end

function findbuffertoreanalyze(l::AbstractLearner)
    dir = getreplaybuffersdir(l)
    ensuredirexists(dir)
    all_buffer_stats = getbufferstats.(
        readdirexpanduser(dir)
    )
    if isempty(all_buffer_stats)
        return nothing
    end
    d = loadplaysummary(l)
    if isnothing(d)
        return nothing
    end
    avg_step = d["avg_step"]
    avg_and_below = filter(
        x->x["step_mean"] <= avg_step,
        all_buffer_stats
    )
    if isempty(avg_and_below)
        return nothing
    end
    rand(avg_and_below)["buffer_dir"]
end

function reanalyzebuffer(
    l::AbstractLearner, 
    buffer_path::String, 
    dir=getdir(l),
    )
    buffer = deserializeexpanduser(buffer_path)
    stats_path = getstatspath(
        buffer,
        getreplaybuffersdir(l, dir)
    )
    savereplaybufferstats(
        l, 
        buffer, 
        stats_path; 
        locked_at=string(Dates.now())
    )
    reanalyze!(l, buffer)
    savereplaybuffer!(
        l,
        buffer,
        dir
    )
    println("released lock. ($(buffer.uuid))")
end


function play!(l::AbstractLearner; sleep_period=30)
    l.is_playing = true
    num_actors = getnumactors(l)
    println("play! (num_actors:$num_actors)...")
    responses = []
    if num_actors > 0 
        for _ in 1:num_actors
            res = @spawn selfplay!(l)
            push!(responses, res)
        end
    end
    buffer = getwritablereplaybuffer(l)
    total_at_last_save = length(buffer)
    while l.is_playing
        loadplayifnewer!(l)
        if !l.play_only
            loadreplaybuffers!(l)
            saveplaysummary(l)
        end
        sleep(sleep_period)
    end
    responses
end

function loadplaysummary(l::AbstractLearner, dir=getdir(l))
    summary_path = "$dir/play-summary.toml"
    if !isfileexpanduser(summary_path)
        return nothing
    end
    TOML.parsefile(expanduser(summary_path))
end

function saveplaysummary(l::AbstractLearner, dir=getdir(l))
    summary_path = "$dir/play-summary.toml"
    ensure_pathto_exists(summary_path)
    open(expanduser(summary_path), "w") do io
        TOML.print(io, getplaysummary(l))
    end
end

getrrvptests(l::AbstractLearner) = gettests(l, RecurrentRVPTests)
randrrvptest(l::AbstractLearner) = getrrvptests(l) |> rand

getpolicyvaluetests(l::AbstractLearner) = gettests(l, PolicyValueTests)
getpolicytests(l::AbstractLearner) = gettests(l, PolicyTests)

function gettests(l::AbstractLearner, T::Type{<:TrainingTests})
    idx = findfirst(ts->typeof(ts) == T, l.tests)
    isnothing(idx) ? [] : l.tests[idx]
end

#@todo randtest(l::AbstractLearner)

randpolicyvaluetest(l::AbstractLearner) = getpolicyvaluetests(l) |> rand
randpolicytest(l::AbstractLearner) = getpolicytests(l) |> rand
get_steps(l::AbstractLearner) = l.stats.step_count.n
randgame(x::AbstractLearner) = rand(x.replay_buffers)
randtrainingsample(x::AbstractLearner, g::MuZeroGame) = TrainingSample(x.config, g)
randtrainingsample(x::AbstractLearner) = randtrainingsample(x, randgame(x))
getnumactors(x::AbstractLearner) = x.config.num_actors
callbacks(x::AbstractLearner) = x.callbacks
training_network(x::AbstractLearner) = x.train
test_network(x::AbstractLearner) = x.test
play_network(x::AbstractLearner) = x.play
learner_bson_path(l) = learner_bson_path(l, getdir(l))
learner_bson_path(l, dir) = "$dir/$(l.uuid)/learner.bson"
getmodelpath(l, dir=getdir(l)) = "$dir/model.bson"
getmodeloptimizerpath(l, dir=getdir(l)) = "$dir/model-optimizer.bson"
learner_images_dir(l) = "$(getdir(l))/$(l.uuid)/images"
getstatspath(l, dir=getdir(l)) = "$dir/stats.bin"
getreplaybuffersdir(l, dir=getdir(l)) = "$dir/replay_buffers"

function getwritablereplaybufferpath(l, dir=getdir(l))
    getpath(
        getwritablereplaybuffer(l),
        getreplaybuffersdir(l, dir)
    )
end

function getwritablereplaybufferstatspath(l, dir=getdir(l))
    getstatspath(
        getwritablereplaybuffer(l),
        getreplaybuffersdir(l, dir)
    )
end

function Learner(
    config::MuZeroConfig;
    kwargs...
    )
    Learner(;
        game_type=config.game_type,
        config=config,
        kwargs...
    )
end

function Learner(;
    game_type,
    config,
    uuid=nothing, 
    dir="./learner",
    tests=[],
    cb=[],
    schedule=nothing,
    play_only=false,
    valuefn_debug=nothing,
    )
    train = new_training_network(config)
    test = new_test_network(config)
    play = new_play_network(config)
    replay_buffers = ReplayBuffers(ReplayBuffer(config))
    #plot_losses_cb = get_plot_losses_cb()
    callbacks = []
    append!(callbacks, cb)

    if !isempty(tests) 
        tests_cb = get_tests_cb()
        push!(callbacks,
            LearnerCallback(
                "InitialRunTests", 
                tests_cb,
                event=:beforetraining
            )
        )
        push!(callbacks,
            LearnerCallback(
                "UpdateTestNetwork", 
                updatetestnet!,
                period=config.checkpoint_interval,
                event=:ontestnetupdate,
            )
        )
        push!(callbacks,
            LearnerCallback(
                "RunTests", 
                tests_cb,
                event=:aftertestnetupdate,
            )
        )
    end
    push!(callbacks, 
        LearnerCallback(
            "InitialLoss", 
            saveloss!,
            event=:beforetraining
        )
    )
    push!(callbacks, 
        LearnerCallback(
            "SaveLoss", 
            saveloss!,
            event=:afterweightupdate
        )
    )
    push!(callbacks, 
        LearnerCallback(
            "SaveLearner", 
            save!,
            period=config.save_interval,
            event=:afterstep
        )
    )
    push!(callbacks,
        LearnerCallback(
            "SaveGamePlayStats", 
            savegameplaystats!,
            period=1000,
            event=:afterstep
        )
    )

    #accuracy_cb,
    #save_losses_cb,
    #save_accuracy_cb,
    #plot_losses_cb, 
    #plot_accuracy_cb, #plot_accuracy_cb = get_plot_accuracy_cb()

    l = Learner(
        config, 
        train, 
        test,
        play,
        replay_buffers; 
        uuid=uuid,
        dir=dir,
        callbacks=callbacks,
        tests=tests,
        play_only=play_only,
        valuefn_debug=valuefn_debug,
    )

    if !isnothing(schedule)
        updateschedule!(l, schedule)
    end

    if !isnothing(uuid)
        if !l.play_only
            loadtrain!(l)
            updatetestnet!(l)
        end
        
        # @todo - Probably need to load these 
        # two only if playing
        loadreplaybuffers!(l)
        loadplay!(l)
        ###

        loadstats!(l)
    else
        l.test = UniformPolicy(l.test)
        l.play = UniformPolicy(l.play)
    end
    l
end

updatetestnet!(l::AbstractLearner) = copy!(l.train, l.test)

mean10(df::DataFrame) = meanx(df, 10) 
mean100(df::DataFrame) = meanx(df, 100)
mean1000(df::DataFrame) = meanx(df, 1000)
mean10000(df::DataFrame) = meanx(df, 10000)
function meanallowmissing(x)
    skipped = skipmissing(x)
    if isempty(collect(skipped))
        return missing
    end
    mean(skipped)
end

function meanx(df::DataFrame, n::Integer)
    df[!, :stepgroup] = df.Step .÷ n
    gdf = groupby(df[2:(end-1), :], :stepgroup)
    colnames = filter(name->!in(name,["stepgroup","Step"]), names(df))
    avgcolnames = map(x->"$(x)Avg", colnames)
    combine(gdf, colnames .=> meanallowmissing .=> avgcolnames)
end

function df_cols(l::AbstractLearner)
    cols = [
        :Step   => (l,step)->step,
        :Loss   => (l,step)->valormissing(l.stats.losses, gettotal, step),
    ]
    if !isnothing(l.stats.accuracies)
        push!(cols, :Acc=>(l,step)->valormissing(l.stats.accuracies, getpct, step))
    end
end

function gettraindf(l::AbstractLearner)
    cols = df_cols(l)
    DataFrame(
        map(step->begin
            colnames = Tuple(map(first, cols))
            colvalues = map(col->col[2](l,step), cols)
            NamedTuple{colnames}(colvalues)
        end, 0:get_steps(l))
    )
end

function showresultsfull(l::AbstractLearner; mean=nothing, max_items=20)
    df = getresultsfull(l, mean=mean)
    len = size(df, 1)
    from = max(1, len-max_items+1)
    to = len
    pretty_table(
       df[from:to, :], 
       show_row_number=false, 
       nosubheader=true, 
       title="Training Results (All columns)",
   )
end

function showresults(l::AbstractLearner; mean=nothing)
    pretty_table(
       getresults(l, mean=mean); 
       show_row_number=false, 
       nosubheader=true, 
       title="Training Results",
   )
end

function getresultsfull(l::AbstractLearner; mean=nothing)
    if isnothing(mean)
        len = max(length(l.stats.losses), 1)
        power = Int(floor(log10(len)))
        mean = 10^power
    end
    meanx(
        gettraindfall(l),
        mean
    )
end

function getresults(l::AbstractLearner; mean=nothing)
    if isnothing(mean)
        power = Int(floor(log10(length(l.stats.losses))))
        mean = 10^power
    end
    meanx(
        select(gettraindfall(l), [:Step, :Loss, :PolicyAcc]),
        mean
    )
end

function showlosses(l::AbstractLearner)
    pretty_table(
        getlosses(l); 
        show_row_number=false, 
        nosubheader=true, 
        title="Losses",
    )
end

function getlosses(l::AbstractLearner)
    select(gettraindfall(l), [:Step, :Loss, :PolicyLoss, :RewardLoss, :ValueLoss])
end

getloss(l::AbstractLearner, step) = valormissing(l.stats.losses, gettotal, step)

function getlosses(l::AbstractLearner)
    map(
        step->getloss(l, step),
        1:get_steps(l)
    )
end
    
function df_cols_full(l::AbstractLearner)
    cols = [
        :Step       => (l,step)->step,
        :Loss       => (l,step)->valormissing(l.stats.losses, gettotal, step),
        :RewardLoss => (l,step)->valormissing(l.stats.losses, getreward, step),
        :ValueLoss  => (l,step)->valormissing(l.stats.losses, getvalue, step),
        :PolicyLoss => (l,step)->valormissing(l.stats.losses, getpolicy, step),
        :TestTime   => (l,step)->valormissing(l.stats.test_times, step),
    ]
    if !isnothing(l.stats.value_errors)
        push!(cols, :ValueErr=>(l,step)->valormissing(l.stats.value_errors, step))
    end
    if !isnothing(l.stats.accuracies)
        push!(cols, :PolicyPass=>(l,step)->valormissing(l.stats.accuracies, getpass, step))
        push!(cols, :PolicyFail=>(l,step)->valormissing(l.stats.accuracies, getfail, step))
        push!(cols, :PolicyAcc=>(l,step)->valormissing(l.stats.accuracies, getpct, step))
    end
    if !isnothing(l.stats.reward_errors)
        push!(cols, :RewardErr=>(l,step)->valormissing(l.stats.reward_errors, step))
    end
    if !isnothing(l.stats.rec_value_errors)
        push!(cols, :RecValueErr=>(l,step)->valormissing(l.stats.rec_value_errors, step))
    end
    if !isnothing(l.stats.rec_accuracies)
        push!(cols, :RecPolicyPass=>(l,step)->valormissing(l.stats.rec_accuracies, getpass, step))
        push!(cols, :RecPolicyFail=>(l,step)->valormissing(l.stats.rec_accuracies, getfail, step))
        push!(cols, :RecPolicyAcc=>(l,step)->valormissing(l.stats.rec_accuracies, getpct, step))
    end
    push!(cols, :Played=>(l,step)->valormissing(l.stats.games_played, step))
    push!(cols, :Reanalyzed=>(l,step)->valormissing(l.stats.games_reanalyzed, step))
    push!(cols, :GameStep=>(l,step)->valormissing(l.stats.games_avg_step, step))
    #push!(cols, :GameStepQuality=>(l,step)->getgamestepquality(l.replay_buffers, step))
    cols
end

function gettraindfall(l::AbstractLearner)
    cols = df_cols_full(l)
    DataFrame(
        map(step->begin
            colnames = Tuple(map(first, cols))
            colvalues = map(col->col[2](l,step), cols)
            NamedTuple{colnames}(colvalues)
        end, 0:get_steps(l))
    )
end

function appendtraincsv(l::AbstractLearner)
    # Step,TrainLoss,PolicyAcc,ValueErr,RewardErr,PolicyLoss,ValueLoss,RewardLoss,Time,TestTime
    # df = DataFrame(Step=[], PolicyAcc=[])
    # losses = Dict()
    # for (step, loss) in enumerate(l.stats.losses)
    #     losses[step-1] = loss
    # end
    # losses = Dict()
    # for acc in l.stats.accuracies
    #     steps[acc.step] = acc.record
    # end
    # for step in 0:length(keys(steps))
    #     push!(df.Step, step)
    #     push!(df.PolicyAcc, x.record)
    # end
    # path = "$(getdir(l))/$(l.uuid)/csv/train.csv"
    # ensure_pathto_exists(path)
    # @show path, df
    # CSV.write(path, df)
end

function persistent_field_names()
    [
        :training_started, 
        :reward_errors,
    ]
end

function getcheckpointdir(l::AbstractLearner, dir=getdir(l))
    "$dir/checkpoints/$(getcheckpointstr(l.train.saved_at))"
end

function ismodelimproved(l::AbstractLearner)
    # @todo
    true
end

function save!(l::AbstractLearner, dir=getdir(l); force=false)
    if force || ismodelimproved(l) 
        savemodel(l, dir)
        #savemodeloptimizer(l, dir)
        savestats(l, dir)
        savecheckpoint(l, dir)
    end
end

function savestats(l::AbstractLearner, dir=getdir(l))
    path = getstatspath(l, dir)
    ensure_pathto_exists(path)
    serializeexpanduser(path, l.stats)
end

function loadstats!(l::AbstractLearner, path=getstatspath(l))
    if isfileexpanduser(path)
        l.stats = deserializeexpanduser(path)
    end
end

function loadtrain!(l::AbstractLearner, path=getmodelpath(l))
    if isfileexpanduser(path)
        try
            load!(l.train, path)
        catch e
            println("Training model could not be loaded.")
        end
    end
end

function getlatestsummaryondisk(l::AbstractLearner, dir=getdir(l))
    summary_path = "$dir/summary.toml"
    if !isfileexpanduser(summary_path)
        return nothing
    end
    TOML.parsefile(expanduser(summary_path))
end

function loadplayifnewer!(l::AbstractLearner, dir=getdir(l))
    function isnewer()
        d = getlatestsummaryondisk(l, dir)
        if isnothing(d)
            return false
        end
        saved_at = d["saved_at"]
        if isnothing(saved_at)
            return false
        end
        if isnothing(l.play)
            return true
        end
        isnothing(l.play.saved_at) || saved_at > l.play.saved_at
    end
    if isnewer()
        model_path = getmodelpath(l, dir)
        loadplay!(l, model_path)
    end
end

function loadplay!(l::AbstractLearner, path=getmodelpath(l))
    if isfileexpanduser(path)
        try
            load!(l.play, path)
            if l.play_only
                println("Loaded new play network (step: $(l.play.step))")
            end
        catch e
            @error "Play model could not be loaded." exception=(e, catch_backtrace())
        end
    end
end

# function loadmodeloptimizer!(l::AbstractLearner, path=getmodeloptimizerpath(l))
#     if isfileexpanduser(path)
#         try
#             d = BSON.load(expanduser(path), @__MODULE__)
#             load!(l.train, d)
#             @show d[:opt]
#             #l.optimizer = d[:optimizer]
#         catch e
#             println("Optimizer/model could not be loaded. Try loading only the the model.")
#         end
#     end
# end


function savewritablereplaybuffer(l::AbstractLearner, dir=getdir(l))
    length_ondisk = 0
    stats_path = getwritablereplaybufferstatspath(l, dir)
    if isfileexpanduser(stats_path)
        length_ondisk = TOML.parsefile(expanduser(stats_path))["length"]
    end
    writable = getwritablereplaybuffer(l)
    if length(writable) > length_ondisk
        println("buffer length: $(length(writable)) - $(atcapacity(writable, l.config) ? "" : "not") at capacity (saving...) \n")
        savereplaybuffer!(
            l,
            writable,
            dir
        )
    end
end

function newwriteable!(l::AbstractLearner)
    new_writable = ReplayBuffer(l.config)
    setbuffer!(l.replay_buffers, new_writable)
    l.replay_buffers.writable_uuid = string(new_writable.uuid)
end

function savereplaybuffer!(
    l::AbstractLearner, 
    buffer::ReplayBuffer,
    dir=getdir(l)
    )
    path = getpath(
        buffer,
        getreplaybuffersdir(l, dir)
    )
    stats_path = getstatspath(
        buffer,
        getreplaybuffersdir(l, dir)
    )
    buffer.saved_at = string(Dates.now())
    updatenetstep!(buffer)
    ensure_pathto_exists(path)
    serializeexpanduser(path, buffer)
    ensure_pathto_exists(stats_path)
    savereplaybufferstats(
        l, 
        buffer, 
        stats_path;
        locked_at=nothing
    )
end

function savereplaybufferstats(
    l::AbstractLearner,
    buffer::ReplayBuffer, 
    stats_path::String;
    locked_at=nothing,
    )
    replay_data = Dict(
        :uuid => string(buffer.uuid),
        :at_capacity => atcapacity(buffer, l.config),
        :length => length(buffer),
        :play_threads => l.config.num_actors,
    )
    fields = [
        :play_start,
        :saved_at,
        :play_finish,
        :reanalyze_start,
        :reanalyze_finish,
        :step_min,
        :reanalyzed,
        :step_max,
        :step_mean,
        :play_duration,
        :reanalyze_duration,
        :single_play_duration,
        :single_reanalyze_duration,
    ]
    for f in fields
        v = getfield(buffer, f)
        if !isnothing(v)
            replay_data[f] = v
        end
    end
    if !isnothing(locked_at)
        replay_data[:locked_at] = locked_at
        println("lock acquired ($(buffer.uuid)) ...")
    end
    ensure_pathto_exists(expanduser(stats_path))
    open(expanduser(stats_path), "w") do io
        TOML.print(io, replay_data)
    end
end

function loadreplaybuffers!(l::AbstractLearner)
    if l.play_only
        return
    end
    writable_path = getwritablereplaybufferpath(l)
    replay_buffers_dir = getreplaybuffersdir(l)
    if isdirexpanduser(replay_buffers_dir)
        readonly_buffer_dirs = filter(
            x->!occursin(writable_path, x), 
            readdirexpanduser(replay_buffers_dir)
        )
        #@show readonly_buffer_dirs
        function isnewbuffer(buffer_dir)
            stats_path = "$buffer_dir/stats.toml"
            buffer_uuid = splitpath(stats_path)[end-1]
            if !hasbuffer(l.replay_buffers, buffer_uuid)
                println("new buffer found ($buffer_uuid)")
                return true
            end
            d = TOML.parsefile(stats_path)
            if !haskey(d, "saved_at")
                @show stats_path
                println("no d.saved_at!")
                return false
            end
            saved_at = d["saved_at"]                
            buff = getbuffer(l.replay_buffers, buffer_uuid)
            if saved_at > buff.saved_at
                println("refreshing buffer $buffer_uuid (new games found)...")
                return true
            end
            false
        end
        new_buffer_dirs = filter(
            x->isnewbuffer(x),
            readonly_buffer_dirs
        )
        for new_buffer_dir in new_buffer_dirs
            buffer_path = "$new_buffer_dir/buffer.bin"
            setbuffer!(l.replay_buffers, buffer_path)
        end
    end
    nothing
end

function savecheckpoint(l, dir)
    saveconfig(l, dir)
    savesummary(l, dir)
    saveresults(l, dir; mean=1000)
    saveresults(l, dir; mean=10000)
end

function saveresults(l::AbstractLearner, dir=getdir(l); mean=1000)
    checkpoint_dir = getcheckpointdir(l, dir)
    results_path = "$checkpoint_dir/results-m$mean.txt"
    ensure_pathto_exists(results_path)
    open(expanduser(results_path), "w") do io
        pretty_table(
            io,
            getresultsfull(l, mean=mean); 
            show_row_number=false, 
            nosubheader=true, 
            title="Training Results (All columns)",
        )
    end
end

function randtrainingresult(x::AbstractLearner)
    if isnothing(x.training_batch)
        println("No training batch defined.")
        return
    end
    sample = x.training_batch |> rand
    initial_output = initial_inference(x.test, sample.image) |> SingleNetworkOutput
    outputs = [initial_output]
    if !isempty(sample.actions)
        # @todo: make this support arbitrarily long rollouts
        recurrent_output = recurrent_inference(x.test, hidden_state(initial_output), sample.actions[1]) |> SingleNetworkOutput
        push!(outputs, recurrent_output)
    end
    TrainingSampleResult(sample, outputs)
end

function getdir(l::AbstractLearner)
    dir = rstrip(l.dir, '/')
    join([dir,string(l.uuid)], "/")
end

function get_plot_losses_cb()
    function fn(l::AbstractLearner)
        losses = l.stats.losses
        if length(losses) <= 100
            plot(losses, fmt = :png)    
        else
            plot(avgx(losses, 100), fmt = :png)
        end
        losses_path = "$(learner_images_dir(l))/losses.png"
        copy_path = "$(learner_images_dir(l))/losses-copy.png"
        ensure_pathto_exists(losses_path)
        savefig(copy_path)
        cp(copy_path, losses_path, force=true)
        rm(copy_path) 
    end
    #LearnerCallback(fn)
end

function updatepolicytestresults!(l::AbstractLearner, acc, fails)
    l.stats.accuracies[get_steps(l)] = acc
    l.stats.policy_test_fails = fails
end

function updatevaluetestresults!(l::AbstractLearner, mse)
    l.stats.value_errors[get_steps(l)] = mse
end

function updaterewardtestresults!(l::AbstractLearner, mse)
    l.stats.reward_errors[get_steps(l)] = mse
end

function updaterecvaluetestresults!(l::AbstractLearner, mse)
    l.stats.rec_value_errors[get_steps(l)] = mse
end

function updaterecpolicytestresults!(l::AbstractLearner, acc, fails)
    l.stats.rec_accuracies[get_steps(l)] = acc
    l.stats.rec_policy_test_fails = fails
end

function runtests!(l::AbstractLearner, tests::RecurrentRVPTests)
    reward_mse, value_mse, acc, fails = gettestfn(tests)(gettestnet(l))
    updaterewardtestresults!(l, reward_mse)
    updaterecvaluetestresults!(l, value_mse)
    updaterecpolicytestresults!(l, acc, fails)
    nothing
end

function runtests!(l::AbstractLearner, tests::PolicyTests)
    acc, fails = gettestfn(tests)(gettestnet(l))
    updatepolicytestresults!(l, acc, fails)
    nothing
end

function runpvtests(l::AbstractLearner)
    runtests(
        gettests(l, PolicyValueTests), 
        gettestnet(l)
    )
end
function runrrvptests(l::AbstractLearner)
    runtests(
        gettests(l, RecurrentRVPTests), 
        gettestnet(l)
    )
end

function runtests!(l::AbstractLearner, tests::ValueTests)
    mse = gettestfn(tests)(gettestnet(l))
    updatevaluetestresults!(l, mse)
end

function runtests!(l::AbstractLearner, tests::RewardTests)
    mse = gettestfn(tests)(gettestnet(l))
    updaterewardtestresults!(l, mse)
end

function runtests!(l::AbstractLearner, tests::PolicyValueTests)
    acc, fails, value_mse = runtests(tests, gettestnet(l))
    updatepolicytestresults!(l, acc, fails)
    updatevaluetestresults!(l, value_mse)
end

function runtests!(l::AbstractLearner, tests::PolicyRewardTests)
    prtestfn = gettestfn(tests)
    acc, fails, reward_mse = prtestfn(gettestnet(l))
    updatepolicytestresults!(l, acc, fails)
    updaterewardtestresults!(l, reward_mse)
end

function runtests!(l::AbstractLearner, tests::ValueRewardTests)
    vrtestfn = gettestfn(tests)
    value_mse, reward_mse = vrtestfn(gettestnet(l))
    updatevaluetestresults!(l, value_mse)
    updaterewardtestresults!(l, reward_mse)
end

function runtests!(l::AbstractLearner, tests::PolicyValueRewardTests)
    rvptestfn = gettestfn(tests)
    acc, fails, value_mse, reward_mse = rvptestfn(gettestnet(l))
    updatepolicytestresults!(l, acc, fails)
    updatevaluetestresults!(l, value_mse)
    updaterewardtestresults!(l, reward_mse)
end

function get_tests_cb()
    function fn(l::AbstractLearner)
        if isempty(l.tests)
            return nothing
        end
        t = @elapsed begin
            for testset in l.tests
                if get_steps(l) % testset.period == 0 && get_steps(l) >= testset.warmup
                    runtests!(l,testset)
                end
            end
        end
        l.stats.test_times[get_steps(l)] = round(t, digits=2)
    end
end

function get_plot_accuracy_cb()
    function fn(l::AbstractLearner)
        accuracies = get_record.(l.stats.accuracies)
        if length(accuracies) <= 100
            plot(accuracies, fmt = :png)    
        else
            plot(avgx(accuracies, 100), fmt = :png)
        end
        accuracy_path = "$(learner_images_dir(l))/accuracy.png"
        copy_path = "$(learner_images_dir(l))/accuracy-copy.png"
        ensure_pathto_exists(accuracy_path)
        savefig(copy_path)
        cp(copy_path, accuracy_path, force=true)
        rm(copy_path) 
    end
    #LearnerCallback(fn)
end

function clearreplaybuffer!(l::AbstractLearner)
    @assert false "punting on this for now"
    # stopplaying!(l)
    # setwritablereplaybuffer!(l, ReplayBuffer(l.config))
    # nothing
end

stopplaying!(l::AbstractLearner) = l.is_playing = false
getwritablereplaybuffer(l::AbstractLearner) = getwritablebuffer(l.replay_buffers)

function selfplay!(l::AbstractLearner)
    println("selfplay! started on thread", threadid(), " ...")
    while l.is_playing
        game_limit_reached = false
        d = loadplaysummary(l)
        if !isnothing(d)
            if d["num_games"] >= d["max_games"]
                game_limit_reached = true
            end
        end
        if timetoreanalyze(l)
            if l.play_only
                println("reanalyzing on ", threadid(), " ...")
            end
            reanalyze!(l)
        elseif !game_limit_reached
            if l.play_only
                println("playing on ", threadid(), " ...")
            end
            savegames!(l, playgames(l))
        else
            sleep(30)
        end
    end
end

function timetoreanalyze(l::AbstractLearner)
    if isnothing(findbuffertoreanalyze(l))
        return false
    end
    d = loadplaysummary(l)
    if isnothing(d)
        return false
    end
    l.play.step > d["avg_step"] && d["step_quality"] < 1.0
end

function reanalyze!(l::AbstractLearner)
    dir = findbuffertoreanalyze(l)
    if !isnothing(dir)
        reanalyzebuffer(l, "$dir/buffer.bin")
    end
end

function atreplaycapacity(l::AbstractLearner)
    atcapacity(getwritablereplaybuffer(l), l.config)
end

function finishplay!(l::AbstractLearner)
    rbs = getwritablereplaybuffer(l)
    newwriteable!(l)
    finishplay!(rbs)
    savereplaybuffer!(l, rbs)
end

function savegames!(l::AbstractLearner, games::Vector{MuZeroGame})
    replay = getwritablereplaybuffer(l)
    replay.games = games
    # @todo - remove mzgame.net_step 
    # (should only be at buffer level)
    replay.step_min = games[1].net_step
    newwriteable!(l)
    finishplay!(replay)
    savereplaybuffer!(l, replay)
end

struct LearnerLossesPNG
    learner::AbstractLearner
end

losses_png(x::AbstractLearner) = x |> LearnerLossesPNG

function Base.show(io::IO, m::MIME"image/png", g::LearnerLossesPNG)
    print(io, read("losses.png", String))
end

randtrainingbatch(x::AbstractLearner) = randtrainingbatch(x.replay_buffers, x.config)

function new_training_batch!(x::AbstractLearner)
    set_training_batch!(x, randtrainingbatch(x))
end

training_batch(x::AbstractLearner) = x.training_batch
training_tensor(x::AbstractLearner) = x |> training_batch |> tensor

# @todo - this should be at the LearnerStats level
function push_time!(l::AbstractLearner, field_name::Symbol, t)
    getfield(l.stats, field_name)[get_steps(l)] = round(t, digits=2)
end

stoptraining!(l::AbstractLearner) = l.is_training = false

function fit!(l::AbstractLearner; steps=nothing)
    l.is_training = true
    println("fit! started on thread ", threadid(), " ...")
    updatetrainingstarted!(l)
    #learning_rate = config.lr_init * config.lr_decay_rate^(get_steps(l) / config.lr_decay_steps)
    starting_step = get_steps(l)
    maxsteps() = isnothing(steps) ? l.config.training_steps : (starting_step + steps)
    while l.is_training && get_steps(l) < maxsteps()
        new_training_batch!(l)
        if get_steps(l) == 0
            runcallbacks(l, :beforetraining)
        end
        runcallbacks(l, :beforeweightupdate)
        t = @elapsed update_weights!(l)
        push_time!(l, :times, t)
        runcallbacks(l, :afterweightupdate)
        runcallbacks(l, :ontestnetupdate)
        runcallbacks(l, :aftertestnetupdate)
        runcallbacks(l, :afterstep)
    end
    save!(l)
    nothing
end

function saveloss!(l::AbstractLearner)
    l.stats.losses[get_steps(l)] = getloss(
        l.train, 
        training_tensor(l),
    )
end

getavggamestep(l::AbstractLearner) = getavgstep(l.replay_buffers)

function savegameplaystats!(l::AbstractLearner)
    l.stats.games_played[get_steps(l)] = length(l.replay_buffers)
    l.stats.games_reanalyzed[get_steps(l)] = calcnumreanalyzed(l.replay_buffers)
    l.stats.games_avg_step[get_steps(l)] = getavggamestep(l)
end

function getcallbacks(l::AbstractLearner, event)
    filter(cb->cb.event == event, l.callbacks)
end

function runcallbacks(l::AbstractLearner, event)
    foreach(cb->cb(l), getcallbacks(l, event))
end

function update_weights!(l::AbstractLearner)
    network = l.train
    batch = training_tensor(l)
    weight_decay = l.config.weight_decay
    local loss
    ps = get_weights(network)
    gs = Zygote.gradient(ps) do
        loss = getloss(network, batch)
        gettotal(loss)
    end
    Flux.update!(l.optimizer, ps, gs)
    increment!(l.stats.step_count)
    l.train.step = l.stats.step_count.n
    loss
end

function export_csv(l::AbstractLearner)
    dir = getdir(l)
    train_path = "$dir/csv/train.csv"
    ensure_pathto_exists(train_path)
    df = DataFrame(
        "Step"=>[1],
        "Loss"=>[1],
        "PolicyAcc"=>[1],
        "ValueErr"=>[1],
        "RewardErr"=>[1],
        "PolicyLoss"=>[1],
        "ValueLoss"=>[1],
        "RewardLoss"=>[1],
        "TestTime"=>[1],
    )
    CSV.write(expanduser(train_path), df)
end

function savemodel(l::AbstractLearner, dir)
    if !l.config.skip_save
        save!(l.train, getmodelpath(l, dir))
    end
end

# function savemodeloptimizer(l::AbstractLearner, dir)
#     if !l.config.skip_save
#         save!(
#             l.train,
#             l.optimizer,
#             getmodeloptimizerpath(l, dir)
#         )
#     end
# end

updatetrainingstarted!(l::AbstractLearner) = l.training_started = now()

function getgamestepquality(l::AbstractLearner)
    d = getlatestsummaryondisk(l)
    if isnothing(d) || !haskey(d,"step_count")
        return 0
    end
    getgamestepquality(
        l.replay_buffers,  
        d["step_count"],
    )
end 

function getplaysummary(l::AbstractLearner)
    Dict(
        "num_games" => length(l.replay_buffers),
        "max_games" => l.config.max_games,
        "step_quality" => getgamestepquality(l),
        "avg_step" => getavggamestep(l),
    )
end

function get_summary(l::AbstractLearner)
    d = Dict(
        "uuid" => string(l.uuid),
        "training_started" => string(l.training_started),
        "saved_at" => l.train.saved_at,
        "step_count" => get_steps(l),
        "avg_game_step" => getavggamestep(l),
    )
    loss = get_loss(l)
    if !isnothing(loss)
        d["loss"] = round(loss, digits=6)
    end
    acc = get_accuracy(l)
    if !isnothing(acc)
        d["accuracy"] = round(getpct(acc), digits=6)
    end
    d
end

get_reward_error(l::AbstractLearner) = lastbykeyornothing(l.stats.reward_errors)
get_accuracy(l::AbstractLearner) = lastbykeyornothing(l.stats.accuracies)
get_loss(l::AbstractLearner) = lastbykeyornothing(l.stats.losses, gettotal)
function getconfigpath(l::AbstractLearner, dir=getdir(l))
     "$(getcheckpointdir(l, dir))/config.txt"
end

function saveconfig(l::AbstractLearner, dir=getdir(l))
    save(l.config, getconfigpath(l, dir))
end

#loadconfig!(l::AbstractLearner) = l.config = load(MuZeroConfig, getconfigpath(l))

function savesummary(l::AbstractLearner, dir)
    summary_path = "$dir/summary.toml"
    ensure_pathto_exists(summary_path)
    open(expanduser(summary_path), "w") do io
        TOML.print(io, get_summary(l))
    end
end

function policy_test_result(x::AbstractLearner)
    sample = x.training_batch |> rand
    initial_output = initial_inference(x.test, sample.image) |> SingleNetworkOutput
    recurrent_output = recurrent_inference(x.test, hidden_state(initial_output), sample.actions[1]) |> SingleNetworkOutput
    PolicyTestResult(sample, [initial_output, recurrent_output])
end

randpolicytestresult(x::AbstractLearner) = x |> getpolicytestresults |> rand
showpolicytestresults(x::AbstractLearner) = display(getpolicytestresults(x))
showpolicyvaluetestresults(x::AbstractLearner) = display(getpolicyvaluetestresults(x))
showrrvptestresults(x::AbstractLearner) = display(getrrvptestresults(x))

function getpvtestresult(l::AbstractLearner, idx)
    getpolicyvaluetestresults(l)[idx]
end

getpolicyvaluetestresult(l::AbstractLearner, idx) = getpolicyvaluetestresults(l)[idx]

function getpolicyvaluetestresults(x::AbstractLearner)
    tests = getpolicyvaluetests(x)
    states = [t.state for t in tests]
    images = tensorize(map(get_image, states))
    output = initial_inference(x.test, images) |> NetworkOutput
    policies = collect(eachcol(policy(output)))
    results = PolicyValueTestResult.(tests, policies)
    PolicyValueTestResults(results)
end

function getrrvptestresults(x::AbstractLearner)
    tests = getrrvptests(x)
    output = recurrentinferencefn(tests)(
        x.test
    )
    RecurrentRVPTestResults(
        RecurrentRVPTestResult.(
            tests, 
            first.(eachcol(reward(output))),
            first.(eachcol(value(output))),
            collect(eachcol(policy(output))),
        )
    )
end

function getpolicyvaluetestresults(x::AbstractLearner)
    tests = getpolicyvaluetests(x)
    states = [t.state for t in tests]
    images = tensorize(map(get_image, states))
    output = initial_inference(x.test, images) |> NetworkOutput
    policies = collect(eachcol(policy(output)))
    results = PolicyValueTestResult.(tests, policies)
    PolicyValueTestResults(results)
end

function getpolicy(l::AbstractLearner, s::AbstractEnvState)
    initial_inference(l.test, s) |> SingleNetworkOutput |> policy
end

function getvalue(l::AbstractLearner, s::AbstractEnvState)
    initial_inference(l.test, s) |> SingleNetworkOutput |> value
end

function getvalue(l::AbstractLearner, s::AbstractEnvState, a::Integer)
    hs = initial_inference(l.test, s) |> SingleNetworkOutput |> hidden_state
    recurrent_inference(l.test, hs, a) |> SingleNetworkOutput |> value
end

function findmuzerogames(l::AbstractLearner, target_state::AbstractEnvState)
    mzgames = []
    for (idx,mzgame) in enumerate(l.replay_buffers)
        for s in mzgame.game.states
            if s == target_state
                println("found... $(idx)")
                push!(mzgames, mzgame)
                break
            end
        end        
    end 
    mzgames
end

function getpolicytestresults(x::AbstractLearner)
    tests = getpolicytests(x)
    states = [t.state for t in tests]
    images = tensorize(map(get_image, states))
    output = initial_inference(x.test, images) |> NetworkOutput
    policies = collect(eachcol(policy(output)))
    results = PolicyTestResult.(tests, policies)
    PolicyTestResults(results)
end

function cycleschedule(; activatation_step=0, min_lr=1e-6, max_lr=1e-2, cycle_steps=10_000)
    (l::AbstractLearner)->begin
        new_lr = trianglelr(
            get_steps(l);
            activatation_step=activatation_step,
            min_lr=min_lr,
            max_lr=max_lr,
            cycle_steps=cycle_steps,
        )
        #println("Change learning rate to: $new_lr")
        l.optimizer[3].eta = new_lr
    end
end

function newcycleschedule(
    l::AbstractLearner; 
    activatation_step=get_steps(l), 
    min_lr=1e-6, 
    max_lr=1e-2, 
    cycle_steps=10_000,
    )
    cycleschedule(; 
        activatation_step=activatation_step,
        min_lr=min_lr,
        max_lr=max_lr,
        cycle_steps=cycle_steps,
    )
end

function updateschedule!(l::AbstractLearner, schedule)
    updatecallback!(
        l,
        LearnerCallback(
            "UpdateLearningRate", 
            schedule,
            event=:beforeweightupdate
        )
    )
end

function updatecallback!(l::AbstractLearner, cb::LearnerCallback)
    idx = findfirst(x->x.name == cb.name, l.callbacks)
    if !isnothing(idx)
        deleteat!(l.callbacks, idx)
    end
    push!(l.callbacks, cb)
end

function trianglelr(step; activatation_step, min_lr=1e-6, max_lr=1e-2, cycle_steps=10_000)
    if step < activatation_step
        return min_lr
    end
    calibrated_step = step - activatation_step
    curr_step = mod1(calibrated_step, cycle_steps)
    halfcycle_steps = cycle_steps ÷ 2
    delta = (max_lr - min_lr) / halfcycle_steps
    if curr_step <= halfcycle_steps
        return min_lr + curr_step*delta
    end
    min_lr + (cycle_steps - curr_step)*delta
end

set_training_batch!(l::AbstractLearner, b::TrainingBatch) = l.training_batch = b
show_game(l::AbstractLearner) = randgame(l) |> display
show_tests(l::AbstractLearner) = show_policy_tests(l)

function show_policy_tests(x::AbstractLearner, n=10)
    indices = 1:n
    len = length(x.policy_tests)
    title_paren = "all $len"
    if len > n
        title_paren = "random batch of $n"
        indices = randperm(length(x.policy_tests))[1:n]
    end
    d = Dict(
        "Test#"=>[],
        "Action(s)"=>[],
    )
    for idx in indices
        t = x.policy_tests[idx]
        push!(d["Test#"], idx)
        push!(d["Action(s)"], join(t.actions, ","))
    end
    df = DataFrame(
        "Test Key"=>d["Test#"],
        "Action(s)"=>d["Action(s)"],
    )
    println()
    println("Policy Tests ($title_paren)\n")
    show(df, show_row_number=false, nosubheader=true, title="")
    println()
end

show_pv_tests(x::AbstractLearner, n=10) = show_policyvalue_tests(x, n)

function show_policyvalue_tests(x::AbstractLearner, n=10)
    indices = 1:n
    len = length(x.policyvalue_tests)
    title_paren = "all $len"
    if len > n
        title_paren = "random batch of $n"
        indices = randperm(length(x.policyvalue_tests))[1:n]
    end
    d = Dict(
        "Test Key"=>[],
        "Value"=>[],
        "Action(s)"=>[],
    )
    for idx in indices
        t = x.policyvalue_tests[idx]
        push!(d["Test Key"], idx)
        push!(d["Value"], t.value)
        push!(d["Action(s)"], join(t.actions, ","))
    end
    df = DataFrame(
        "Test Key"=>d["Test Key"],
        "Value Test"=>d["Value"],
        "Action Test"=>d["Action(s)"],
    )
    println()
    println("Policy-Value Tests ($title_paren)\n")
    show(df, show_row_number=false, nosubheader=true, title="")
    println()
end

function show_reward_tests(x::AbstractLearner, n=10)
    indices = 1:n
    len = length(x.reward_tests)
    title_paren = "all $len"
    if len > n
        title_paren = "random batch of $n"
        indices = randperm(length(x.reward_tests))[1:n]
    end
    d = Dict(
        "Idx"=>[],
        "Action"=>[],
        "Reward Test"=>[],
    )
    for idx in indices
        t = x.reward_tests[idx]
        push!(d["Idx"], idx)
        push!(d["Action"], t.action)
        push!(d["Reward Test"], t.reward)
    end
    df = DataFrame(
        "Test Key"=>d["Idx"],
        "Action"=>d["Action"],
        "Reward Test"=>d["Reward Test"],
    )
    println()
    println("Reward Tests ($title_paren)\n")
    show(df, show_row_number=false, nosubheader=true, title="")
    println()
end

gettestnet(l::AbstractLearner) = l.test
getplaynet(l::AbstractLearner) = l.play

function play_game(l::AbstractLearner, action_idxs::Vector{<:Integer})
    action_list = getactions(l.config.game_type)
    actions = map(idx->action_list[idx], action_idxs)
    play_game(l, plan=actions)
end

function newgames(l::AbstractLearner, i)
    step = getplaynet(l).step
    [new_game(l.config, step) for _ in 1:i]
end

movesmaxedout(config::MuZeroConfig, game) = movesmaxedout(game, config.max_moves)
isactive(config::MuZeroConfig, game) = !is_terminal(game) && !movesmaxedout(config, game)

function playgames(l::AbstractLearner; n=100, plans=[])
    config = l.config
    net = getplaynet(l)
    games = newgames(l, n)
    active_games = view(games, :)
    active_plans = view(plans, :)
    while !isempty(active_games)
        roots = Node.(0, get_state.(active_games))
        net_output = initial_inference(net, active_games) |> NetworkOutput
        expand_node!(roots, legal_actions.(active_games), net_output, zeros(Float32, length(roots)))
        add_exploration_noise(config, roots)
        run_mcts(net, config, roots, action_history.(active_games))
        store_search_statistics.(active_games, roots)
        next_actions = getnextactions(
            roots,
            active_plans;
            num_moves=length.(history.(active_games)),
            num_steps=net.step,
            temp_fn=config.visit_softmax_temperature_fn
        )  
        apply!.(active_games, next_actions)
        active_idxs = isactive.(config, active_games)
        active_games = view(active_games, active_idxs)
        if !isempty(active_plans)
            active_plans = view(active_plans, active_idxs)
        end
    end
    games
end

function play_game(l::AbstractLearner; show_root=false, plan=[])
    net = getplaynet(l)
    game = new_game(l.config, net.step)
    while !is_terminal(game) && !movesmaxedout(l.config, game)
        root = Node(0, get_state(game))
        net_output = initial_inference(net, game) |> SingleNetworkOutput
        expand_node!(root, legal_actions(game), net_output, Float32(0))
        add_exploration_noise(l.config, root)
        run_mcts(net, l.config, root, action_history(game))
        store_search_statistics(game, root)
        action = select_action(
            root;
            num_moves=length(history(game)),
            num_steps=net.step,
            temp_fn=l.config.visit_softmax_temperature_fn,
            plan=plan
        )  
        apply!(game, action)
    end
    game
end

function reanalyze!(l::AbstractLearner, b::ReplayBuffer)
    try
        b.reanalyze_start = string(now())
        b.games = playgames(
            l;
            n=length(b), 
            plans=gethistory.(getgame.(b))
        )
        b.reanalyze_finish = string(now())
        b.reanalyze_duration = Dates.value(
            now() - DateTime(b.reanalyze_start)
        )
        b.single_reanalyze_duration = Int32(
            round(b.reanalyze_duration / length(b))
        )
        incrementreanalyzed!(b)
        updatenetstep!(b)
    catch e
        @error "reanalyze errored out." exception=(e, catch_backtrace())
    end
end

function reanalyze(l::AbstractLearner, mzg::MuZeroGame)
    play_game(l, plan=gethistory(mzg.game)) 
end

function run_mcts(
    net::Network,
    config::MuZeroConfig, 
    root::Node, 
    action_history::ActionHistory;
    valuefn_debug=nothing
    )
    min_max_stats = MinMaxStats(config.known_bounds)
    logging = false #length(action_history) == 3
    for sim_idx in 1:config.num_simulations
        if logging
            println()
            @show sim_idx
        end
        history = clone(action_history)
        node = root
        search_path = [node]
        while expanded(node)
            action, node = select_child!(
                node, 
                min_max_stats,
                config,
            )
            add_action!(history, action)
            push!(search_path, node)
        end
        parent = search_path[end-1]
        inference = recurrent_inference(net, parent.hidden_state, last_action(history)) |> SingleNetworkOutput
        r = reward(inference)
        expand_node!(node, legal_actions(config, history), inference, r)
        v = value(inference)
        if !isnothing(valuefn_debug)
            v = valuefn_debug[getnextstate(parent.state, last_action(history))]
        end
        backpropagate!(
            search_path, 
            v, 
            to_play(history),
            config.discount, 
            min_max_stats
        )
    end
end

function run_mcts(
    net::Network,
    config::MuZeroConfig, 
    roots::Vector{Node}, 
    action_histories::Vector{ActionHistory}
    )
    min_max_stats_set = MinMaxStats(config.known_bounds, length(roots))
    logging = false

    for sim_idx in 1:config.num_simulations
        if logging
            println()
            @show sim_idx
        end
        histories = clone.(action_histories)
        nodes = roots
        search_paths = [[node] for node in nodes]
        nodes = searchexpanded!.(
            search_paths, 
            nodes, 
            histories, 
            min_max_stats_set,
            config,
        )
        lens = length.(search_paths)
        parents = getindex.(search_paths, lens .- 1)

        inference = recurrent_inference(
            net, 
            gethiddenstate.(parents), 
            last_action.(histories)
        ) |> NetworkOutput

        r = reward(inference)
        v = value(inference)

        expand_node!(
            nodes, 
            legal_actions.(config.game_type, histories), 
            inference, 
            r[1, :]
        )
        
        backpropagate!.(
            search_paths, 
            r[1, :], 
            to_play.(histories),
            config.discount, 
            min_max_stats_set
        )
    end
end

function searchexpanded!(
    search_path, 
    node, 
    history, 
    min_max_stats,
    config,
    )
    while expanded(node)
        action, node = select_child!(
            node,
            min_max_stats,
            config,
            )
        add_action!(history, action)
        push!(search_path, node)
    end
    node
end
