using Printf

# @todo does Node need to be mutable?
mutable struct Node
    visit_count::Int32
    prior::Float32
    value_sum::Float32
    children::Dict{Int32,Node}
    hidden_state::Union{Nothing,AbstractArray{Float32, 3}}
    reward::Float32
    info::String
    state::AbstractEnvState # @todo - remove state altogether
    prior_policy::Dict
    action_scores::Dict
    function Node(prior, state; visit_count=0, value_sum=0, 
        children=Dict(), hidden_state=nothing, reward=0)
        new(visit_count, prior, value_sum, children, hidden_state, reward, "", state, Dict(), Dict())
    end
end

gethiddenstate(n::Node) = n.hidden_state

gettoplay(n::Node) = n.state.to_play

function json(x::Node, visit_pr=0.0)
    d = Dict()
    noisy_prior_policy = Dict()
    for (action, node) in x.children
        noisy_prior_policy[action] = node.prior
    end
    d["visit_pr"] = visit_pr
    d["visit_count"] = x.visit_count
    d["to_play"] = get_player_label(x)
    d["prior"] = x.prior
    d["value_sum"] = x.value_sum
    d["prior_policy"] = x.prior_policy
    d["noisy_prior_policy"] = noisy_prior_policy
    d["reward"] = x.reward
    d["children"] = Dict()
    d["action_scores"] = x.action_scores
    total_child_visit_count = sum([child.visit_count for child in values(x.children)])
    for (action, child) in x.children
        visit_pr = child.visit_count / total_child_visit_count
        d["children"][action] = json(child, visit_pr)
    end
    d
end

get_player_label(x::Node) = x.state.to_play.idx == PLAYER_ID ? "A" : "B"

function Base.show(io::IO, x::Node)
    show_indented(io, x)
end

function show_indented(io::IO, x::Node, level=0, action=nothing)
    #show(io, x.channels)
    print(" "^(level*5))
    action_label = isnothing(action) ? "" : "A($action)"
    hex_state = "hex_str" #hex_str(x.state)
    @printf "%s %s pr: %.2f {c: %i, r: %i, vs: %.3f}|%s %s" action_label hex_state x.prior x.visit_count x.reward x.value_sum get_player_label(x) x.info
    println(io)
    actions = [action for action in keys(x.children)]
    for action in sort(actions)
        child = x.children[action]
        show_indented(io, child, level+1, action)
    end
end

expanded(x::Node) = length(x.children) > 0

function calc_value(x::Node)
    Float32(x.visit_count == 0 ? 0.0 : x.value_sum / x.visit_count)
end

function priors(x::Node)
    sort([(action, child.prior) for (action, child) in x.children])
end

#=
We expand a node using the value, reward and policy prediction obtained from
the neural network.
=#

function expand_node!(node::Node, actions::Vector, output::SingleNetworkOutput, reward::Number)
    p_logits = policy_logits(output)
    hidden = hidden_state(output)
    node.hidden_state = hidden
    node.reward = Float32(reward)

    for (action, pr) in enumerate(softmax(p_logits))
        node.prior_policy[action] = pr
    end

    if !isempty(actions)
        policy = Dict([a=>exp(p_logits[a.idx]) for a in actions])
        policy_sum = sum(values(policy))
        for (action, p) in policy
            node.children[action.idx] = Node(p / policy_sum, getnextstate(node.state, action)) #reward=get_reward(node.state, action)
        end
    end
    
end

function expand_node!(
    nodes::Vector{Node}, 
    actionSets::Vector,
    output::NetworkOutput, 
    reward::Vector{Float32}
    )
    setfield!.(
        nodes, 
        :hidden_state, 
        eachslice(hidden_state(output), dims=4)
    )
    setfield!.(
        nodes, 
        :reward, 
        Float32.(reward)
    )
    #### @todo - efficiency could be much improved
    p_logits_tensor = policy_logits(output)
    for (idx, actions) in enumerate(actionSets)
        node = nodes[idx]
        p_logits = view(p_logits_tensor, :, idx)
        for (action, pr) in enumerate(softmax(p_logits))
            node.prior_policy[action] = pr
        end
        if !isempty(actions)
            policy = Dict([a=>exp(p_logits[a.idx]) for a in actions])
            policy_sum = sum(values(policy))
            for (action, p) in policy
                node.children[action.idx] = Node(p / policy_sum, getnextstate(node.state, action))
            end
        end
    end
    ####
end