using DataFrames
using Grids
using Crayons

action_space_size(g::MuZeroGame)                        = action_space_size(g.game)
to_play(g::MuZeroGame)                                  = gettoplay(g.game)
to_play_at(g::MuZeroGame, idx::Integer)                 = to_play_at(g.game, idx) 
next_to_play(g::MuZeroGame)                             = getnexttoplay(g.game)
discount(g::MuZeroGame)                                 = discount(g.game)
reward_matrix(g::MuZeroGame)                            = reward_matrix(g.game)
value_matrix(g::MuZeroGame, idx::Integer; kwargs...)    = value_matrix(g.game, idx; kwargs...)
uniform_policy(g::MuZeroGame)                           = uniform_policy(g.game)
history(g::MuZeroGame)                                  = gethistory(g.game)
is_terminal(g::MuZeroGame)                              = isterminal(g.game)
legal_actions(g::MuZeroGame)                            = getlegalactions(g.game)
action_history(g::MuZeroGame, idx)                      = action_history(g.game, idx)
action_history(g::MuZeroGame)                           = action_history(g.game)
apply!(g::MuZeroGame, action::Integer)                  = step!(g.game, action)
make_image(g::MuZeroGame, state_idx::Integer)           = make_image(g.game, state_idx)
#getstate(g::MuZeroGame, state_idx::Integer)             = get_state(g, state_idx)
get_state(g::MuZeroGame, state_idx::Integer)            = getstate(g.game, state_idx)
make_image_label(g::MuZeroGame, state_idx::Integer)     = make_image_label(g.game, state_idx)
get_state(g::MuZeroGame)                                = getstate(g.game)

struct MuZeroGameSVG
    game::MuZeroGame
end

getgame(mzg::MuZeroGame) = mzg.game
alltrue(itr) = length(itr)==0 || all(identity, itr)
is_terminal(games::Vector) = is_terminal.(games) |> alltrue
movesmaxedout(game::MuZeroGame, max_moves) = length(history(game)) >= max_moves
movesmaxedout(games::Vector, max_moves) = movesmaxedout.(games, max_moves) |> alltrue

function svg(g::MuZeroGame)
    MuZeroGameSVG(g)
end

function Base.show(io::IO, m::MIME"image/svg+xml", g::MuZeroGameSVG)
    show(io, m, getmoves(g.game.game))
end


Base.getindex(x::MuZeroGame, idx::Integer) = getindex(x.game, idx)

function df(x::MuZeroGame; range=1:length(x.game.states), limit=3)
    g = x.game
    state_col = "State #"
    action_col = "Action"
    reward_col = "Reward"
    d = Dict(
        state_col=>[],
        action_col=>[],
        reward_col=>[],
    )
    for s_idx in range
        action = s_idx > length(g.actions) ? nothing : g.actions[s_idx]
        reward = s_idx > length(g.rewards) ? nothing : g.rewards[s_idx]
        push!(d[state_col], s_idx)
        push!(d[action_col], isnothing(action) ? nothing : action.idx)
        push!(d[reward_col], reward)
    end
    DataFrame(
        state_col=>d[state_col],
        action_col=>d[action_col],
        reward_col=>d[reward_col],
    )
end

function Base.show(io::IO, m::MIME"text/plain", x::MuZeroGame)
    child_visit_strings = []
    for visits in x.child_visits
        from = 1 
        limit = 5
        to = min(limit, length(visits))
        actions = reverse(sortperm(visits))[from:to]
        prs = visits[actions]
        action_prs = map(
            x->begin
                action, pr = x
                pr_s = @sprintf "%.2f" pr
                "$action ($(pr_s))"
            end, 
            zip(actions, prs)
        )
        push!(child_visit_strings, join(action_prs, ", "))
    end
    game_df = df(x)
    root_values_str = map(root_value->begin
        @sprintf "%.2f" root_value  
    end, x.root_values)
    game_df[!, "V(S)"] = [root_values_str..., nothing]
    game_df[!, "Sampled Actions (Top 5)"] = [child_visit_strings..., nothing]
    game_df = game_df[!, [1, 4, 5, 2, 3]]
    println()
    pretty_game(
        io,
        game_df; 
        title="MuZero Game History"
    )
end

function pretty_game(io::IO, df::DataFrame; title)
    reward_hl = Highlighter((data, i, j)->(j==5) && !isnothing(data[i,j]) && data[i,j] > 0 , crayon"green");
    pretty_table(
        io,
        df; 
        show_row_number=false, 
        nosubheader=true, 
        title=title,
        highlighters=(reward_hl),
    )
end

# @todo make this more readable
# @todo use exclamation point to indicate it is mutating
function store_search_statistics(game::MuZeroGame, root::Node)
    sum_visits = sum([child.visit_count for child in values(root.children)])
    action_space = [idx for idx in 1:action_space_size(game)]
    push!(game.child_visits, 
        [in(a, keys(root.children)) ? (root.children[a].visit_count / sum_visits) : 0 for a in action_space]
    )
    push!(game.root_values, calc_value(root))
end

# @todo This function needs extensive unit testing for edge cases, especially for non-endgame rewards.
function make_target(game::MuZeroGame, state_index::Integer, num_unroll_steps::Integer, td_steps::Integer)
    # The value target is the discounted root value of the search tree N steps
    # into the future, plus the discounted sum of all rewards until then.
    targets = []
    rewards = reward_matrix(game)
    num_rewards = size(rewards, 1)

    for (target_id, current_index) in enumerate(state_index:(state_index + num_unroll_steps))
        if current_index <= num_rewards
            player = to_play_at(game, current_index-1)
            next_player = to_play_at(game, current_index)
            bootstrap_index = current_index + td_steps
            value = 0
            if bootstrap_index <= length(game.root_values)
                @assert false "This should not be executing for board games"
                value = game.root_values[bootstrap_index] * discount(game)^td_steps
            end
            # For simplicity the network always predicts the most recently received
            # reward, even for the initial representation network where we already
            # know this reward.
            reward = reward_matrix(game)[current_index, player]
            value += value_matrix(game, current_index, td_steps=td_steps)[1][next_player]
            try
                push!(targets, Target(value, reward, policy_matrix(game)[current_index]))
            catch e
                @show game.uuid, current_index, policy_matrix(game)
                throw(e)
            end
        else
            # States past the end of games are treated as absorbing states.
            push!(targets, Target(0, 0, uniform_policy(game)))
        end
    end
    targets
end

function policy_matrix(g::MuZeroGame)
    [g.child_visits..., uniform_policy(g)]
end

function Base.show(io::IO, m::MIME"text/plain", g::MuZeroGame)
    show(io, "text/plain", g.game)
    println(io, "Root Values: $(join(g.root_values, ", "))")
    action_visits = zip(g.game.history, g.child_visits)
    child_visits_str = join(["[$(join(child_visits, ", "))] -> A($action)" for (action, child_visits) in action_visits], "\n")
    println(io, "Child Visits:\n$child_visits_str")
end

# Base.copy(g::AbstractBoardGame) = copy_at(g, length(g.states))

# function copy_at(g::AbstractBoardGame, state_idx::Integer)
#     typeof(g)(
#         copy(g.history[1:(state_idx-1)]),
#         copy(g.states[1:state_idx]),
#         copy(g.rewards[1:(state_idx-1)]),
#         # @todo child_visits and root_values should not be 
#         # copied for AbstractEnv
#         copy(g.child_visits[1:(state_idx-1)]),
#         copy(g.root_values[1:(state_idx-1)]),
#     )
# end