State = Tuple[int, int, int, int]
Action = int

class CartPoleAbstractAgent(metaclass=abc.ABCMeta):
    def __init__(self, buckets=(1, 2, 6, 12), discount=0.98, lr_min=0.1, epsilon_min=0.1):
        self.env = gym.make('CartPole-v0')

        env = self.env
        # [position, velocity, angle, angular velocity]
        self.dims_config = [(env.observation_space.low[0], env.observation_space.high[0], 1),
                            (-0.5, 0.5, 1),
                            (env.observation_space.low[2], env.observation_space.high[2], 6),
                            (-math.radians(50) / 1., math.radians(50) / 1., 12)]
        self.q = np.zeros(buckets + (self.env.action_space.n,))
        self.pi = np.zeros_like(self.q)
        self.pi[:] = 1.0 / env.action_space.n

    def to_bin_idx(self, val: float, lower: float, upper: float, bucket_num: int) -> int:
        percent = (val + abs(lower)) / (upper - lower)
        return min(bucket_num - 1, max(0, int(round((bucket_num - 1) * percent))))

    def discretize(self, obs: np.ndarray) -> State:
        discrete_states = tuple([self.to_bin_idx(obs[d], *self.dims_config[d]) for d in range(len(obs))])
        return discrete_states

def train(self, num_episodes=2000):
    for e in range(num_episodes):
        print(e)
        s: State = self.discretize(self.env.reset())

        self.adjust_learning_rate(e)
        self.adjust_epsilon(e)
        done = False

        while not done:
            action: Action = self.choose_action(s)
            obs, reward, done, _ = self.env.step(action)
            s_next: State = self.discretize(obs)
            a_next = self.choose_action(s_next)
            self.update_q(s, action, reward, s_next, a_next)
            s = s_next

def choose_action(self, state) -> Action:
    if np.random.random() < self.epsilon:
        return self.env.action_space.sample()
    else:
        return np.argmax(self.q[state])

class CartPoleAbstractAgent(metaclass=abc.ABCMeta):
    @abc.abstractmethod
    def update_q(self, s: State, a: Action, r, s_next: State, a_next: Action):
        pass

class SarsaAgent(CartPoleAbstractAgent):

    def update_q(self, s: State, a: Action, r, s_next: State, a_next: Action):
        self.q[s][a] += self.lr * (r + self.discount * (self.q[s_next][a_next]) - self.q[s][a])

class QLearningAgent(CartPoleAbstractAgent):

    def update_q(self, s: State, a: Action, r, s_next: State, a_next: Action):
        self.q[s][a] += self.lr * (r + self.discount * np.max(self.q[s_next]) - self.q[s][a])

class ExpectedSarsaAgent(CartPoleAbstractAgent):

    def update_q(self, s: State, a: Action, r, s_next: State, a_next: Action):
        self.q[s][a] = self.q[s][a] + self.lr * (r + self.discount * np.dot(self.pi[s_next], self.q[s_next]) - self.q[s][a])
        # update pi[s]
        best_a = np.random.choice(np.where(self.q[s] == max(self.q[s]))[0])
        n_actions = self.env.action_space.n
        self.pi[s][:] = self.epsilon / n_actions
        self.pi[s][best_a] = 1 - (n_actions - 1) * (self.epsilon / n_actions)

State = Tuple[int, int, bool]
Action = bool
Reward = float
Actions = np.ndarray
ActionValue = Dict[State, Actions]
Policy = Dict[State, Actions]

policy: Policy
A: ActionValue = policy[state]
action = np.random.choice([0, 1], p=A/sum(A))

