import sys
import pygame
pygame.init()
display = pygame.display.set_mode((800,600))
clock = pygame.time.Clock()

while True:
	for event in pygame.event.get():
		if event.type == pygame.QUIT:
			sys.exit(0)
		else:
			pygame.display.update()
			clock.tick(1)

class PyGameBoard:

	def __init__(self, connectNGame: ConnectNGame):
		self.connectNGame = connectNGame
		pygame.init()

	def next_user_input(self) -> Tuple[int, int]:
		self.action = None
		while not self.action:
			self.check_event()
			self._render()
			self.clock.tick(60)
		return self.action
  
  def move(self, r: int, c: int) -> int:
		return self.connectNGame.move(r, c)
  
if __name__ == '__main__':
	connectNGame = ConnectNGame()
	pygameBoard = PyGameBoard(connectNGame)
	while not pygameBoard.isGameOver():
		pos = pygameBoard.next_user_input()
		pygameBoard.move(*pos)

	pygame.quit()

def check_event(self):
	for e in pygame.event.get():
		if e.type == pygame.QUIT:
			pygame.quit()
			sys.exit(0)
		elif e.type == pygame.MOUSEBUTTONDOWN:
			self._handle_user_input(e)
    
def _handle_user_input(self, e: Event) -> Tuple[int, int]:
	origin_x = self.start_x - self.edge_size
	origin_y = self.start_y - self.edge_size
	size = (self.board_size - 1) * self.grid_size + self.edge_size * 2
	pos = e.pos
	if origin_x <= pos[0] <= origin_x + size and origin_y <= pos[1] <= origin_y + size:
		if not self.connectNGame.gameOver:
			x = pos[0] - origin_x
			y = pos[1] - origin_y
			r = int(y // self.grid_size)
			c = int(x // self.grid_size)
			valid = self.connectNGame.checkAction(r, c)
			if valid:
				self.action = (r, c)
				return self.action

class ConnectNGym(gym.Env):

	def reset(self) -> ConnectNGame:
		"""Resets the state of the environment and returns an initial observation.

		Returns:
			observation (object): the initial observation.
		"""
		raise NotImplementedError


	def step(self, action: Tuple[int, int]) -> Tuple[ConnectNGame, int, bool, None]:
		"""Run one timestep of the environment's dynamics. When end of
		episode is reached, you are responsible for calling `reset()`
		to reset this environment's state.

		Accepts an action and returns a tuple (observation, reward, done, info).

		Args:
			action (object): an action provided by the agent

		Returns:
			observation (object): agent's observation of the current environment
			reward (float) : amount of reward returned after previous action
			done (bool): whether the episode has ended, in which case further step() calls will return undefined results
			info (dict): contains auxiliary diagnostic information (helpful for debugging, and sometimes learning)
		"""
		raise NotImplementedError



	def render(self, mode='human'):
		"""
		Renders the environment.

		The set of supported modes varies per environment. (And some
		environments do not support rendering at all.) By convention,
		if mode is:

		- human: render to the current display or terminal and
			return nothing. Usually for human consumption.
		- rgb_array: Return an numpy.ndarray with shape (x, y, 3),
			representing RGB values for an x-by-y pixel image, suitable
			for turning into a video.
		- ansi: Return a string (str) or StringIO.StringIO containing a
			terminal-style text representation. The text can include newlines
			and ANSI escape sequences (e.g. for colors).

		Note:
		Make sure that your class's metadata 'render.modes' key includes
		the list of supported modes. It's recommended to call super()
		in implementations to use the functionality of this method.

		Args:
			mode (str): the mode to render with
		"""
		raise NotImplementedError

def reset(self) -> ConnectNGame

def step(self, action: Tuple[int, int]) -> Tuple[ConnectNGame, int, bool, None]

def render(self, mode='human')

class ConnectNGym(gym.Env):

	def __init__(self, pygameBoard: PyGameBoard, isGUI=True, displaySec=2):
		self.pygameBoard = pygameBoard
		self.isGUI = isGUI
		self.displaySec = displaySec
		self.action_space = spaces.Discrete(pygameBoard.board_size * pygameBoard.board_size)
		self.observation_space = spaces.Discrete(pygameBoard.board_size * pygameBoard.board_size)
		self.seed()
		self.reset()

	def reset(self) -> ConnectNGame:
		self.pygameBoard.connectNGame.reset()
		return copy.deepcopy(self.pygameBoard.connectNGame)

	def step(self, action: Tuple[int, int]) -> Tuple[ConnectNGame, int, bool, None]:
		# assert self.action_space.contains(action)

		r, c = action
		reward = REWARD_NONE
		result = self.pygameBoard.move(r, c)
		if self.pygameBoard.isGameOver():
			reward = result

		return copy.deepcopy(self.pygameBoard.connectNGame), reward, not result is None, None

	def render(self, mode='human'):
		if not self.isGUI:
			self.pygameBoard.connectNGame.drawText()
			time.sleep(self.displaySec)
		else:
			self.pygameBoard.display(sec=self.displaySec)

	def get_available_actions(self) -> List[Tuple[int, int]]:
		return self.pygameBoard.getAvailablePositions()

def similarStatus(self, status: Tuple[Tuple[int, ...]]) -> List[Tuple[Tuple[int, ...]]]:
	ret = []
	rotatedS = status
	for _ in range(4):
		rotatedS = self.rotate(rotatedS)
		ret.append(rotatedS)
	return ret

def rotate(self, status: Tuple[Tuple[int, ...]]) -> Tuple[Tuple[int, ...]]:
	N = len(status)
	board = [[ConnectNGame.AVAILABLE] * N for _ in range(N)]

	for r in range(N):
		for c in range(N):
			board[c][N - 1 - r] = status[r][c]

	return tuple([tuple(board[i]) for i in range(N)])

class PlannedMinimaxStrategy(Strategy):
	def __init__(self, game: ConnectNGame):
		super().__init__()
		self.game = copy.deepcopy(game)
		self.dpMap = {}  # game_status => result, move
		self.result = self.minimax(game.getStatus())


	def action(self, game: ConnectNGame) -> Tuple[int, Tuple[int, int]]:
		game = copy.deepcopy(game)

		player = game.currentPlayer
		bestResult = player * -1  # assume opponent win as worst result
		bestMove = None
		for move in game.getAvailablePositions():
			game.move(*move)
			status = game.getStatus()
			game.undo()

			result = self.dpMap[status]

			if player == ConnectNGame.PLAYER_A:
				bestResult = max(bestResult, result)
			else:
				bestResult = min(bestResult, result)
			# update bestMove if any improvement
			bestMove = move if bestResult == result else bestMove
			print(f'move {move} => {result}')

		return bestResult, bestMove

class BaseAgent(object):
	def __init__(self):
		pass

	def act(self, game: PyGameBoard, available_actions):
		return random.choice(available_actions)

class AIAgent(BaseAgent):
	def __init__(self, strategy: Strategy):
		self.strategy = strategy

	def act(self, game: PyGameBoard, available_actions):
		result, move = self.strategy.action(game.connectNGame)
		assert move in available_actions
		return move

class HumanAgent(BaseAgent):
	def __init__(self):
		pass

	def act(self, game: PyGameBoard, available_actions):
		return game.next_user_input()

def play_ai_vs_ai(env: ConnectNGym):
	plannedMinimaxAgent = AIAgent(PlannedMinimaxStrategy(env.pygameBoard.connectNGame))
	play(env, plannedMinimaxAgent, plannedMinimaxAgent)


def play(env: ConnectNGym, agent1: BaseAgent, agent2: BaseAgent):
	agents = [agent1, agent2]

	while True:
		env.reset()
		done = False
		agent_id = -1
		while not done:
			agent_id = (agent_id + 1) % 2
			available_actions = env.get_available_actions()
			agent = agents[agent_id]
			action = agent.act(pygameBoard, available_actions)
			_, reward, done, info = env.step(action)
			env.render(True)

			if done:
				print(f'result={reward}')
				time.sleep(3)
				break


if __name__ == '__main__':
	pygameBoard = PyGameBoard(connectNGame=ConnectNGame(board_size=3, N=3))
	env = ConnectNGym(pygameBoard)
	env.render(True)

	play_ai_vs_ai(env)

# AC
from typing import List

class Solution:
    def tictactoe(self, moves: List[List[int]]) -> str:
        board = [[0] * 3 for _ in range(3)]
        for idx, xy in enumerate(moves):
            player = 1 if idx % 2 == 0 else -1
            board[xy[0]][xy[1]] = player

        turn = 0
        row, col = [0, 0, 0], [0, 0, 0]
        diag1, diag2 = False, False
        for r in range(3):
            for c in range(3):
                turn += board[r][c]
                row[r] += board[r][c]
                col[c] += board[r][c]
                if r == c:
                    diag1 += board[r][c]
                if r + c == 2:
                    diag2 += board[r][c]

        oWin = any(row[r] == 3 for r in range(3)) or any(col[c] == 3 for c in range(3)) or diag1 == 3 or diag2 == 3
        xWin = any(row[r] == -3 for r in range(3)) or any(col[c] == -3 for c in range(3)) or diag1 == -3 or diag2 == -3

        return "A" if oWin else "B" if xWin else "Draw" if len(moves) == 9 else "Pending"

# AC
from typing import List

class Solution:
    def checkWin(self, r: int, c: int) -> bool:
        north = self.getConnectedNum(r, c, -1, 0)
        south = self.getConnectedNum(r, c, 1, 0)

        east = self.getConnectedNum(r, c, 0, 1)
        west = self.getConnectedNum(r, c, 0, -1)

        south_east = self.getConnectedNum(r, c, 1, 1)
        north_west = self.getConnectedNum(r, c, -1, -1)

        north_east = self.getConnectedNum(r, c, -1, 1)
        south_west = self.getConnectedNum(r, c, 1, -1)

        if (north + south + 1 >= 3) or (east + west + 1 >= 3) or \
                (south_east + north_west + 1 >= 3) or (north_east + south_west + 1 >= 3):
            return True
        return False

    def getConnectedNum(self, r: int, c: int, dr: int, dc: int) -> int:
        player = self.board[r][c]
        result = 0
        i = 1
        while True:
            new_r = r + dr * i
            new_c = c + dc * i
            if 0 <= new_r < 3 and 0 <= new_c < 3:
                if self.board[new_r][new_c] == player:
                    result += 1
                else:
                    break
            else:
                break
            i += 1
        return result

    def tictactoe(self, moves: List[List[int]]) -> str:
        self.board = [[0] * 3 for _ in range(3)]
        for idx, xy in enumerate(moves):
            player = 1 if idx % 2 == 0 else -1
            self.board[xy[0]][xy[1]] = player

        # only check last move
        r, c = moves[-1]
        win = self.checkWin(r, c)
        if win:
            return "A" if len(moves) % 2 == 1 else "B"

        return "Draw" if len(moves) == 9 else "Pending"

# AC
from typing import List

class Solution:

    def convertCell(self, c:str):
        return 1 if c == 'X' else -1 if c == 'O' else 0

    def validTicTacToe(self, board: List[str]) -> bool:
        turn = 0
        row, col = [0, 0, 0], [0, 0, 0]
        diag1, diag2 = False, False
        for r in range(3):
            for c in range(3):
                turn += self.convertCell(board[r][c])
                row[r] += self.convertCell(board[r][c])
                col[c] += self.convertCell(board[r][c])
                if r == c:
                    diag1 += self.convertCell(board[r][c])
                if r + c == 2:
                    diag2 += self.convertCell(board[r][c])

        xWin = any(row[r] == 3 for r in range(3)) or any(col[c] == 3 for c in range(3)) or diag1 == 3 or diag2 == 3
        oWin = any(row[r] == -3 for r in range(3)) or any(col[c] == -3 for c in range(3)) or diag1 == -3 or diag2 == -3
        if (xWin and turn == 0) or (oWin and turn == 1):
            return False
        return (turn == 0 or turn == 1) and (not xWin or not oWin)

# AC
class TicTacToe:

    def __init__(self, n:int):
        """
        Initialize your data structure here.
        :type n: int
        """
        self.row, self.col, self.diag1, self.diag2, self.n = [0] * n, [0] * n, 0, 0, n

    def move(self, row:int, col:int, player:int) -> int:
        """
        Player {player} makes a move at ({row}, {col}).
        @param row The row of the board.
        @param col The column of the board.
        @param player The player, can be either 1 or 2.
        @return The current winning condition, can be either:
                0: No one wins.
                1: Player 1 wins.
                2: Player 2 wins.
        """
        if player == 2:
            player = -1

        self.row[row] += player
        self.col[col] += player
        if row == col:
            self.diag1 += player
        if row + col == self.n - 1:
            self.diag2 += player

        if self.n in [self.row[row], self.col[col], self.diag1, self.diag2]:
            return 1
        if -self.n in [self.row[row], self.col[col], self.diag1, self.diag2]:
            return 2
        return 0

class ConnectNGame:

    PLAYER_A = 1
    PLAYER_B = -1
    AVAILABLE = 0
    RESULT_TIE = 0
    RESULT_A_WIN = 1
    RESULT_B_WIN = -1

    def __init__(self, N:int = 3, board_size:int = 3):
        assert N <= board_size
        self.N = N
        self.board_size = board_size
        self.board = [[ConnectNGame.AVAILABLE] * board_size for _ in range(board_size)]
        self.gameOver = False
        self.gameResult = None
        self.currentPlayer = ConnectNGame.PLAYER_A
        self.remainingPosNum = board_size * board_size
        self.actionStack = []

    def move(self, r: int, c: int) -> int:
        """

        :param r:
        :param c:
        :return: None: game ongoing
        """
        assert self.board[r][c] == ConnectNGame.AVAILABLE
        self.board[r][c] = self.currentPlayer
        self.actionStack.append((r, c))
        self.remainingPosNum -= 1
        if self.checkWin(r, c):
            self.gameOver = True
            self.gameResult = self.currentPlayer
            return self.currentPlayer
        if self.remainingPosNum == 0:
            self.gameOver = True
            self.gameResult = ConnectNGame.RESULT_TIE
            return ConnectNGame.RESULT_TIE
        self.currentPlayer *= -1

    def undo(self):
        if len(self.actionStack) > 0:
            lastAction = self.actionStack.pop()
            r, c = lastAction
            self.board[r][c] = ConnectNGame.AVAILABLE
            self.currentPlayer = ConnectNGame.PLAYER_A if len(self.actionStack) % 2 == 0 else ConnectNGame.PLAYER_B
            self.remainingPosNum += 1
            self.gameOver = False
            self.gameResult = None
        else:
            raise Exception('No lastAction')

    def getAvailablePositions(self) -> List[Tuple[int, int]]:
        return [(i,j) for i in range(self.board_size) for j in range(self.board_size) if self.board[i][j] == ConnectNGame.AVAILABLE]

    def getStatus(self) -> Tuple[Tuple[int, ...]]:
        return tuple([tuple(self.board[i]) for i in range(self.board_size)])

class Strategy(ABC):

    def __init__(self):
        super().__init__()

    @abstractmethod
    def action(self, game: ConnectNGame) -> Tuple[int, Tuple[int, int]]:
        pass

class MinimaxStrategy(Strategy):
    def action(self, game: ConnectNGame) -> Tuple[int, Tuple[int, int]]:
        self.game = copy.deepcopy(game)
        result, move = self.minimax()
        return result, move

    def minimax(self) -> Tuple[int, Tuple[int, int]]:
        game = self.game
        bestMove = None
        assert not game.gameOver
        if game.currentPlayer == ConnectNGame.PLAYER_A:
            ret = -math.inf
            for pos in game.getAvailablePositions():
                move = pos
                result = game.move(*pos)
                if result is None:
                    assert not game.gameOver
                    result, oppMove = self.minimax()
                game.undo()
                ret = max(ret, result)
                bestMove = move if ret == result else bestMove
                if ret == 1:
                    return 1, move
            return ret, bestMove
        else:
            ret = math.inf
            for pos in game.getAvailablePositions():
                move = pos
                result = game.move(*pos)
                if result is None:
                    assert not game.gameOver
                    result, oppMove = self.minimax()
                game.undo()
                ret = min(ret, result)
                bestMove = move if ret == result else bestMove
                if ret == -1:
                    return -1, move
            return ret, bestMove

class CountingMinimaxStrategy(Strategy):
    def action(self, game: ConnectNGame) -> Tuple[int, Tuple[int, int]]:
        self.game = copy.deepcopy(game)
        self.dpMap = {}
        result, move = self.minimax(game.getStatus())
        return result, move

    def minimax(self, gameStatus: Tuple[Tuple[int, ...]]) -> Tuple[int, Tuple[int, int]]:
        # print(f'Current {len(strategy.dpMap)}')

        if gameStatus in self.dpMap:
            return self.dpMap[gameStatus]

        game = self.game
        bestMove = None
        assert not game.gameOver
        if game.currentPlayer == ConnectNGame.PLAYER_A:
            ret = -math.inf
            for pos in game.getAvailablePositions():
                move = pos
                result = game.move(*pos)
                if result is None:
                    assert not game.gameOver
                    result, oppMove = self.minimax(game.getStatus())
                    self.dpMap[game.getStatus()] = result, oppMove
                else:
                    self.dpMap[game.getStatus()] = result, move
                game.undo()
                ret = max(ret, result)
                bestMove = move if ret == result else bestMove
            self.dpMap[gameStatus] = ret, bestMove
            return ret, bestMove
        else:
            ret = math.inf
            for pos in game.getAvailablePositions():
                move = pos
                result = game.move(*pos)

                if result is None:
                    assert not game.gameOver
                    result, oppMove = self.minimax(game.getStatus())
                    self.dpMap[game.getStatus()] = result, oppMove
                else:
                    self.dpMap[game.getStatus()] = result, move
                game.undo()
                ret = min(ret, result)
                bestMove = move if ret == result else bestMove
            self.dpMap[gameStatus] = ret, bestMove
            return ret, bestMove


if __name__ == '__main__':
    tic_tac_toe = ConnectNGame(N=3, board_size=3)
    strategy = CountingMinimaxStrategy()
    strategy.action(tic_tac_toe)
    print(f'Game States Number {len(strategy.dpMap)}')

class AlphaBetaStrategy(Strategy):
    def action(self, game: ConnectNGame) -> Tuple[int, Tuple[int, int]]:
        self.game = game
        result, move = self.alpha_beta(self.game.getStatus(), -math.inf, math.inf)
        return result, move

    def alpha_beta(self, gameStatus: Tuple[Tuple[int, ...]], alpha:int=None, beta:int=None) -> Tuple[int, Tuple[int, int]]:
        game = self.game
        bestMove = None
        assert not game.gameOver
        if game.currentPlayer == ConnectNGame.PLAYER_A:
            ret = -math.inf
            for pos in game.getAvailablePositions():
                move = pos
                result = game.move(*pos)
                if result is None:
                    assert not game.gameOver
                    result, oppMove = self.alpha_beta(game.getStatus(), alpha, beta)
                game.undo()
                alpha = max(alpha, result)
                ret = max(ret, result)
                bestMove = move if ret == result else bestMove
                if alpha >= beta or ret == 1:
                    return ret, move
            return ret, bestMove
        else:
            ret = math.inf
            for pos in game.getAvailablePositions():
                move = pos
                result = game.move(*pos)
                if result is None:
                    assert not game.gameOver
                    result, oppMove = self.alpha_beta(game.getStatus(), alpha, beta)
                game.undo()
                beta = min(beta, result)
                ret = min(ret, result)
                bestMove = move if ret == result else bestMove
                if alpha >= beta or ret == -1:
                    return ret, move
            return ret, bestMove

class AlphaBetaDPStrategy(Strategy):
    def action(self, game: ConnectNGame) -> Tuple[int, Tuple[int, int]]:
        self.game = game
        self.alphaBetaStack = [(-math.inf, math.inf)]
        result, move = self.alpha_beta_dp(self.game.getStatus())
        return result, move

    @lru_cache(maxsize=None)
    def alpha_beta_dp(self, gameStatus: Tuple[Tuple[int, ...]]) -> Tuple[int, Tuple[int, int]]:
        alpha, beta = self.alphaBetaStack[-1]
        game = self.game
        bestMove = None
        assert not game.gameOver
        if game.currentPlayer == ConnectNGame.PLAYER_A:
            ret = -math.inf
            for pos in game.getAvailablePositions():
                move = pos
                result = game.move(*pos)
                if result is None:
                    assert not game.gameOver
                    self.alphaBetaStack.append((alpha, beta))
                    result, oppMove = self.alpha_beta_dp(game.getStatus())
                    self.alphaBetaStack.pop()
                game.undo()
                alpha = max(alpha, result)
                ret = max(ret, result)
                bestMove = move if ret == result else bestMove
                if alpha >= beta or ret == 1:
                    return ret, move
            return ret, bestMove
        else:
            ret = math.inf
            for pos in game.getAvailablePositions():
                move = pos
                result = game.move(*pos)
                if result is None:
                    assert not game.gameOver
                    self.alphaBetaStack.append((alpha, beta))
                    result, oppMove = self.alpha_beta_dp(game.getStatus())
                    self.alphaBetaStack.pop()
                game.undo()
                beta = min(beta, result)
                ret = min(ret, result)
                bestMove = move if ret == result else bestMove
                if alpha >= beta or ret == -1:
                    return ret, move
            return ret, bestMove