function getrewardtests(T::Type{Tac})
    tests = []
    for (idx,s) in enumerate(states(T))
        action_rewards = map(a->(a, get_reward(s, a)), getlegalactions(s))
        reward_tests = map(action_reward->begin
            a, r = action_reward
            RewardTest(
                idx,
                s,
                get_image(s),
                a,
                get_action_image(T, a),
                r
            )
        end, action_rewards)
        append!(tests, reward_tests)
    end
    RewardTests(tests)
end

function MuZero.makepolicytests(T::Type{Tac})
    value_fn = getvaluefn(T)
    actionvaluefn = getactionvaluefn(T; value_fn=value_fn)
    tests = []
    for (idx,s) in enumerate(getstates(T))
        action_values = map(a->begin
            value = actionvaluefn(s,a)
            (a.idx, value)
        end, getlegalactions(s))
        max_value = maximum(x->x[2], action_values)
        ideal_actions = map(first, filter(x->x[2] == max_value, action_values))
        pt = PolicyTest(
            idx,
            s,
            get_image(s),
            ideal_actions
        )
        push!(tests, pt)
    end
    PolicyTests(tests)
end

function getoptimalactionsfn(
    T::Type{Tac}; 
    actionvaluefn=getactionvaluefn(T)
    )
    v = Dict()
    for s in getstates(T)
        action_values = map(a->begin
            (a, actionvaluefn(s,a))
        end, getlegalactions(s))
        max_value = maximum(x->x[2], action_values)
        ideal_actions = map(first, filter(x->x[2] == max_value, action_values))
        v[s] = ideal_actions
    end
    v
end

function makerrvptests(T::Type{Tac})
    value_fn = getvaluefn(T)
    optimalactionsfn = getoptimalactionsfn(
        T; 
        actionvaluefn=getactionvaluefn(T; value_fn=value_fn)
    )
    tests = []
    for (idx,s) in enumerate(getstates(T))
        for a in getlegalactions(s)
            reward_next, s_next = dynamics(s, a)
            if !isterminal(s_next)
                rrvp_test = RecurrentRVPTest(
                    "$idx-$(a.idx)",
                    s,
                    s_next,
                    get_image(s),
                    a,
                    get_action_image(T, a),
                    map(x->x.idx, optimalactionsfn[s_next]),
                    value_fn[s_next],
                    reward_next,
                )
                push!(tests, rrvp_test)
            end
        end
    end
    RecurrentRVPTests(tests)
end

# @todo - this should be a MuZero function 
# (i.e. MuZero.makepolicyvaluetests)
function makepolicyvaluetests(T::Type{Tac})
    value_fn = getvaluefn(T)
    actionvaluefn = getactionvaluefn(T; value_fn=value_fn)
    tests = []
    for (idx,s) in enumerate(getstates(T))
        action_values = map(a->begin
            value = actionvaluefn(s,a)
            (a.idx, value)
        end, getlegalactions(s))
        max_value = maximum(x->x[2], action_values)
        ideal_actions = map(first, filter(x->x[2] == max_value, action_values))
        actual_value = value_fn[s]
        pt = PolicyValueTest(
            idx,
            s,
            get_image(s),
            ideal_actions,
            actual_value
        )
        push!(tests, pt)
    end
    PolicyValueTests(tests)
end

function getoptimalpolicyfn(T::Type{Tac})
    actionvaluefn = getactionvaluefn(T)
    (s)->begin
        action_values = map(a->begin
            value = actionvaluefn(s,a)
            (a.idx, value)
        end, getlegalactions(s))
        max_value = maximum(x->x[2], action_values)
        ideal_actions = map(first, filter(x->x[2] == max_value, action_values))
        idealpolicy = zeros(Float16, length(getactions(s)))
        idealpolicy[ideal_actions] .= 1/length(ideal_actions)
        idealpolicy
    end
end

