using Flux: ADAM, Momentum

#=
Each self-play job is independent of all others; it takes the latest network
snapshot, produces a game and makes it available to the training job by
writing it to a shared replay buffer.
=#

using Base.Threads

struct SelfplayCallback
    period::Integer
    fn::Function
end

(x::SelfplayCallback)() = x.fn()

function getnextactions(
    nodes::Vector{Node},
    plans::AbstractArray;
    num_moves,
    num_steps, 
    temp_fn,
    )
    if !isempty(plans)
        return getfield.(popfirst!.(plans), :idx)
    end
    select_action.(
        nodes;
        num_moves=num_moves,
        num_steps=num_steps,
        temp_fn=temp_fn
    )  
end

#=
Core Monte Carlo Tree Search algorithm.
To decide on an action, we run N simulations, always starting at the root of
the search tree and traversing the tree according to the UCB formula until we
reach a leaf node.
=#

function select_action(
    node::Node;
    num_moves,
    num_steps, 
    temp_fn,
    )
    softmax_sample(
        [(child.visit_count, action) for (action, child) in node.children], 
        temp_fn(num_moves, num_steps)
    )
end


# Select the child with the highest UCB score.
function select_child!(
    node::Node, 
    min_max_stats::MinMaxStats,
    config,
    )
    score_action_child = [
        (
            ucbscore(node, child, min_max_stats, config), 
            action, 
            child
        ) 
        for (action, child) in node.children
    ]
    scores = [x[1] for x in score_action_child]
    max_score = maximum(scores)
    selected = rand(score_action_child[scores .== max_score ])
    _, action, child = selected
    action, child
end


#=
The score for a node is based on its value, plus an exploration bonus based on
the prior.
=#

function ucbscore(
    parent::Node, 
    child::Node, 
    min_max_stats::MinMaxStats,
    config;
    action=nothing,
    )
    discount = config.discount
    pb_c_init = config.pb_c_init
    pb_c_base = config.pb_c_base
    pb_c = log((parent.visit_count + pb_c_base + 1) / pb_c_base) + pb_c_init
    pb_c *= sqrt(parent.visit_count) / (child.visit_count + 1)
    prior_score = pb_c * child.prior
    value_score = 0
    neg_nextstate_value = -1*(discount * normalize(min_max_stats, calc_value(child)))
    if child.visit_count > 0
        value_score = child.reward + neg_nextstate_value
    end
    ucb_score = prior_score + value_score
    if !isnothing(action)
        if parent.state.board == [0 2 0; 0 0 1; 0 1 0]
            @show action, value_score, child.value_sum, child.visit_count #child.prior, prior_score, ucb_score
        end
    end
    ucb_score
end

#= 
At the start of each search, we add dirichlet noise to the prior of the root
to encourage the search to explore new actions.
=#
using Distributions

function addexplorationnoise(
    node::Node; 
    root_dirichlet_alpha, 
    root_exploration_fraction
    )
    actions = collect(keys(node.children))
    noise = rand(
        Dirichlet(
            length(actions), 
            root_dirichlet_alpha
        )
    )
    for (a, n) in zip(actions, noise)
        node.children[a].prior = (
            node.children[a].prior 
            * (1 - root_exploration_fraction) 
            + (n * root_exploration_fraction)
        )
    end
end

function add_exploration_noise(
    config::MuZeroConfig, 
    nodes::Vector{Node}
    )
    addexplorationnoise.(
        nodes;
        root_dirichlet_alpha=config.root_dirichlet_alpha,
        root_exploration_fraction=config.root_exploration_fraction
    )
end

function add_exploration_noise(config::MuZeroConfig, node::Node)
    addexplorationnoise(
        node;
        root_dirichlet_alpha=config.root_dirichlet_alpha,
        root_exploration_fraction=config.root_exploration_fraction
    )
end

# function play_me()
#     network = test_network
#     g = config.game_type()
#     to_play = PLAYER_ID
#     output = initial_inference(network, current_image(g)) |> SingleNetworkOutput
#     function get_next(output)
#         to_play = (to_play == PLAYER_ID) ? OPPONENT_ID : PLAYER_ID
#         if to_play == OPPONENT_ID
#             println(); print_game(g); println();
#             next = OPPONENT_ID
#             while !islegal(g, next)
#                 println(); print("Your move [1-9]: ")
#                 next = parse(Int, readline())
#             end
#             return next
#         end
#         greedy_action(output)
#     end
#     while !is_terminal(g)
#         next = get_next(output)
#         if !islegal(g, next)
#             println("bot made illegal move... ($next)")
#             return
#         end
#         apply!(g, next)
#         output = recurrent_inference(network, hidden_state(output), [next]) |> SingleNetworkOutput
#     end
#     println(); print_game(g); println();
#     println(); println("Done.")
#     g
# end

# function dynamics_passed(game_type, network)
#     g = game_type()
#     # @todo This needs attention (player 1 is can only be first for certain games)
#     player = rand(1:2) == PLAYER_ID ? PLAYER_ID : OPPONENT_ID
#     output = NetworkOutput(initial_inference(network, current_image(g)))
#     function get_next(output)
#         @assert false "Fix logic below for OPPONENT_ID"
#         player *= OPPONENT_ID
#         if player == OPPONENT_ID
#             return rand(legal_actions(g))
#         end
#         greedy_action(output)
#     end
#     while !is_terminal(g)
#         next = get_next(output)
#         if !islegal(g, next)
#             return false, g
#         end
#         apply!(g, next)
#         output = recurrent_inference(network, output.hidden_state, [next]) |> SingleNetworkOutput
#     end
#     true, g
# end


function check_dynamics(config::MuZeroConfig, network)
    passes = 0
    fails = 0
    failed_games = []
    for _ in 1:1000
        res, g = dynamics_passed(config.game_type, network)
        if res
            passes += 1
        else 
            push!(failed_games, g)
            fails += 1
        end
    end
    passes, fails, failed_games
end

#=
At the end of a simulation, we propagate the evaluation all the way up the
tree to the root.
=#

function backpropagate!(
    search_path::Vector{Node}, 
    last_value::Float32, 
    to_play::Integer,
    discount::Float32, 
    min_max_stats::MinMaxStats
    )
    value = last_value
    #backprops = []
    for node in reverse(search_path)
        # @todo check on the logic below and make sure value_sum is correct
        node.value_sum += (gettoplay(node).idx == to_play) ? value : -value
        #node.info = "$(node.to_play == to_play ? "v" : "-v") -> $(node.value_sum)"
        node.visit_count += 1
        update!(min_max_stats, calc_value(node))
        value = Float32(-1)*node.reward + discount * value
        #push!(backprops, (state=key(node.state), value_sum=node.value_sum))
    end
    #backprops
    return
end

# MuZero training is split into two independent parts: Network training and
# self-play data generation.
# These two parts only communicate by transferring the latest network checkpoint
# from the training to the self-play, and the finished games from the self-play
# to the training.

using Base.Threads: @spawn, @threads

launch_job(f, args...) = f(args...)