def mc_control_exploring_starts(env: BlackjackEnv, num_episodes, discount_factor=1.0) \
        -> Tuple[ActionValue, Policy]:
    states = list(product(range(10, 22), range(1, 11), (True, False)))
    policy = {s: np.ones(env.action_space.n) * 1.0 / env.action_space.n for s in states}
    Q = defaultdict(lambda: np.zeros(env.action_space.n))
    returns_sum = defaultdict(float)
    returns_count = defaultdict(float)

    for episode_i in range(1, num_episodes + 1):
        s0 = random.choice(states)
        reset_env_with_s0(env, s0)
        episode_history = gen_custom_s0_stochastic_episode(policy, env, s0)

        G = 0
        for t in range(len(episode_history) - 1, -1, -1):
            s, a, r = episode_history[t]
            G = discount_factor * G + r
            if not any(s_a_r[0] == s and s_a_r[1] == a for s_a_r in episode_history[0: t]):
                returns_sum[s, a] += G
                returns_count[s, a] += 1.0
                Q[s][a] = returns_sum[s, a] / returns_count[s, a]
                best_a = np.argmax(Q[s])
                policy[s][best_a] = 1.0
                policy[s][1-best_a] = 0.0

    return Q, policy

def reset_env_with_s0(env: BlackjackEnv, s0: State) -> BlackjackEnv:
    env.reset()
    player_sum = s0[0]
    oppo_sum = s0[1]
    has_usable = s0[2]

    env.dealer[0] = oppo_sum
    if has_usable:
        env.player[0] = 1
        env.player[1] = player_sum - 11
    else:
        if player_sum > 11:
            env.player[0] = 10
            env.player[1] = player_sum - 10
        else:
            env.player[0] = 2
            env.player[1] = player_sum - 2
    return env

def mc_control_epsilon_greedy(env: BlackjackEnv, num_episodes, discount_factor=1.0, epsilon=0.1) \
        -> Tuple[ActionValue, Policy]:
    returns_sum = defaultdict(float)
    returns_count = defaultdict(float)

    states = list(product(range(10, 22), range(1, 11), (True, False)))
    policy = {s: np.ones(env.action_space.n) * 1.0 / env.action_space.n for s in states}
    Q = defaultdict(lambda: np.zeros(env.action_space.n))

    def update_epsilon_greedy_policy(policy: Policy, Q: ActionValue, s: State):
        policy[s] = np.ones(env.action_space.n, dtype=float) * epsilon / env.action_space.n
        best_action = np.argmax(Q[s])
        policy[s][best_action] += (1.0 - epsilon)

    for episode_i in range(1, num_episodes + 1):
        episode_history = gen_stochastic_episode(policy, env)

        G = 0
        for t in range(len(episode_history) - 1, -1, -1):
            s, a, r = episode_history[t]
            G = discount_factor * G + r
            if not any(s_a_r[0] == s and s_a_r[1] == a for s_a_r in episode_history[0: t]):
                returns_sum[s, a] += G
                returns_count[s, a] += 1.0
                Q[s][a] = returns_sum[s, a] / returns_count[s, a]
                update_epsilon_greedy_policy(policy, Q, s)

    return Q, policy

from gym.envs.toy_text import BlackjackEnv

State = Tuple[int, int, bool]
Action = bool
Reward = float
StateValue = Dict[State, float]
DeterministicPolicy = Callable[[State], Action]

class BlackjackEnv(gym.Env):

    def step(self, action):
        assert self.action_space.contains(action)
        if action:  # hit: add a card to players hand and return
            self.player.append(draw_card(self.np_random))
            if is_bust(self.player):
                done = True
                reward = -1.
            else:
                done = False
                reward = 0.
        else:  # stick: play out the dealers hand, and score
            done = True
            while sum_hand(self.dealer) < 17:
                self.dealer.append(draw_card(self.np_random))
            reward = cmp(score(self.player), score(self.dealer))
            if self.natural and is_natural(self.player) and reward == 1.:
                reward = 1.5
        return self._get_obs(), reward, done, {}

    def _get_obs(self):
        return (sum_hand(self.player), self.dealer[0], usable_ace(self.player))

def gen_episode_data(policy: DeterministicPolicy, env: BlackjackEnv) -> List[Tuple[State, Action, Reward]]:
    episode_history = []
    state = env.reset()
    done = False
    while not done:
        action = policy(state)
        next_state, reward, done, _ = env.step(action)
        episode_history.append((state, action, reward))
        state = next_state
    return episode_history

def fixed_policy(observation):
    """
    sticks if the player score is >= 20 and hits otherwise.
    """
    score, dealer_score, usable_ace = observation
    return 0 if score >= 20 else 1

