import numpy as np
from src.bitboard import get_bit, bit_to_row_col, popcount

class OthelloGame:
    def __init__(self):
        # Initial Board Setup (A1 = MSB, H8 = LSB)
        # Black pieces: D5 (35), E4 (28) -> 0x0000000810000000
        # White pieces: D4 (36), E5 (27) -> 0x0000001008000000
        self.player_bb = 0x0000000810000000 # Black starts
        self.opponent_bb = 0x0000001008000000
        self.turn = 1 # 1: Black, -1: White

    def get_valid_moves(self, player, opponent):
        """Calculates valid moves for 'player' against 'opponent'."""
        empty = ~(player | opponent) & 0xFFFFFFFFFFFFFFFF
        
        # Consistent with MSB=A1:
        # North: << 8. South: >> 8.
        # West: << 1 (mask A). East: >> 1 (mask H).
        mask_h = 0x0101010101010101
        mask_a = 0x8080808080808080
        
        # Directions
        shifts = [
             (lambda x: (x & ~mask_h) >> 1), # East
             (lambda x: (x & ~mask_a) << 1), # West
             (lambda x: (x << 8) & 0xFFFFFFFFFFFFFFFF), # North
             (lambda x: (x >> 8) & 0xFFFFFFFFFFFFFFFF), # South
             (lambda x: (x & ~mask_h) << 7), # NE (N+E -> <<8 + >>1 = <<7)
             (lambda x: (x & ~mask_a) << 9), # NW (N+W -> <<8 + <<1 = <<9)
             (lambda x: (x & ~mask_h) >> 9), # SE (S+E -> >>8 + >>1 = >>9)
             (lambda x: (x & ~mask_a) >> 7)  # SW (S+W -> >>8 + <<1 = >>7)
        ]
        
        valid_moves = 0
        for shift_func in shifts:
            candidates = shift_func(player) & opponent
            for _ in range(6): # Max 6 opponent pieces can be in between
                candidates |= shift_func(candidates) & opponent
            valid_moves |= shift_func(candidates) & empty
            
        return valid_moves

    def apply_move(self, player, opponent, move_bit):
        """Calculates new boards after move_bit."""
        if move_bit == 0:
            return player, opponent
            
        flipped = 0
        mask_h = 0x0101010101010101
        mask_a = 0x8080808080808080
        
        shifts = [
             (lambda x: (x & ~mask_h) >> 1), # East
             (lambda x: (x & ~mask_a) << 1), # West
             (lambda x: (x << 8) & 0xFFFFFFFFFFFFFFFF), # North
             (lambda x: (x >> 8) & 0xFFFFFFFFFFFFFFFF), # South
             (lambda x: (x & ~mask_h) << 7), # NE
             (lambda x: (x & ~mask_a) << 9), # NW
             (lambda x: (x & ~mask_h) >> 9), # SE
             (lambda x: (x & ~mask_a) >> 7)  # SW
        ]
        
        for shift_func in shifts:
            mask = shift_func(move_bit)
            potential_flips = 0
            while mask & opponent:
                potential_flips |= mask
                mask = shift_func(mask)
            if mask & player:
                flipped |= potential_flips
                
        new_player = player | move_bit | flipped
        new_opponent = opponent & ~flipped
        return new_player, new_opponent

    def play_move(self, move_bit):
        if move_bit != 0:
            self.player_bb, self.opponent_bb = self.apply_move(self.player_bb, self.opponent_bb, move_bit)
        
        # Turn always swaps (even on pass)
        self.player_bb, self.opponent_bb = self.opponent_bb, self.player_bb
        self.turn *= -1

    def is_terminal(self):
        p_moves = self.get_valid_moves(self.player_bb, self.opponent_bb)
        o_moves = self.get_valid_moves(self.opponent_bb, self.player_bb)
        return (p_moves == 0) and (o_moves == 0)