solver.py

from transpositionTable import TranspositionTable


class Solver:
    __slots__ = ("board", "nodes", "columnOrder", "transpositionTable", "size")

    EXACT = 0
    LOWERBOUND = 1
    UPPERBOUND = 2

    def __init__(self, board):
        self.board = board
        self.nodes = 0
        self.columnOrder = (3, 4, 2, 5, 1, 6, 0)
        self.transpositionTable = TranspositionTable()
        self.size = board.size

    def negamax(self, alpha, beta):
        self.nodes += 1

        board = self.board
        move_number = board.moveNumber
        if move_number == self.size:
            return 0

        ai_position = board.aiPosition
        human_position = board.humanPosition
        mask = ai_position | human_position
        current_position = human_position if board.player == 1 else ai_position
        opponent_position = ai_position if board.player == 1 else human_position
        compute_winning_position = board.compute_winning_position

        possible_moves = (mask + board.bottom_mask) & board.board_mask
        if compute_winning_position(current_position, mask) & possible_moves:
            return (self.size + 1 - move_number) // 2

        alpha_orig = alpha
        key = current_position + mask
        tt_entry = self.transpositionTable.retrieve(key)
        tt_best_move = -1
        if tt_entry is not None:
            tt_value, tt_flag, tt_best_move = tt_entry
            if tt_flag == self.EXACT:
                return tt_value
            if tt_flag == self.LOWERBOUND:
                if tt_value > alpha:
                    alpha = tt_value
            else:
                if tt_value < beta:
                    beta = tt_value
            if alpha >= beta:
                return tt_value

        opponent_win = compute_winning_position(opponent_position, mask)
        forced_moves = possible_moves & opponent_win
        if forced_moves:
            if forced_moves & (forced_moves - 1):
                possible_mask = 0
            else:
                possible_mask = forced_moves
        else:
            possible_mask = possible_moves

        possible_mask &= ~(opponent_win >> 1)
        if possible_mask == 0:
            value = -(self.size - 1 - move_number) // 2
            self.transpositionTable.store(key, value, self.EXACT, -1)
            return value

        min_score = -(self.size - move_number) // 2
        if alpha < min_score:
            alpha = min_score
            if alpha >= beta:
                return alpha

        max_score = (self.size - 1 - move_number) // 2
        if beta > max_score:
            beta = max_score
            if alpha >= beta:
                return beta

        beta_bound = beta
        column_masks = board.COLUMN_MASKS
        negamax = self.negamax
        play = board.play
        undo = board.undo
        change_player = board.change_player

        best_move = -1
        best_score = -self.size

        tt_move_mask = 0
        if tt_best_move >= 0:
            tt_move_mask = possible_mask & column_masks[tt_best_move]

        if tt_move_mask:
            play(tt_move_mask)
            change_player()
            score = -negamax(-beta, -alpha)
            change_player()
            undo(tt_move_mask)

            best_move = tt_best_move
            best_score = score
            if score > alpha:
                alpha = score
            if alpha >= beta:
                self.transpositionTable.store(key, best_score, self.LOWERBOUND, best_move)
                return best_score

        ordered_moves = [0] * board.columns
        ordered_masks = [0] * board.columns
        move_scores = [0] * board.columns
        move_count = 0
        for move in self.columnOrder:
            if move == tt_best_move:
                continue

            move_mask = possible_mask & column_masks[move]
            if not move_mask:
                continue

            score = compute_winning_position(current_position | move_mask, mask | move_mask).bit_count()
            insert_at = move_count
            while insert_at > 0 and score > move_scores[insert_at - 1]:
                ordered_moves[insert_at] = ordered_moves[insert_at - 1]
                ordered_masks[insert_at] = ordered_masks[insert_at - 1]
                move_scores[insert_at] = move_scores[insert_at - 1]
                insert_at -= 1

            ordered_moves[insert_at] = move
            ordered_masks[insert_at] = move_mask
            move_scores[insert_at] = score
            move_count += 1

        for index in range(move_count):
            move = ordered_moves[index]
            move_mask = ordered_masks[index]
            play(move_mask)
            change_player()
            score = -negamax(-beta, -alpha)
            change_player()
            undo(move_mask)

            if score > best_score:
                best_score = score
                best_move = move

            if score > alpha:
                alpha = score
                if alpha >= beta:
                    break

        if best_score <= alpha_orig:
            flag = self.UPPERBOUND
        elif best_score >= beta_bound:
            flag = self.LOWERBOUND
        else:
            flag = self.EXACT

        self.transpositionTable.store(key, best_score, flag, best_move)
        return best_score

    def solve(self, weak=False):
        board = self.board
        mask = board.aiPosition | board.humanPosition
        current_position = board.humanPosition if board.player == 1 else board.aiPosition
        possible_moves = (mask + board.bottom_mask) & board.board_mask
        if board.compute_winning_position(current_position, mask) & possible_moves:
            return (self.size + 1 - board.moveNumber) // 2

        self.nodes = 0
        self.transpositionTable.clear()

        if weak:
            min_score, max_score = -1, 1
        else:
            min_score = -(self.size - board.moveNumber) // 2
            max_score = (self.size - 1 - board.moveNumber) // 2

        while min_score < max_score:
            med = min_score + (max_score - min_score) // 2
            if med <= 0 and min_score // 2 < med:
                med = min_score // 2
            elif med >= 0 and max_score // 2 > med:
                med = max_score // 2

            result = self.negamax(med, med + 1)
            if result <= med:
                max_score = result
            else:
                min_score = result

        return min_score