using DataStructures: OrderedDict
using Flux, CuArrays
export Game, showAll, validMoves, randomPlace, place, getGameState, int, PLAYERS

## Game setting
# board size
BOARD_SIZE = Int8(9)
ACTION_SIZE = Int8(9)
HISTORY = Int8(4)

# trainsformation of matrix mapping
ALL_MOVES = (
    original_dont_move  = Int8[1, 2, 3, 4, 5, 6, 7, 8, 9],
    original_dim1_flip  = Int8[7, 8, 9, 4, 5, 6, 1, 2, 3],
    right_rotate_1x90°  = Int8[7, 4, 1, 8, 5, 2, 9, 6, 3],
    right_rot1x90_flip  = Int8[9, 6, 3, 8, 5, 2, 7, 4, 1],
    right_rotate_2x90°  = Int8[9, 8, 7, 6, 5, 4, 3, 2, 1],
    right_rot2x90_flip  = Int8[3, 2, 1, 6, 5, 4, 9, 8, 7],
    right_rotate_3x90°  = Int8[3, 6, 9, 2, 5, 8, 1, 4, 7],
    right_rot3x90_flip  = Int8[1, 4, 7, 2, 5, 8, 3, 6, 9]
)

# transition from players
PLAYERS = Dict{Union{String, Int8}, Union{String, Int8}}("O"=>Int8(1), Int8(1)=>"O", "X"=>Int8(-1), Int8(-1)=>"X", zero(Int8)=>".")
# NEXT_TURN = Dict{Union{String, Int8}, Union{String, Int8}}("O"=>"X", Int8(1)=>Int8(-1), "X"=>"O", Int8(-1)=>Int8(1))

# The state we need to evaluate winning situtaions
STATE = Array{Array{Int8,1}}([
        [1, 2, 3], [4, 5, 6], [7, 8, 9],
        [1, 4, 7], [2, 5, 8], [3, 6, 9],
        [1, 5, 9], [3, 5, 7]
        ])

# give an action, get all possible winning situations
SLOT = Array{Array{Array{Int8,1},1},1}([
        [[1, 2, 3], [1, 4, 7], [1, 5, 9]],
        [[1, 2, 3], [2, 5, 8]],
        [[1, 2, 3], [3, 5, 7], [3, 6, 9]],
        [[1, 4, 7], [4, 5, 6]],
        [[4, 5, 6], [2, 5, 8], [1, 5, 9], [3, 5, 7]],
        [[4, 5, 6], [3, 6, 9]],
        [[1, 4, 7], [7, 8, 9], [3, 5, 7]],
        [[7, 8, 9], [2, 5, 8]],
        [[7, 8, 9], [3, 6, 9], [1, 5, 9]]
        ])

# offset move -> actual board index
OFFSET = Dict{Tuple{Int8,Int8},Tuple{Int8,Int8}}(
  (3, 5) => (2, 8), (3, 6) => (2, 9), (6, 9) => (6, 9), (4, 4) => (5, 1), (9, 4) => (8, 7), (9, 1) => (7, 7), (8, 9) => (9, 6), (3, 1) => (1, 7), (4, 5) => (5, 2),
  (9, 9) => (9, 9), (2, 4) => (2, 4), (6, 5) => (5, 8), (4, 9) => (6, 3), (8, 4) => (8, 4), (1, 2) => (1, 2), (9, 2) => (7, 8), (5, 1) => (4, 4), (9, 8) => (9, 8),
  (3, 4) => (2, 7), (9, 6) => (8, 9), (8, 2) => (7, 5), (2, 3) => (1, 6), (2, 6) => (2, 6), (5, 4) => (5, 4), (2, 1) => (1, 4), (2, 5) => (2, 5), (8, 5) => (8, 5),
  (5, 9) => (6, 6), (1, 4) => (2, 1), (3, 9) => (3, 9), (2, 8) => (3, 5), (5, 8) => (6, 5), (4, 8) => (6, 2), (7, 8) => (9, 2), (1, 1) => (1, 1), (7, 7) => (9, 1),
  (7, 9) => (9, 3), (5, 7) => (6, 4), (7, 2) => (7, 2), (2, 2) => (1, 5), (6, 3) => (4, 9), (1, 7) => (3, 1), (4, 7) => (6, 1), (6, 1) => (4, 7), (6, 8) => (6, 8),
  (8, 8) => (9, 5), (7, 1) => (7, 1), (1, 3) => (1, 3), (7, 5) => (8, 2), (2, 9) => (3, 6), (3, 2) => (1, 8), (7, 6) => (8, 3), (9, 3) => (7, 9), (3, 7) => (3, 7),
  (6, 6) => (5, 9), (5, 5) => (5, 5), (2, 7) => (3, 4), (4, 1) => (4, 1), (8, 3) => (7, 6), (4, 2) => (4, 2), (1, 6) => (2, 3), (4, 3) => (4, 3), (9, 7) => (9, 7),
  (5, 6) => (5, 6), (8, 6) => (8, 6), (4, 6) => (5, 3), (8, 1) => (7, 4), (3, 3) => (1, 9), (6, 2) => (4, 8), (7, 4) => (8, 1), (1, 9) => (3, 3), (8, 7) => (9, 4),
  (6, 4) => (5, 7), (5, 2) => (4, 5), (1, 8) => (3, 2), (6, 7) => (6, 7), (7, 3) => (7, 3), (5, 3) => (4, 6), (3, 8) => (3, 8), (1, 5) => (2, 2), (9, 5) => (8, 8)
)

