using Printf
import Flux
import GridGames

OVERRIDE_MSG = "This should be overridden in concrete implemention of AbstractEnv"

get_marker(::AbstractEnv, ::Integer, ::Integer)   = @assert false OVERRIDE_MSG
get_coord(::AbstractEnv, action::Integer)         = @assert false OVERRIDE_MSG
getstates(::Type{<:AbstractEnv})                  = @assert false OVERRIDE_MSG
getstatesplus(::Type{<:AbstractEnv})              = @assert false OVERRIDE_MSG

makepolicytests(::Type{<:AbstractEnv}) = nothing
legal_actions(g::AbstractEnv) = legal_actions(get_state(g))
svg(s::AbstractEnv) = BoardGameSVG(s)
get_board_state(x::AbstractEnv) = get_state(x).board_state
get_state(x::AbstractEnv) = x.states[end]
get_state(x::AbstractEnv, idx::Integer) = x.states[idx]
num_rows(T::Type{<:AbstractEnv}) = num_rows(getstatetype(T))
num_rows(g::AbstractEnv) = num_rows(typeof(g))
num_cols(T::Type{<:AbstractEnv}) = num_cols(getstatetype(T))
num_cols(g::AbstractEnv) = num_cols(typeof(g))
num_cells(T::Type{<:AbstractEnv}) = num_cells(getstatetype(T))
num_cells(g::AbstractEnv) = num_cells(typeof(g))
image_scale(x::AbstractEnv) = image_scale(getstatetype(x))
image_title_offset(x::AbstractEnv) = image_title_offset(getstatetype(x))
player1_markers(x::AbstractEnv) = player1_markers(getstatetype(x))
player2_markers(x::AbstractEnv) = player2_markers(getstatetype(x))
getstatetype(x::AbstractEnv) = GridGames.getstatetype(typeof(x))
to_play(x::AbstractEnv) = to_play(x, length(x.history) + 1)
to_play(x::AbstractEnv, state_idx::Integer) = isodd(state_idx) ? PLAYER_1 : PLAYER_2
get_pieces(g::AbstractEnv) = get_pieces(getstatetype(g))
get_pieces(T::Type{<:AbstractEnv}) = get_pieces(getstatetype(T))
action_space_size(g::AbstractEnv) = action_space_size(getstatetype(g))
action_space_size(T::Type{<:AbstractEnv}) = action_space_size(GridGames.getstatetype(T))
Base.:-(g1::AbstractEnv, g2::AbstractEnv) = get_state(g1) - get_state(g2)
piece(g::AbstractEnv, name::Symbol) = piece(getstatetype(s), name)
piece(T::Type{<:AbstractEnv}, name::Symbol) = piece(getstatetype(T), name)
init_state(g::AbstractEnv, pieces::Vector{Piece}) = init_state(getstatetype(g), pieces)
init_state(T::Type{<:AbstractEnv}, pieces::Vector{Piece}) = init_state(getstatetype(T), pieces)

function MuZero.policypredsvg(g::AbstractEnv, test_result::PolicyTestResult)
    test = test_result.test
    policy = test_result.prediction
    g = test.game
    sorted_actions = sortperm(policy)
    top3_actions = reverse(sorted_actions)[1:3]
    top3 = map(x->begin
        pr_idx, action = x
        pred_pr = @sprintf "%.2f" policy[action]
        g_copy = copy(g)
        is_max = in(action, test.actions)
        step!(g_copy, action)
        highlight_coord = diff_coord(g_copy, g)
        highlight_color = is_max ? "#AADBC0" : "#ECABA7"
        gridsvg(
            get_state(g_copy);
            title="prediction #$pr_idx",
            title_offset=(-210,-195),
            tr_header="pr=$(pred_pr)",
            tr_offset=(-310,-295),
            highlight=highlight_coord,
            highlight_color=highlight_color,
        )
    end, enumerate(top3_actions))

    GridSVG(top3)
end

function diff_coord(g1::AbstractEnv, g2::AbstractEnv)
    diff = (g1 - g2)
    isempty(diff) ? nothing : diff[1]
end

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

function copy_at(g::AbstractEnv, state_idx::Integer)
    typeof(g)(
        copy(gethistory(g)[1:(state_idx-1)]),
        copy(g.rewards[1:(state_idx-1)]),
        copy(g.states[1:state_idx]),
    )
end

# @todo change this to get_image(::AbstractEnv, ::Action)
get_action_image(g::AbstractEnv, a::Integer) = get_action_image(GridGames.getstatetype(g), a)
get_action_image(T::Type{<:AbstractEnv}, a::Integer) = get_action_image(GridGames.getstatetype(T), a)
get_action_image(T::Type{<:AbstractEnv}, a::AbstractEnvAction) = get_action_image(T, a.idx)

function get_image(g::AbstractEnv, state_idx=length(g.states))
    @assert 0 < state_idx <= length(g.states) "state_idx is $state_idx, but the number of states is $(length(g.states))"
    get_image(get_state(g, state_idx))
end

get_reward(g::AbstractEnv, action::Integer) = get_reward(get_state(g), action)

