"""
Created by: Brandon Schmidt
Last modified: 10/23/25

"""
import multiprocessing
from multiprocessing import Pool, cpu_count
from typing import List, Tuple

def _eval_root_move_worker(args: Tuple[int,int,int,int,Tuple[int,int,int,int,int,int,int],int]):
    """
    Worker for evaluating a single root move.
    args = (move, side_to_move, R, Y, cols_tuple, maxDepth)
    Returns (move, numeric_eval)
    """
    move, side_to_move, R, Y, cols_tuple, maxDepth = args

    # Reconstruct a fresh Bitboard in this process
    bb = Bitboard()
    bb.R = R
    bb.Y = Y
    # bb.move is a bool-like (1 = red to move). We set it so the worker can play the root move.
    bb.move = bool(side_to_move)
    bb.moveStack = []
    # restore the column map
    bb.colMap = {
        0: cols_tuple[0],
        1: cols_tuple[1],
        2: cols_tuple[2],
        3: cols_tuple[3],
        4: cols_tuple[4],
        5: cols_tuple[5],
        6: cols_tuple[6],
    }

    # Create a private Search instance for this worker (its TT is private)
    search = Search(maxDepth)
    # ensure fresh transposition table
    try:
        search.transposition = {}
    except Exception:
        pass

    # Play the candidate root move
    if bb.move:
        bb.moveR(move)
    else:
        bb.moveY(move)

    # Run the search from this child node
    posEval = search.go(bb, 1)

    # if search returns a tuple/list, extract numeric value
    if isinstance(posEval, (tuple, list)):
        posEval_val = posEval[0]
    else:
        posEval_val = posEval

    return (move, posEval_val)


# This is the game class. It allows the user to play against the Engine. 
class Game:
    def __init__(self):
        self.board = Bitboard()
        self.moves = Moves()
        self.evaluation = Evaluation()
        self.engine = Engine()
        self.lookup = {"R":0b1, "Y":0b0}

    def readFile(self, fileName):
        try:
            with open(fileName, 'r') as file:
                boardFile = file.read()
            lines = boardFile.splitlines()
            for line in reversed(lines):
                line = line.replace(" ", "")
                col = 0
                for char in line:
                    if char.upper() == "R":
                        self.board.moveR(col)
                    if char.upper() == "Y":
                        self.board.moveY(col)
                    col+=1
        except:
            return

    def writeFile(self, fileName):
        p1 = self.board.R
        p2 = self.board.Y
        rows = 6
        cols = 7

        with open(fileName, 'w') as file:
            for r in reversed(range(rows)):
                line = []
                for c in range(cols):
                    bit_index = r * cols + c
                    mask = 1 << bit_index
                    if p1 & mask:
                        line.append("R")
                    elif p2 & mask:
                        line.append("Y")
                    else:
                        line.append(".")
                file.write(" ".join(line)+"\n")
    
    def play(self):
        self.readFile("Connect4.txt")
        self.engine.findBest(self.board)
        if self.engine.bestMove!=None:
            if self.board.move:
                self.board.moveR(self.engine.bestMove)
            else:
                self.board.moveY(self.engine.bestMove)
            self.writeFile("Connect4.txt")

"""
This is the Engine class, it incorporates the Search class to
find the best move in the position. It works similar to the Search
class, but it returns a move instead of an evaluation.
"""
class Engine:
    """
    Parallel root-split engine. Drop-in replacement for your original Engine class.
    """

    def __init__(self, max_threads: int = 7, search_max_depth: int = 10):
        self.moves = Moves()
        self.evaluation = Evaluation()
        self.search_max_depth = int(search_max_depth)
        self.bestMove = None
        self.bestEval = 0
        self.max_threads = max(1, int(max_threads))
        self._pool = None

    def _make_pool(self, procs: int):
        # lazily create Pool (avoids creating processes at import time)
        if self._pool is None:
            # NOTE: On Windows, Pool will use 'spawn' start method (requires __main__ guard).
            self._pool = Pool(processes=procs)
        return self._pool

    def close_pool(self):
        if self._pool is not None:
            try:
                self._pool.close()
                self._pool.join()
            except Exception:
                pass
            self._pool = None

    def findBest(self, bitboardObject):
        """
        Root-parallel findBest. Uses up to min(max_threads, len(moves), cpu_count()) workers.
        Safe fallback to sequential evaluation if multiprocessing fails.
        """

        moveList = self.moves.get(bitboardObject)
        if not moveList:
            self.bestMove, self.bestEval = None, 0
            return

        if self.evaluation.wins(bitboardObject):
            return

        # Pack serializable state for workers
        R = bitboardObject.R
        Y = bitboardObject.Y
        cols_tuple = (
            bitboardObject.colMap.get(0, 0),
            bitboardObject.colMap.get(1, 0),
            bitboardObject.colMap.get(2, 0),
            bitboardObject.colMap.get(3, 0),
            bitboardObject.colMap.get(4, 0),
            bitboardObject.colMap.get(5, 0),
            bitboardObject.colMap.get(6, 0),
        )
        # side_to_move: 1 -> red (maximize), 0 -> yellow (minimize)
        side_to_move = 1 if bool(getattr(bitboardObject, "move", 1)) else 0

        # Number of worker processes: at most number of moves and CPUs and max_threads
        avail_cpus = cpu_count()
        procs = min(self.max_threads, len(moveList), avail_cpus)

        # Build arguments list for workers
        args_list = [
            (move, side_to_move, R, Y, cols_tuple, self.search_max_depth)
            for move in moveList
        ]

        results: List[Tuple[int, float]] = []

        # Try parallel evaluation; on any exception, fall back to sequential
        try:
            if procs > 1:
                pool = self._make_pool(procs)
                # map the lightweight worker over candidate moves
                results = pool.map(_eval_root_move_worker, args_list)
            else:
                # single-process path (deterministic and safe)
                results = [_eval_root_move_worker(a) for a in args_list]
        except Exception:
            # fallback: sequential evaluation
            results = [_eval_root_move_worker(a) for a in args_list]

        # Choose the best move depending on side to move
        if side_to_move:  # Red to move -> maximize
            best = max(results, key=lambda x: x[1])
        else:             # Yellow to move -> minimize
            best = min(results, key=lambda x: x[1])

        self.bestMove, self.bestEval = best[0], best[1]
        return
                    
