using MuZero

#=
@boardgame Tac (3,3)
@actions 1:9
@pieces [
    (name=:X, icon=(color="black")),
    (name=:O, icon=(color="red"))
]
=#

struct TacState <: AbstractBoardState
    board_state::BoardState{9}
    to_play::Integer
end

TacState() = TacState(BoardState{9}(), PLAYER_1)
TacState(encoded::String) = decode(TacState, encoded)
TacState(::Vector{Nothing}, to_play::Integer) = TacState()
TacState(pieces::Vector{Piece}, to_play::Integer) = init_state(TacState, pieces, to_play)
TacState(pieces::Vector{Union{Nothing, Piece}}, to_play::Integer) = init_state(TacState, pieces, to_play)

#@todo change this to player "A" and "B"
const PLAYER_1_PIECES = [Piece(:X, "X", icon_symbol="X", icon_color="black", icon_size=60)]
const PLAYER_2_PIECES = [Piece(:O, "O", icon_symbol="O", icon_color="red", icon_size=60)]
const ALL_PIECES = [PLAYER_1_PIECES..., PLAYER_2_PIECES...]

MuZero.get_pieces(::Type{TacState}) = ALL_PIECES
MuZero.num_rows(::Type{TacState}) = 3
MuZero.num_cols(::Type{TacState}) = 3
MuZero.image_scale(::Type{TacState}) = 750
MuZero.image_title_offset(::Type{TacState}) = (75,45)
MuZero.actions(::Type{TacState}) = collect(1:9)
MuZero.num_players(::Type{TacState}) = 2

struct TacStateSVG
    s::TacState
end

struct Tac <: AbstractBoardGame
    history::Vector{<:Integer}
    states::Vector{TacState}
    rewards::Vector{Float32}
    function Tac(history, states, rewards)
        @assert length(history) == length(rewards)
        @assert length(states) == (length(history) + 1)
        new(history, states, rewards)
    end
end

function Tac()
    Tac(
        Vector{Int8}(),
        [TacState()],
        Vector{Float32}(),
    )
end

Tac(actions::Vector{<:Integer}) = init_game(Tac, actions)

MuZero.getstatetype(::Type{Tac}) = TacState

################################################

# @todo change to get_coord(::TacState, ::AbstractAction)
function MuZero.get_coord(s::TacState, action::Integer)
    action_to_coord = (
        a7=(1,1), a8=(1,2), a9=(1,3),
        a4=(2,1), a5=(2,2), a6=(2,3),
        a1=(3,1), a2=(3,2), a3=(3,3))
    action_to_coord[Symbol("a$action")]
end

# @todo simplify this to get_marker(::Tac, ::AbstractAction)
# @todo break out into get_piece(::Tac, ::AbstractAction)
function MuZero.get_marker(s::TacState, player::Integer, action::Integer)
    row_idx, col_idx = get_coord(s, action)
    name = player == PLAYER_1 ? :X : :O
    piece = get_piece(s, name)
    Marker(row_idx, col_idx, piece)
end

function MuZero.get_winner(s::TacState)
    for player in [PLAYER_1, PLAYER_2]
        target = player * ones(UInt8, 3)
        for row_idx in 1:num_rows(s)
            if s[row_idx, 1:3] == target
                return player
            end
        end
        for col_idx in 1:num_cols(s)
            if s[1:3, col_idx] == target
                return player
            end
        end
        if s[1, 1] == s[2, 2] == s[3, 3] == player
            return player
        end
        if s[3, 1] == s[2, 2] == s[1, 3] == player
            return player
        end
    end
    nothing
end

function MuZero.legal_actions(s::TacState)
    if game_is_won(s)
        return []
    end
    filter(a->begin
        coord = get_coord(s, a)
        iszero(s[coord...])
    end, actions(s))
end

function MuZero.getconfig(::Type{Tac})
    config = make_board_game_config(
        Tac, 
        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,
    )
    config.batch_size = 10
    config.num_simulations = 10
    config.training_steps = 12000
    config
end