struct ExploitabilitySolver{G<:Game, K}
    "Game upon which a solver is to be evaluated"
    game        :: G

    "EXPLOITING Player"
    p           :: Int

    "Child information states of being in information state `I` and taking action `a`"
    Ia_children :: Dict{Tuple{K, Int}, Set{K}}

    "Reach probability weighted utility of being in information state `I` and taking action `a`"
    utility     :: Dict{K,Vector{Float64}}

    "π⁻ⁱ(h₀ → I)"
    prob        :: Dict{K,Float64}

    "Information states where player `p` is first allowed to act"
    roots       :: Set{K}
end

function reset!(e_sol::ExploitabilitySolver)
    for U in values(e_sol.utility)
        fill!(U, 0.0)
    end
    for I in keys(e_sol.prob)
        e_sol.prob[I] = 0.0
    end
    nothing
end

function ExploitabilitySolver(game::Game, p::Int)
    I = infokeytype(game)
    e_sol = ExploitabilitySolver(
        game,
        other_player(p),
        Dict{Tuple{I,Int},Set{I}}(),
        Dict{I,Vector{Float64}}(),
        Dict{I,Float64}(),
        Set{I}()
    )
    return populate!(e_sol)
end

function ExploitabilitySolver(sol::AbstractCFRSolver, p::Int)
    return ExploitabilitySolver(sol.game, p)
end

function populate!(e_sol::ExploitabilitySolver)
    _populate!(e_sol, initialhist(e_sol.game), nothing, 0)
    return e_sol
end

function _populate!(e_sol::ExploitabilitySolver, h, last_I, last_a_idx::Int)
    game = e_sol.game
    if isterminal(game, h)
        return nothing
    elseif iszero(player(game, h))
        for a in chance_actions(game, h)
            _populate!(e_sol, next_hist(game, h, a), last_I, last_a_idx)
        end
        return nothing
    end

    A = actions(game, h)
    I = infokey(game, h)

    if player(game, h) == e_sol.p
        if isnothing(last_I)
            push!(e_sol.roots, I)
        else
            Ia_children = get(e_sol.Ia_children,(last_I, last_a_idx), nothing)
            if isnothing(Ia_children)
                e_sol.Ia_children[(last_I, last_a_idx)] = Set{typeof(I)}()
            else
                push!(Ia_children, I)
            end
        end
        if !haskey(e_sol.utility, I)
            e_sol.utility[I] = zeros(Float64, length(A))
            e_sol.prob[I] = 0.0
        end
        for (k,a) in enumerate(A)
            _populate!(e_sol, next_hist(game, h, a), I, k)
        end
    else
        for a in A
            _populate!(e_sol, next_hist(game, h, a), last_I, last_a_idx)
        end
    end
    return nothing
end

function utility!(
    e_sol::ExploitabilitySolver,
    solver::AbstractCFRSolver,
    h,
    last_I,
    last_a_idx::Int,
    π_ni::Float64)

    (;game, p) = e_sol
    if isterminal(game, h)
        # π_I doesn't represent true probability, as it's being built incrementally
        # π_I increases every time it is "hit"
        # π_I = e_sol.prob[last_I]
        # e_sol.utility[last_I][last_a_idx] += (π_ni/π_I)*utility(game, p, h)
        e_sol.utility[last_I][last_a_idx] += π_ni*utility(game, p, h)
        return nothing
    elseif iszero(player(game, h)) # chance player
        A = chance_actions(game, h)
        inv_l_a = inv(length(A))
        for a in A
            utility!(e_sol, solver, next_hist(game, h, a), last_I, last_a_idx, π_ni*inv_l_a)
        end
        return nothing
    end

    A = actions(game, h)
    I = infokey(game, h)

    if player(game, h) == p
        e_sol.prob[I] += π_ni
        for (k,a) in enumerate(A)
            h′ = next_hist(game, h, a)
            utility!(e_sol, solver, h′, I, k, π_ni)
        end
        return nothing
    else
        σ = strategy(solver, I)
        for (k,a) in enumerate(A)
            h′ = next_hist(game, h, a)
            utility!(e_sol, solver, h′, last_I, last_a_idx, π_ni*σ[k])
        end
        return nothing
    end
end

function exploitability(e_sol::ExploitabilitySolver, sol::AbstractCFRSolver)
    reset!(e_sol)
    h0 = initialhist(e_sol.game)
    utility!(e_sol, sol, h0, nothing, 0, 1.0)

    for (I,U) in e_sol.utility
        U ./= e_sol.prob[I]
    end

    u_exp = 0.0
    for I in e_sol.roots
        u_exp += e_sol.prob[I]*best_response(e_sol, I)
    end
    u_current = evaluate(sol, e_sol.p)
    return u_exp - u_current
end

function best_response(e_sol::ExploitabilitySolver, I)
    util = e_sol.utility[I]
    π_A = e_sol.prob[I]

    Q_max = -Inf
    for a_i in eachindex(util)
        children = get(e_sol.Ia_children, (I,a_i), nothing)
        Qa_i = 0.0
        Qa_i += π_A*util[a_i]
        if !isnothing(children)
            for I′ in children
                Qa_i += e_sol.prob[I′]*best_response(e_sol, I′)
            end
        end
        if Qa_i > Q_max
            Q_max = Qa_i
        end
    end
    return inv(π_A)*Q_max
end