function getactionvaluefn(T::Type{Tac}; value_fn=getvaluefn(T))
    action_values = Dict()
    stateactionhash(s,a) = hash("$(hash(s))-$(hash(a))")
    for s in getstates(T)
        for a in getlegalactions(s)
            reward, s_next = dynamics(s, a)
            value = reward + -1*value_fn[s_next]
            action_values[stateactionhash(s,a)] = value
        end
    end
    (s,a)->action_values[stateactionhash(s,a)]
end

function MuZero.getstates(T::Type{Tac})
    terminal(s) = !isnothing(getwinner(s)) || length(getlegalactions(s)) == 0
    nonterminal(s) = !isterminal(s)
    filter(nonterminal, getstatesplus(T))
end

using Combinatorics

function Chess()
    board = emptyboard(ChessState)
    board[1, 1:8] = getpieceidx(Chess, [
        :rook_white, :knight_white,
        :bishop_white, :king_white, 
        :queen_white, :bishop_white, 
        :knight_white, :rook_white,
    ])
    board[2, 1:8] = getpieceidx(Chess, [
        :pawn_white, :pawn_white, 
        :pawn_white, :pawn_white, 
        :pawn_white, :pawn_white, 
        :pawn_white, :pawn_white,
    ])
    board[7, 1:8] = getpieceidx(Chess, [
        :pawn_black, :pawn_black, 
        :pawn_black, :pawn_black, 
        :pawn_black, :pawn_black, 
        :pawn_black, :pawn_black,
    ])
    board[8, 1:8] = getpieceidx(Chess, [
        :rook_black, :knight_black, 
        :bishop_black, :queen_black, 
        :king_black, :bishop_black, 
        :knight_black, :rook_black,
    ])
    player = getplayer(ChessState, :white)
    states = [ChessState(board, player)]
    Chess([],[],states)
end

function MuZero.getstatesplus(T::Type{Tac})
    toplay(pieces) = isodd(length(filter(isnothing, pieces))) ? getplayer(T, :X) : getplayer(T, :O)
    perms(pieces) = pieces |> permutations |> collect |> unique
    a0 = [nothing for _ in 1:9]
    a1 = [:X, [nothing for _ in 1:8]...]
    a2 = [:X, :O, [nothing for _ in 1:7]...]
    a3 = [:X, :O, :X, [nothing for _ in 1:6]...]
    a4 = [:X, :O, :X, :O, [nothing for _ in 1:5]...]
    a5 = [:X, :O, :X, :O, :X, [nothing for _ in 1:4]...]
    a6 = [:X, :O, :X, :O, :X, :O, [nothing for _ in 1:3]...]
    a7 = [:X, :O, :X, :O, :X, :O, :X, [nothing for _ in 1:2]...]
    a8 = [:X, :O, :X, :O, :X, :O, :X, :O, nothing]
    a9 = [:X, :O, :X, :O, :X, :O, :X, :O, :X]
    all = []
    for aX in [a0, a1, a2, a3, a4, a5, a6, a7, a8, a9]
        append!(all, perms(aX))
    end
    [TacState(getboard(T, pieces), toplay(pieces)) for pieces in all]
end

function getvaluefn(T::Type{Tac}; max_i=typemax(Int), θ=Float32(0.1))
    v = Dict()
    for s in getstatesplus(T)
        v[s]=Float32(0)
    end
    function player_action_value(state, a)
        reward, s_next = dynamics(state, a)
        reward + -1 * v[s_next]
    end
    for i in 1:max_i
        Δ, num_updated = Float32(0), 0
        for state in getstates(T)
            v_prev = v[state]
            legal = getlegalactions(state)
            state_action_values = map(a->player_action_value(state, a), legal)
            max_found = maximum(state_action_values)
            if v[state] != max_found  
                v[state] = max_found
                num_updated += 1
            end
            Δ = max(Δ, abs(v_prev - v[state]))
        end
        vs = v |> values |> collect
        pos = vs[vs .> 0] |> length
        neg = vs[vs .< 0] |> length
        #println("Δ=$Δ($i), $pos, $neg (num_updated: $num_updated)")
        i += 1
        (Δ >= θ && i <= max_i) || break
    end
    v
end