import copy
from typing import List, Literal, Optional, Tuple

Player = Literal['X', 'O']
WinStatus = Literal[Player, 'TIE', None]


class ConnectNGame:
    """
    A game of Connect N, of width x height, where N is the number of pieces in a row/column/diagonal to win.
    """

    def __init__(self, width, height, n):
        self.width = width
        self.height = height
        self.n = n
        self.board = [[' ' for _ in range(width)] for _ in range(height)]

    def __str__(self):
        return '\n'.join(['|' + '|'.join(row) + '|' for row in self.board])

    def drop(self, column, player: Player) -> bool:
        if column < 0 or column >= self.width:
            return False
        for row in range(self.height - 1, -1, -1):
            if self.board[row][column] == ' ':
                self.board[row][column] = player
                return True
        return False

    def is_won(self) -> WinStatus:
        # Check rows
        for row in self.board:
            for i in range(self.width - self.n + 1):
                if row[i] != ' ' and all(row[i] == row[j] for j in range(i + 1, i + self.n)):
                    return row[i]

        # Check columns
        for j in range(self.width):
            for i in range(self.height - self.n + 1):
                if self.board[i][j] != ' ' and all(self.board[i][j] == self.board[k][j] for k in range(i + 1, i + self.n)):
                    return self.board[i][j]

        # Check diagonals
        for i in range(self.height - self.n + 1):
            for j in range(self.width - self.n + 1):
                if self.board[i][j] != ' ' and all(self.board[i][j] == self.board[i + k][j + k] for k in range(1, self.n)):
                    return self.board[i][j]

        for i in range(self.height - self.n + 1):
            for j in range(self.n - 1, self.width):
                if self.board[i][j] != ' ' and all(self.board[i][j] == self.board[i + k][j - k] for k in range(1, self.n)):
                    return self.board[i][j]

        # Check for tie
        if all(self.board[i][j] != ' ' for i in range(self.height) for j in range(self.width)):
            return 'TIE'

        return None

    def possible_moves(self) -> List[int]:
        return [col for col in range(self.width) if self.board[0][col] == ' ']

    def heuristic(self, player: Player) -> float:
        """
        Returns a heuristic score [-0.9, 0.9] for the board state.
        """
        score = 0

        # center column preference
        center_column = [self.board[i][self.width // 2]
                         for i in range(self.height)]
        center_count = center_column.count(player)
        score += center_count * 0.3

        # check rows, columns, and diagonals for potential wins
        for row in range(self.height):
            for col in range(self.width):
                if self.board[row][col] == ' ':
                    continue

                # horizontal potential
                if col + self.n <= self.width:
                    window = [self.board[row][c]
                              for c in range(col, col + self.n)]
                    score += self.evaluate_window(window, player)

                # vertical potential
                if row + self.n <= self.height:
                    window = [self.board[r][col]
                              for r in range(row, row + self.n)]
                    score += self.evaluate_window(window, player)

                # positive diagonal
                if col + self.n <= self.width and row + self.n <= self.height:
                    window = [self.board[row + i][col + i]
                              for i in range(self.n)]
                    score += self.evaluate_window(window, player)

                # negative diagonal
                if col - self.n >= -1 and row + self.n <= self.height:
                    window = [self.board[row + i][col - i]
                              for i in range(self.n)]
                    score += self.evaluate_window(window, player)

        return score

    def evaluate_window(self, window, player):
        opponent = 'O' if player == 'X' else 'X'
        score = 0
        if window.count(player) == self.n - 1 and window.count(' ') == 1:
            score += 0.5
        if window.count(player) == self.n - 2 and window.count(' ') == 2:
            score += 0.2
        if window.count(opponent) == self.n - 1 and window.count(' ') == 1:
            score -= 0.4
        return score

    def score_position(self, status: WinStatus, player: Player) -> float:
        """
        Assign scores to the board state. 
        Win is 1, loss is -1, tie (or ongoing) is heuristic.
        """
        status = self.is_won()
        if status == player:
            return 1
        elif status == 'TIE':
            return 0
        elif status is None:
            return self.heuristic(player)
        else:
            return -1

    def ai(self, depth: int, maximizing: bool, player: Player, alpha: float = float('-inf'), beta: float = float('inf')) -> Tuple[float, Optional[int]]:
        """ 
        Implements an AI that picks the "best" move using Minimax with Alpha-Beta pruning.
        Returns a tuple of (score, column).
        """
        opponent = 'O' if player == 'X' else 'X'

        status = self.is_won()
        if depth == 0 or status is not None:
            return self.score_position(status, player), None

        if maximizing:
            max_score = float('-inf')
            best_column = None
            for move in self.possible_moves():
                temp_game = copy.deepcopy(self)
                temp_game.drop(move, player)
                score, _ = temp_game.ai(
                    depth - 1, False, opponent, alpha, beta)
                if score > max_score:
                    max_score = score
                    best_column = move
                alpha = max(alpha, score)
                if alpha >= beta:
                    break
            return max_score, best_column
        else:
            min_score = float('inf')
            best_column = None
            for move in self.possible_moves():
                temp_game = copy.deepcopy(self)
                temp_game.drop(move, opponent)
                score, _ = temp_game.ai(depth - 1, True, player, alpha, beta)
                if score < min_score:
                    min_score = score
                    best_column = move
                beta = min(beta, score)
                if beta <= alpha:
                    break
            return min_score, best_column

    def best_move(self, player: Player, depth=4) -> int:
        """ Returns the best column for the player using Minimax. """
        _, best_column = self.ai(depth, False, player)
        if best_column is None:
            best_column = self.possible_moves()[0]
        return best_column