function init_game(T::Type{<:AbstractEnv}, actions::Vector{<:Integer})
    g = T()
    foreach(a->step!(g,a), actions)
    g
end

function coord_to_byte_idx(T::Type{<:AbstractEnv}, coord::Tuple{<:Integer, <:Integer})
    coord_to_byte_idx(getstatetype(T), coord)
end

function get_piece(g::AbstractEnv, name::Symbol)
    get_pieces(g)[get_piece_idx(g, name)]
end

get_piece_idx(g::AbstractEnv, name::Symbol) = get_piece_idx(getstatetype(g), name)


# @todo remove all dependencies on "apply!" (and use step! instead)
apply!(g::AbstractEnv, a::Integer) = step!(g, a)
isterminal(g::AbstractEnv) = length(getlegalactions(g)) == 0
islegal(g::AbstractEnv, action::Integer) = in(action, legal_actions(g))
game_is_won(g::AbstractEnv) = game_is_won(get_state(g))
get_winner(g::AbstractEnv) = get_state(g) |> get_winner
#next_to_play(g::AbstractEnv) = to_play(g) == PLAYER_1 ? PLAYER_2 : PLAYER_1
discount(::AbstractEnv) = Float32(1)

# @todo eventually remove these
is_terminal(g::AbstractEnv) = isterminal(g)
current_image(g::AbstractEnv) = get_image(g)
make_image(g::AbstractEnv, i::Integer) = get_image(g, i)
to_play_at(g::AbstractEnv, i::Integer) = to_play(g, i)
function make_image_label(g::AbstractEnv, image_idx::Integer)
    if 0 < image_idx <= length(g.states)
        actions = gethistory(g)[1:(image_idx-1)]
        return join(actions, ",")
        players = player_history(g)[1:(image_idx-1)]
        return join(map(pair->join(pair), zip(players, actions)), ",")
    end
    @assert false "image_idx is $image_idx, but the number of states is $(length(g.states))"
end
function player_history(g::AbstractEnv)
    players = ["A","B"]
    map(idx->players[(idx-1)%2+1], 1:length(g.history))
end
function reward_matrix(g::AbstractEnv)
    rewards_with_initial_zero = [0., g.rewards...]
    m = zeros(Float32, length(rewards_with_initial_zero), 2)
    for (idx,reward) in enumerate(rewards_with_initial_zero)
        m[idx, 1] = iseven(idx) ? reward : -1*reward
        m[idx, 2] = iseven(idx) ? -1*reward : reward
    end
    m
end
function value_matrix(g::AbstractEnv, from::Integer=1; td_steps)
    to = min(length(gethistory(g)) + 1, from + td_steps)
    value_matrix_support(g, from, to)
end
function value_matrix_support(g::AbstractEnv, from::Integer=1, to=(length(g.history)+1))
    reward_from = from + 1
    reward_to = to
    rewards = reward_matrix(g)[reward_from:reward_to, :]
    function cumvalue(next_value, row)
        i, reward = row
        reward * discount(g)^i + next_value
    end
    mat = accumulate(
        cumvalue, 
        enumerate(eachrow(reverse(rewards; dims=1))); 
        init=[0., 0.]) |> reverse
    [mat..., [0.,0.]]
end

# function step!(g::AbstractEnv, actions::Vector{<:Integer})
#     map(a->step!(g, a), actions)
# end

# function step!(g::AbstractEnv, action::Integer)
#     if !isterminal(g) && islegal(g, action)
#         push!(g.rewards, get_reward(g, action))
#         push!(g.states, next_state(g, action))
#         push!(g.history, action)
#         return g.rewards[end]
#     end
#     nothing
# end

next_state(g::AbstractEnv, action::Integer) = next_state(get_state(g), action)

function Base.getindex(g::AbstractEnv, row_idx::Integer, col_idx::Integer)
    get_state(g)[row_idx, col_idx]
end

function Base.getindex(g::AbstractEnv, ::Colon, col_idx::Integer)
    get_state(g)[:, col_idx]
end

function Base.getindex(g::AbstractEnv, row_idx::Integer, ::Colon)
    get_state(g)[row_idx, :]
end

function Base.getindex(g::AbstractEnv, row_idx::Integer, col_range::UnitRange{<:Integer})
    get_state(g)[row_idx, col_range]
end

function Base.getindex(g::AbstractEnv, row_range::UnitRange{<:Integer}, col_idx::Integer)
    get_state(g)[row_range, col_idx]
end

function Base.show(io::IO, g::AbstractEnv)
    game_type_name = string(typeof(g))
    state_type_name = string(getstatetype(g))
    print(io, "$game_type_name{$state_type_name{\"$(encode(get_state(g)))\"}}")
end

# function Base.show(io::IO, m::MIME"text/plain", g::AbstractEnv)
#     show(io, "text/plain", get_state(g))
#     println(io)
#     if game_is_won(g)
#         println(io, "Winner: Player $(get_winner(g))")
#     else
#         println(io, "To Play: $(to_play(g))")
#     end
#     println(io, "History: $(join(g.history, ", "))")
# end