"""
This is the Search class, it creates a game tree and searches it
to find the best path.
"""
class Search:
    # Creates a Moves and an Evaluation object to assist in
    # the search. It also takes a number as input, which is
    # the maximum search depth.
    def __init__(self, maxDepth):
        self.moves = Moves()
        self.evaluation = Evaluation()
        self.maxDepth = maxDepth

    # The main method of this class, which uses recursion to
    # search a game tree and find the best path forward.
    def go(self, bitboardObject, depth, alpha=float('-inf'), beta=float('inf')):
        # If the position has a win for red, return the maximum
        # evaluation minus the search depth
        if self.evaluation.winCheck(bitboardObject.R):
            return 0xFFFF - depth
        
        # If the position has a win for yellow, return the minimum
        # evaluation plus the search depth
        if self.evaluation.winCheck(bitboardObject.Y):
            return -0xFFFF + depth

        # If the depth is greater than the maximum depth, stop
        # searching and return a general position evaluation.
        if depth>self.maxDepth:
            return self.evaluation.get(bitboardObject)

        # If there are no moves, the game is a draw
        moves = self.moves.get(bitboardObject)
        if moves==[]:
            return 0

        # Red to move, maximize the game evaluation
        if bitboardObject.move:
            maxEval = -0xFFFF
            for move in moves:
                bitboardObject.moveR(move)
                moveEval = self.go(bitboardObject, depth+1, alpha, beta)
                bitboardObject.undo()
                maxEval = max(maxEval, moveEval)
                alpha = max(alpha, maxEval)
                if alpha >= beta:
                    break
            return maxEval
        # Yellow to move, minimize the game evaluation
        else:
            minEval = 0xFFFF
            for move in moves:
                bitboardObject.moveY(move)
                moveEval = self.go(bitboardObject, depth+1, alpha, beta)
                bitboardObject.undo()
                minEval = min(minEval, moveEval)
                beta = min(beta, minEval)
                if beta <= alpha:
                    break
            return minEval


