export EDPolicy

## definition
"""
    EDPolicy(opponent, learner, explorer)

Exploitability Descent(ED) algorithm directly updates the player’s policy against a worst 
case opponent(best reponse) in a two-player zero sum game. On each iteration, the ED algorithm 
performs the following update for each player:

1. Construct a (deterministic) best response policy of the opponent to the current policy;
2. Compute the value of every action in every state when playing current policy vs the 
   best response;
3. Update player's current policy to do better vs opponent's best response by performing 
   a policy-gradient update.

# Keyword Arguments

- `opponent::Any`, the opponent's name.
- `learner::NeuralNetworkApproximator`, used to get the value of each action. 
- `explorer::AbstractExplorer`

# Ref

[Computing Approximate Equilibria in Sequential Adversarial Games by Exploitability Descent](https://arxiv.org/abs/1903.05614)
"""
mutable struct EDPolicy{P<:NeuralNetworkApproximator, E<:AbstractExplorer}
    opponent::Any
    learner::P
    explorer::E
end

Flux.functor(x::EDPolicy) = (learner = x.learner,), y -> begin
    x = @set x.learner = y.learner
    x
end

## interactions with the environment
function (π::EDPolicy)(env::AbstractEnv)
    s = state(env)
    s = send_to_device(device(π.learner), Flux.unsqueeze(s, ndims(s) + 1))
    logits = π.learner(s) |> vec |> send_to_host
    ActionStyle(env) isa MinimalActionSet ? π.explorer(logits) : 
        π.explorer(logits, legal_action_space_mask(env))
end
# set the `_device` function for convenience transferring the variable to the corresponding device.
_device(π::EDPolicy, x) = send_to_device(device(π.learner), x)

function RLBase.prob(π::EDPolicy, env::AbstractEnv; to_host::Bool = true)
    s = @ignore state(env) |> x -> Flux.unsqueeze(x, ndims(x) + 1) |> x -> _device(π, x)
    logits = π.learner(s) |> vec
    mask = @ignore legal_action_space_mask(env) |> x -> _device(π, x)
    p = ActionStyle(env) isa MinimalActionSet ? prob(π.explorer, logits) : prob(π.explorer, logits, mask)
    to_host ? p |> send_to_host : p
end

function RLBase.prob(π::EDPolicy, env::AbstractEnv, action)
    A = action_space(env)
    P = prob(π, env)
    @assert length(A) == length(P)
    if A isa Base.OneTo
        P[action]
    else
        for (a, p) in zip(A, P)
            if a == action
                return p
            end
        end
        @error "action[$action] is not found in action space[$(action_space(env))]"
    end
end

## update policy
function RLBase.update!(
    π::EDPolicy, 
    Opponent_BR::BestResponsePolicy, 
    env::AbstractEnv,
    player::Any,
)
    reset!(env)

    # construct policy vs best response
    policy_vs_br = PolicyVsBestReponse(
        MultiAgentManager(
            NamedPolicy(player, π),
            NamedPolicy(π.opponent, Opponent_BR),
            ),
        env,
        player,
    )
    info_states = collect(keys(policy_vs_br.info_reach_prob))
    cfr_reach_prob = collect(values(policy_vs_br.info_reach_prob)) |> x -> _device(π, x)

    gs = gradient(Flux.params(π.learner)) do
        # Vector of shape `(length(info_states), 1)`
        # compute expected reward from the start of `e` with policy_vs_best_reponse
        # baseline = ∑ₐ πᵢ(s, a) * q(s, a)
        baseline = @ignore Flux.stack(([values_vs_br(policy_vs_br, e)] for e in info_states), 1) |> x -> _device(π, x)
        
        # Vector of shape `(length(info_states), length(action_space))`
        # compute expected reward from the start of `e` when playing each action.
        q_values = Flux.stack((q_value(π, policy_vs_br, e) for e in info_states), 1)

        advantage = q_values .- baseline
        # Vector of shape `(length(info_states), length(action_space))`
        # get the prob of each action with `e`, i.e., πᵢ(s, a).
        policy_values = Flux.stack((prob(π, e, to_host = false) for e in info_states), 1)

        # get each info_state's loss
        # ∑ₐ πᵢ(s, a) * (q(s, a) - baseline), where baseline = ∑ₐ πᵢ(s, a) * q(s, a).
        loss_per_state = - sum(policy_values .* advantage, dims = 2)

        sum(loss_per_state .* cfr_reach_prob)
    end
    update!(π.learner, gs)
end

## Supplement struct for Computing related results when player's policy versus opponent's best_response.
struct PolicyVsBestReponse{E, P<:MultiAgentManager}
    info_reach_prob::Dict{E, Float64}
    player::Any
    policy::P
end

function PolicyVsBestReponse(policy, env, player)
    E = typeof(env)

    p = PolicyVsBestReponse(
        Dict{E, Float64}(),
        player,
        policy,
    )
    
    e = copy(env)
    RLBase.reset!(e)
    get_cfr_prob!(p, e)
    p
end

function get_cfr_prob!(p::PolicyVsBestReponse, env::AbstractEnv, reach_prob = 1.0)
    if !is_terminated(env)
        if current_player(env) == p.player
            p.info_reach_prob[env] = reach_prob

            for a in legal_action_space(env)
                get_cfr_prob!(p, child(env, a), reach_prob)
            end
        elseif current_player(env) == chance_player(env)
            for (a, pₐ) in zip(action_space(env), prob(env))
                if pₐ > 0
                    get_cfr_prob!(p, child(env, a), reach_prob * pₐ)
                end
            end
        else # for opponent
            a = p.policy(env) # deterministic
            get_cfr_prob!(p, child(env, a), reach_prob)
        end
    end
end

function values_vs_br(p::PolicyVsBestReponse, env::AbstractEnv)
    if @ignore is_terminated(env)
        @ignore reward(env, p.player)

    elseif @ignore current_player(env) == p.player
        v = 0.0
        A, P = @ignore action_space(env), prob(p.policy.agents[p.player].policy, env)
        for (a, pₐ) in zip(A, P)
            if pₐ > 0
                v += pₐ * values_vs_br(p, @ignore child(env, a))
            end
        end
        v
    # for game which has two or more rounds.
    elseif @ignore current_player(env) == chance_player(env)
        v = 0.0
        A, P = @ignore (action_space(env), prob(env))
        for (a, pₐ) in zip(A, P)
            if pₐ > 0
                v += pₐ * values_vs_br(p, @ignore child(env, a))
            end
        end
        v
    else # for opponent
        a = @ignore p.policy(env) # deterministic
        values_vs_br(p, @ignore child(env, a))
    end
end

function q_value(π::EDPolicy, p::PolicyVsBestReponse, env::AbstractEnv)
    P, A = prob(π, env) , @ignore action_space(env)
    v = []
    for (a, pₐ) in zip(A, P)
        value = pₐ == 0 ? pₐ : values_vs_br(p, @ignore child(env, a))
        v = vcat(v..., value)
    end
    _device(π, v)
end