byte_idx_to_coord(g::AbstractEnv, byte_idx::Integer) = byte_idx_to_coord(getstatetype(g), byte_idx)
byte_idx_to_coord(T::Type{<:AbstractEnv}, byte_idx::Integer) = byte_idx_to_coord(getstatetype(T), byte_idx)

# function gridsvg(g::AbstractEnv; main_title="image", next_action=nothing, header=nothing, tr_header=nothing)
#     function cell_fn(_,_,row,col)
#         piece_idx = s[row,col]
#         if piece_idx == 0
#             return nothing
#         end
#         m = all_markers(s)[piece_idx]
#         (value=m.icon_symbol, color=m.icon_color, size=m.icon_size)
#     end
#     GridSVG(
#         Grid(num_cols(s), num_rows(s), cell_fn), 
#         main_title=main_title,
#         main_title_offset=image_title_offset(s),
#         tr_header=tr_header,
#         header=header,
#         aspect_ratio=(num_cols(s), num_rows(s)),
#         scale=image_scale(s))
# end

function Base.show(io::IO, m::MIME"image/svg+xml", g::AbstractBoardGameSVG)
    show(io, m, gridsvg(g.x))
end

actions(g::AbstractEnv) = actions(getstatetype(g))
actions(T::Type{<:AbstractEnv}) = actions(getstatetype(T))
uniformpolicy(T::Type{<:AbstractEnv}) = (s)->rand(actions(T))

function get_boardgame_config(game_type::Type{<:AbstractEnv}; 
    network_type, action_space_size, max_moves, root_dirichlet_alpha, 
    lr_init, support_size, observation_shape, 
    stacked_observations, hidden_state_size, checkpoint_interval,
    num_actors, weight_decay, save_interval, num_unroll_steps)
    
    function visit_softmax_temperature(num_moves, num_steps)
        if num_moves < 30
            return 1.0
        end
        0.0  # Play according to the max.
    end

    tictactoe_visit_softmax_temperature(num_moves, num_steps) = 1.0
    
    return MuZeroConfig(
        game_type,
        network_type=network_type,
        action_space_size=action_space_size,
        max_moves=max_moves,
        discount=1.0,
        root_dirichlet_alpha=root_dirichlet_alpha,
        num_simulations=100,
        training_steps=2000000,
        batch_size=64,
        td_steps=max_moves, # Always use Monte Carlo return.
        num_actors=num_actors,
        lr_init=lr_init,
        lr_decay_steps=400e3,
        visit_softmax_temperature_fn=tictactoe_visit_softmax_temperature,
        known_bounds=KnownBounds(-1, 1),
        support_size=support_size,
        observation_shape=observation_shape,
        stacked_observations=stacked_observations,
        hidden_state_size=hidden_state_size,
        checkpoint_interval=checkpoint_interval,
        weight_decay=weight_decay, 
        save_interval=save_interval,
        num_unroll_steps=num_unroll_steps,
    )     
end

function getconfig(g::Type{<:AbstractEnv})
    get_boardgame_config(
        g, 
        network_type=Incremental,
        max_moves=9, 
        root_dirichlet_alpha=0.3, 
        action_space_size=9, # Number of all possible actions in the game 
        support_size = 10, # Value and reward are scaled (with almost sqrt) and encoded on a vector with a range of -support_size to support_size. Choose it so that support_size <= sqrt(max(abs(discounted reward)))
        observation_shape = (3, 3, 3),
        stacked_observations = 0, # Number of previous observations and previous actions to add to the current observation
        hidden_state_size = 32, # The 1-dimensional size of the encoded game state
        checkpoint_interval = 100, #  Number of training steps before using the model for self-playing
        num_actors = 3,
        weight_decay = 1e-3,
        lr_init=0.002f0,
        save_interval = 10000,
        num_unroll_steps=1,
    )
end

# function Base.show(io::IO, m::MIME"image/svg+xml", g::AbstractEnvSVG)
#     game = g.x
#     svgs = map(action_idx->gridsvg(game, action_idx), 2:length(game.history))
#     fig = GridSVG(svgs, :, 2;
#         main_title="Game", 
#         main_title_offset=(70,40),
#         main_title_size=40)
#     show(io, m, fig)
# end

function gridsvg(g::AbstractEnv)
    title_offset = (-210, -200)
    title_size = 30

    svgs = map(x->begin
        (action_idx, state) = x
        curr = state
        prev = g.states[action_idx]
        highlight_coord = diff_coord(curr, prev)
        gridsvg(
            state;
            highlight=highlight_coord,
            title="action $action_idx", 
            title_offset=title_offset,
            title_size=title_size)
    end, enumerate(g.states[2:end]))

    if isempty(svgs)
        push!(svgs, 
            gridsvg(
                g.states[1];
                title="initial state", 
                title_offset=title_offset,
                title_size=title_size
            ))
    end

    fig = GridSVG(svgs, :, 2;
        main_title="Game", 
        main_title_offset=(70,40),
        main_title_size=40)


    # state = get_state(g, action_idx)
    # player = to_play(g, action_idx)
    # next_action = g.history[action_idx]
    # gridsvg(state;
    #     title="action $action_idx", 
    #     title_offset=(-210, -200),
    #     title_size=30)
end
