include("EDPolicy.jl")

export EDManager

## definition
"""
    EDManager(agents::Dict{Any, EDPolicy})

A special MultiAgentManager in which all agents use Exploitability Descent(ED) algorithm to play the game.
"""
mutable struct EDManager <: AbstractPolicy
    agents::Dict{Any, EDPolicy}
end

## interactions
function (π::EDManager)(env::AbstractEnv)
    player = current_player(env)
    if player == chance_player(env)
        rand(legal_action_space(env))
    else
        π.agents[player](env)
    end
end

RLBase.prob(π::EDManager, env::AbstractEnv, args...) = prob(π.agents[current_player(env)], env, args...)

## run function
function Base.run(
    π::EDManager,
    env::AbstractEnv,
    stop_condition = StopAfterEpisode(1),
    hook::AbstractHook = EmptyHook(),
)
    @assert NumAgentStyle(env) == MultiAgent(2) "ED algorithm only support 2-players games."
    @assert UtilityStyle(env) isa ZeroSum "ED algorithm only support zero-sum games."

    is_stop = false

    while !is_stop
        RLBase.reset!(env)
        hook(PRE_EPISODE_STAGE, π, env)

        for (player, policy) in π.agents
            # construct opponent's best response policy.
            oppo_best_response = BestResponsePolicy(π, env, policy.opponent)
            # update player's policy by using policy-gradient.
            update!(policy, oppo_best_response, env, player)
        end

        # run one episode for update stop_condition
        RLBase.reset!(env)
        while !is_terminated(env)
            π(env) |> env
        end

        if stop_condition(π, env)
            is_stop = true
            break
        end
        hook(POST_EPISODE_STAGE, π, env)
    end
    hook(POST_EXPERIMENT_STAGE, π, env)
    hook
end