"""
This is the Evaluation class, it handles win checks as well as
general position evaluations. 
"""
class Evaluation:
    # Initializes a heatmap that gives each square a value.
    # This value is the number of wins possible using that square
    def __init__(self):
        self.lookup = (
            3,  4,  5,  7,  5,  4,  3,
            4,  6,  8,  10, 8,  6,  4,
            6,  8,  11, 13, 11, 8,  6,
            6,  8,  11, 13, 11, 8,  6,
            4,  6,  8,  10, 8,  6,  4,
            3,  4,  5,  7,  5,  4,  3)

    # Returns true if the game state has a win
    def wins(self, bitboardObject):
        return self.winCheck(bitboardObject.R) or self.winCheck(bitboardObject.Y)

    # Shifts the board horizontally, vertically, and diagonally
    # to check for a win. bbnl and bbnr keep the method from
    # wrapping around to the next row and declaring a false win.
    def winCheck(self, bb):
        bbnl = bb&0b111111011111101111110111111011111101111110
        bbnr = bb&0b011111101111110111111011111101111110111111
        if bb & (bbnl >> 1) & (bbnl >> 2) & (bbnl >> 3):
            return True
        if bb & (bb >> 7) & (bb >> 14) & (bb >> 21):
            return True
        if bb & (bbnr >> 6) & (bbnr >> 12) & (bbnr >> 18):
            return True
        if bb & (bbnl >> 8) & (bbnl >> 16) & (bbnl >> 24):
            return True
        return False
    
    # Looks up every squares value and adds it to the score
    def scoreBitboard(self, bb):
        total = 0
        while bb:
            lsb = bb & -bb
            total += self.lookup[lsb.bit_length() - 1]
            bb &= bb - 1
        return total

    # Finds all of the potential wins in the position
    def getThreats(self, bb, occ):
        threats = 0
        bbnl = bb&0b111111011111101111110111111011111101111110
        bbnr = bb&0b011111101111110111111011111101111110111111
        empty = (~occ) & 0x3FFFFFFFFFF
        threats |= empty & (bb << 7) & (bb << 14) & (bb << 21)
        threats |= (bbnl >> 3) & (bbnl >> 2) & (bbnl >> 1) & empty
        threats |= (bbnl >> 2) & (bbnl >> 1) & empty & (bbnr << 1)
        threats |= (bbnl >> 1) & empty & (bbnr << 1) & (bbnr << 2)
        threats |= empty & (bbnr << 1) & (bbnr << 2) & (bbnr << 3)
        threats |= (bbnr >> 18) & (bbnr >> 12) & (bbnr >> 6) & empty
        threats |= (bbnr >> 12) & (bbnr >> 6) & empty & (bbnl << 6)
        threats |= (bbnr >> 6) & empty & (bbnl << 6) & (bbnl << 12)
        threats |= empty & (bbnl << 6) & (bbnl << 12) & (bbnl << 18)
        threats |= (bbnl >> 24) & (bbnl >> 16) & (bbnl >> 8) & empty
        threats |= (bbnl >> 16) & (bbnl >> 8) & empty & (bbnr << 8)
        threats |= (bbnl >> 8) & empty & (bbnr << 8) & (bbnr << 16)
        threats |= empty & (bbnr << 8) & (bbnr << 16) & (bbnr << 24)
        return threats  

    # Gives a board state a score based on likelihood winning.
    def get(self, bitboardObject):
        sumR = self.scoreBitboard(bitboardObject.R)
        sumY = self.scoreBitboard(bitboardObject.Y)
        occ  = bitboardObject.R | bitboardObject.Y
        threatsR = self.getThreats(bitboardObject.R, occ)
        threatsY = self.getThreats(bitboardObject.Y, occ)
        sumR += threatsR.bit_count() * 20
        sumY += threatsY.bit_count() * 20
        return sumR - sumY


"""
This is the Moves class. It's method function "get" will return
all of the possible moves for a given board.
"""
class Moves:
    # Simply finds all of the columns that aren't full
    # and returns a sorted list of those columns.
    def get(self, bitboardObject):
        moveList=[]
        bb = bitboardObject.R | bitboardObject.Y
        comb =  bb         & 0x7F
        comb &= (bb >> 7)  & 0x7F
        comb &= (bb >> 14) & 0x7F
        comb &= (bb >> 21) & 0x7F
        comb &= (bb >> 28) & 0x7F
        comb &= (bb >> 35) & 0x7F
        bb = 0b1111111 & ~comb
        while bb:
            lsb = bb & -bb
            index = (lsb.bit_length() - 1)
            moveList.append(index)
            bb &= bb - 1
        moveList.sort(key=lambda m: abs(m - 3))
        return moveList


"""
This is the Bitboard class, it stores the board state and all of 
the current moves that have been made. It's method functions allow
you to make and undo moves on the board.
"""
class Bitboard:
    def __init__(self):
        # Bitboard for the red pieces
        self.R = 0
        # Bitboard for the yellow pieces
        self.Y = 0
        # Player to move, where 0b1 represents
        # red and 0b0 represents yellow
        self.move = 0b1
        # Storage of moves
        self.moveStack = []
        # Keeps track of the column heights
        self.colMap = {
            0:0, 1:0, 2:0, 3:0,
            4:0, 5:0, 6:0}

    # Places a red piece into the specified column
    def moveR(self, column):
        self.moveStack.append(column)
        self.R |= 1 << (self.colMap[column]*7+column)
        self.colMap[column]+=1
        self.move ^= 0b1

    # Places a yellow piece into the specified column
    def moveY(self, column):
        self.moveStack.append(column)
        self.Y |= 1 << (self.colMap[column]*7+column)
        self.colMap[column]+=1
        self.move ^= 0b1

    # Undoes the last move
    def undo(self):
        column=self.moveStack.pop()
        self.move ^= 0b1
        self.colMap[column]-=1
        self.R &= ~(1 << (self.colMap[column]*7+column))
        self.Y &= ~(1 << (self.colMap[column]*7+column))


if __name__ == "__main__":
    game = Game()
    game.play()