int(x::Union{Char, String}, T=Int) = parse(T, x)
convertState(state) = [int(i, Int8) for i in state]


mutable struct Game
    boards::AbstractArray{Float32,3}
    vmoves::Dict{Int8, AbstractArray}
    nextPlayer::Int8
    steps::Int8
    winner::Int8
    state::String
    function Game()
        boards = zeros(Float32, (BOARD_SIZE, BOARD_SIZE, 2HISTORY+1))
        vmoves = Dict(Int8(1)=>trues(9), Int8(-1)=>trues(9), Int8(0)=>zeros(Int8, 9,9))
        new(boards, vmoves, Int8(-1),  Int8(1), Int8(0), "")
    end
end

function nozero(x::AbstractArray)
    for i in x
        i==0 && return false
    end
    return true
end

@inline function rmInvalid!(self::Game, x::Integer, y::Integer)
    if @inbounds self.vmoves[0][x,y] != 0
        self.winner = -self.nextPlayer
        return
    end
    @inbounds self.vmoves[0][x,y] = self.nextPlayer

    @inbounds for triple in SLOT[y]
        @inbounds code = sum(self.vmoves[0][x, triple])
        if  code == 2self.nextPlayer
            @inbounds self.vmoves[-self.nextPlayer][x] = false
        end
        if @inbounds code ==-self.nextPlayer && nozero(self.vmoves[0][x, triple])
            count = 0
            for t in STATE
                @inbounds count+= (sum(self.vmoves[0][x, t]) == -2self.nextPlayer)
                count>0 && break
            end
            count >0 && continue
            @inbounds self.vmoves[self.nextPlayer][x] = true
        end
        if @inbounds code==3self.nextPlayer
            self.winner = self.nextPlayer
            return
        end
    end
    return
end


function Game(state::String)
    self = Game()
    isempty(state) && return self
    moves = convertState(state)
    num_moves =  length(moves)
    subboard = moves[1]
    update_bound = max(2, num_moves - HISTORY+1)
    upp = max(1, HISTORY+2-num_moves)
    @sync for i=2:num_moves
        @inbounds rmInvalid!(self, subboard, moves[i])
        @inbounds x,y = OFFSET[(subboard, moves[i])]
        @inbounds self.boards[x, y, upp:HISTORY] .= self.nextPlayer
        @inbounds self.boards[x, y, HISTORY+1:end-(upp-1)] .= -self.nextPlayer
        upp += (i ≥ update_bound);
        self.nextPlayer *= -1;
        @inbounds subboard = moves[i]
    end
    self.steps = num_moves
    @inbounds self.boards[:,:,HISTORY+1] .= self.nextPlayer
    self.state = state
    return self
end


function Base.show(io::IO, self::Game)
    "print out the board"
    compact = get(io, :compact, false)
    if !compact
        board = self.boards[:,:,HISTORY]
        for i=1:9
            print(io, PLAYERS[board[i,1]]," ", PLAYERS[board[i,2]]," ",PLAYERS[board[i,3]]," | ")
            print(io, PLAYERS[board[i,4]]," ", PLAYERS[board[i,5]]," ",PLAYERS[board[i,6]]," | ")
            print(io, PLAYERS[board[i,7]]," ", PLAYERS[board[i,8]]," ",PLAYERS[board[i,9]],"\n")
            if i in (3, 6)
                print(io, "------+-------+------\n")
            end
        end
        showState = isempty(self.state) ? "0" : self.state
        print(io, "\nSubboard: ",showState[end],"; next Player: ", PLAYERS[self.nextPlayer], ";\n")
        print(io, "\nGame States: ",showState,"; Steps: ",self.steps-1,"; ")
        iszero(self.winner) ? print(io, "Game continue...\n") : print(io, PLAYERS[self.winner], " wins\n")
    end