def mc_prediction_first_visit(policy: DeterministicPolicy, env: BlackjackEnv,
                              num_episodes, discount_factor=1.0) -> StateValue:
    returns_sum = defaultdict(float)
    returns_count = defaultdict(float)

    for episode_i in range(1, num_episodes + 1):
        episode_history = gen_episode_data(policy, env)

        G = 0
        for t in range(len(episode_history) - 1, -1, -1):
            s, a, r = episode_history[t]
            G = discount_factor * G + r
            if not any(s_a_r[0] == s for s_a_r in episode_history[0: t]):
                returns_sum[s] += G
                returns_count[s] += 1.0

    V = defaultdict(float)
    V.update({s: returns_sum[s] / returns_count[s] for s in returns_sum.keys()})
    return V

def mc_prediction_every_visit(policy: DeterministicPolicy, env: BlackjackEnv,
                              num_episodes, discount_factor=1.0) -> StateValue:
    returns_sum = defaultdict(float)
    returns_count = defaultdict(float)

    for episode_i in range(1, num_episodes + 1):
        episode_history = gen_episode_data(policy, env)

        G = 0
        for t in range(len(episode_history) - 1, -1, -1):
            s, a, r = episode_history[t]
            G = discount_factor * G + r
            returns_sum[s] += G
            returns_count[s] += 1.0

    V = defaultdict(float)
    V.update({s: returns_sum[s] / returns_count[s] for s in returns_sum.keys()})
    return V

[[ 0. -1. -2. -3.]
 [-1. -2. -3. -2.]
 [-2. -3. -2. -1.]
 [-3. -2. -1.  0.]]

def action_value(env: GridWorldEnv, state: State, V: StateValue, gamma=1.0) -> ActionValue:
    q = np.zeros(env.nA)
    for a in range(env.nA):
        for prob, next_state, reward, done in env.P[state][a]:
            q[a] += prob * (reward + gamma * V[next_state])
    return q

def policy_improvement(env: GridWorldEnv, policy: Policy, V: StateValue, gamma=1.0) -> bool:
    policy_stable = True

    for s in range(env.nS):
        old_action = np.argmax(policy[s])
        Q_s = action_value(env, s, V)
        best_action = np.argmax(Q_s)
        policy[s] = np.eye(env.nA)[best_action]

        if old_action != best_action:
            policy_stable = False
    return policy_stable


def policy_iteration(env: GridWorldEnv, policy: Policy, gamma=1.0) -> Tuple[Policy, StateValue]:
    iter = 0
    while True:
        V = policy_evaluate(policy, env, gamma)
        policy_stable = policy_improvement(env, policy, V)
        iter += 1

        if policy_stable:
            return policy, V

def value_iteration(env:GridWorldEnv, gamma=1.0, theta=0.0001) -> Tuple[Policy, StateValue]:
    V = np.zeros(env.nS)
    while True:
        delta = 0
        for s in range(env.nS):
            action_values = action_value(env, s, V, gamma=gamma)
            best_action_value = np.max(action_values)
            delta = max(delta, np.abs(best_action_value - V[s]))
            V[s] = best_action_value
        if delta < theta:
            break

    policy = np.zeros([env.nS, env.nA])
    for s in range(env.nS):
        action_values = action_value(env, s, V, gamma=gamma)
        best_action = np.argmax(action_values)
        policy[s, best_action] = 1.0

    return policy, V

def build_reverse_mapping(env:GridWorldEnv) -> Dict[State, Set[State]]:
    MAX_R, MAX_C = env.shape[0], env.shape[1]
    mapping = {s: set() for s in range(0, MAX_R * MAX_C)}
    action_delta = {Action.UP: (-1, 0), Action.DOWN: (1, 0), Action.LEFT: (0, -1), Action.RIGHT: (0, 1)}
    for s in range(0, MAX_R * MAX_C):
        r = s // MAX_R
        c = s % MAX_R
        for a in list(Action):
            neighbor_r = min(MAX_R - 1, max(0, r + action_delta[a][0]))
            neighbor_c = min(MAX_C - 1, max(0, c + action_delta[a][1]))
            s_ = neighbor_r * MAX_R + neighbor_c
            mapping[s_].add(s)
    return mapping