end

function checkWin(self::Game, move::Integer)
    @inbounds x = int(self.state[end])
    @inbounds for triple in SLOT[move]
        @inbounds if sum(self.vmoves[0][x, triple]) == 2self.nextPlayer
            return true
        end
    end
    return false
end

function validMoves(self::Game)
    "Returns all the legal moves for current subboard"
    isempty(self.state) && return keys(OFFSET)
    @inbounds x = int(self.state[end])
    @inbounds available_moves = Array{Int8}(findall(iszero, self.vmoves[0][x,:]))
    if @inbounds !self.vmoves[-self.nextPlayer][x]
        return filter(available_moves) do move
            checkWin(self, move)
        end
    end
    validmoves = filter(available_moves) do move
        @inbounds self.vmoves[self.nextPlayer][move]
    end
    @inbounds if isempty(validmoves)
        return isempty(available_moves) ? [Int8(1)] : available_moves
    end
    return validmoves
end


function place(self::Game, move::Union{Tuple, Integer})
    !iszero(self.winner) && return false
    if isempty(self.state)
        @inbounds rmInvalid!(self, move[1], move[2])
        @inbounds self.state = "$(move[1])$(move[2])"
    else
        @inbounds rmInvalid!(self, int(self.state[end]), move)
        x, y = OFFSET[(int(self.state[end]), move)]
        if self.boards[x, y, HISTORY] != 0
            self.winner = -self.nextPlayer
            return false
        end
        self.state = "$(self.state)$move"
    end
    num_moves = length(self.state)
    iters = max(1, HISTORY-num_moves+2)
    player = self.nextPlayer
    for i=HISTORY:-1:iters
        @inbounds subboard = int(self.state[end+i-HISTORY-1])
        @inbounds x, y = OFFSET[(subboard, int(self.state[end+i-HISTORY]))]
        @inbounds self.boards[x, y, i] = player
        @inbounds self.boards[x, y, end-i+1] = -player
        player *= -1
    end
    self.steps += 1
    self.nextPlayer *= -1
    @inbounds self.boards[:,:,HISTORY+1] .= self.nextPlayer
    return true
end


function randomPlace(self::Game, times::Integer=1)
    for i = 1:times
        !iszero(self.winner) && break
        moves = validMoves(self)
        place(self, rand(moves))
    end
end


function conv(state::String, typ::Int)
    result = ""
    for (i,e) in enumerate(state)
        result *= string(ALL_MOVES[typ][int(e)])
    end
    return result
end


function getGameState(self::Game)
    return reshape(self.boards, (BOARD_SIZE, BOARD_SIZE, 2HISTORY+1,1))
end


function getAugmentation(self::Game, p::Array{Float32, 1})
    """ The order is
        {game.state => M1, flip1 => M2,
         rotr1x90   => M3, flip2 => M4
         rotr2x90   => M5, flip3 => M6
         rotr3x90   => M7, flip4 => M8}
    """
    # strings = String[];
    # push!(strings, conv(self.state, 2))
    data = Array{Array{Float32, 3}, Array{Float32, 1}, Integer}[]
    gameState1 = self.boards
    push!(data, [gameState1, p, game.nextPlayer])
    # flip along with dim1
    # states = cat(dims=4, reverse(gameState1, dims=1))
    @inbounds push!(data, [reverse(gameState1, dims=1), p[ALL_MOVES[2]], game.nextPlayer])
    for i=3:2:8
        # right rotate 90°
        # push!(strings, conv(self.state, i))
        gameState1 = mapslices(rotr90, gameState1, dims=[1,2])
        @inbounds push!(data, [gameState1, p[ALL_MOVES[i]], game.nextPlayer])
        # states = cat(dims=4, states, gameState1)
        # flip along with dim1
        # push!(strings, conv(self.state, i+1))
        @inbounds push!(data, [reverse(gameState1, dims=1), p[ALL_MOVES[i]], game.nextPlayer])
        # states = cat(dims=4, states, reverse(gameState1, dims=1))
    end
    return data
end