def value_iteration_async(env:GridWorldEnv, gamma=1.0, theta=0.0001) -> Tuple[Policy, StateValue]:
    mapping = build_reverse_mapping(env)

    V = np.zeros(env.nS)
    changed_state_set = set(s for s in range(env.nS))

    iter = 0
    while len(changed_state_set) > 0:
        changed_state_set_ = set()
        for s in changed_state_set:
            action_values = action_value(env, s, V, gamma=gamma)
            best_action_value = np.max(action_values)
            v_diff = np.abs(best_action_value - V[s])
            if v_diff > theta:
                changed_state_set_.update(mapping[s])
                V[s] = best_action_value
        changed_state_set = changed_state_set_
        iter += 1

    policy = np.zeros([env.nS, env.nA])
    for s in range(env.nS):
        action_values = action_value(env, s, V, gamma=gamma)
        best_action = np.argmax(action_values)
        policy[s, best_action] = 1.0

    return policy, V

MDP = Dict[State, Dict[Action, List[Tuple[Prob, State, Reward, bool]]]]

# P[state][action] = [
#    (prob1, next_state1, reward1, is_done),
#    (prob2, next_state2, reward2, is_done), ...]

class Action(Enum):
    UP = 0
    DOWN = 1
    LEFT = 2
    RIGHT = 3

State = int
Reward = float
Prob = float
Policy = Dict[State, Dict[Action, Prob]]
Value = List[float]
StateSet = Set[int]
NonTerminalStateSet = Set[int]
MDP = Dict[State, Dict[Action, List[Tuple[Prob, State, Reward, bool]]]]
# P[s][a] = [(prob, next_state, reward, is_done), ...]

class GridWorldEnv(discrete.DiscreteEnv):
    """
    Grid World environment described in Sutton and Barto Reinforcement Learning 2nd, chapter 4.
    """

    def __init__(self, shape=[4,4]):
        self.shape = shape
        nS = np.prod(shape)
        nA = len(list(Action))
        MAX_R = shape[0]
        MAX_C = shape[1]
        self.grid = np.arange(nS).reshape(shape)
        isd = np.ones(nS) / nS

        # P[s][a] = [(prob, next_state, reward, is_done), ...]
        P: MDP = {}
        action_delta = {Action.UP: (-1, 0), Action.DOWN: (1, 0), Action.LEFT: (0, -1), Action.RIGHT: (0, 1)}
        for s in range(0, MAX_R * MAX_C):
            P[s] = {a.value : [] for a in list(Action)}
            is_terminal = self.is_terminal(s)
            if is_terminal:
                for a in list(Action):
                    P[s][a.value] = [(1.0, s, 0, True)]
            else:
                r = s // MAX_R
                c = s % MAX_R
                for a in list(Action):
                    neighbor_r = min(MAX_R-1, max(0, r + action_delta[a][0]))
                    neighbor_c = min(MAX_C-1, max(0, c + action_delta[a][1]))
                    s_ = neighbor_r * MAX_R + neighbor_c
                    P[s][a.value] = [(1.0, s_, -1, False)]

        super(GridWorldEnv, self).__init__(nS, nA, P, isd)

def policy_evaluate(policy: Policy, env: GridWorldEnv, gamma=1.0, theta=0.0001):
    V = np.zeros(env.nS)
    while True:
        delta = 0
        for s in range(env.nS):
            v = 0
            for a, action_prob in enumerate(policy[s]):
                for prob, next_state, reward, done in env.P[s][a]:
                    v += action_prob * prob * (reward + gamma * V[next_state])
            delta = max(delta, np.abs(v - V[s]))
            V[s] = v
        if delta < theta:
            break
    return np.array(V)

[[  0.         -13.99931242 -19.99901152 -21.99891199]
 [-13.99931242 -17.99915625 -19.99908389 -19.99909436]
 [-19.99901152 -19.99908389 -17.99922697 -13.99942284]
 [-21.99891199 -19.99909436 -13.99942284   0.        ]]

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