partially introduced castling, debating move rewrite

app/chess_sim/test.py

@@ -1,6 +1,6 @@
 from enum import Enum, auto
 from dataclasses import dataclass
-from typing import Optional, List, Tuple
+from typing import Iterator, Optional, List, Tuple
 import random
 
 class PieceType(Enum):
@@ -16,6 +16,13 @@ class Color(Enum):
     WHITE = auto()
     BLACK = auto()
 
+    @property
+    def opposite(self) -> "Color": 
+        if self == Color.WHITE:
+            return Color.BLACK
+        else:
+            return Color.WHITE
+
 
 @dataclass
 class BoardPos:
@@ -24,6 +31,13 @@ class BoardPos:
     def __str__(self) -> str:
         return f"({self.p[0]},{self.p[1]})"
     
+    @property
+    def x (self) -> int:
+        return self.p[0]
+    @property
+    def y (self) -> int:
+        return self.p[1]
+
 
 @dataclass
 class Piece:
@@ -120,7 +134,7 @@ class ChessBoard:
         src_field = self.fields[sr][sc]
 
         # validate move using legal moves for the player color
-        legal = self.moves_basic_checked(color)
+        legal = self.moves_basic(color)
         if not any(m.m_from.p == move.m_from.p and m.m_to.p == move.m_to.p for m in legal):
             return False
 
@@ -211,25 +225,20 @@ class ChessBoard:
         return 0 <= rr < 8 and 0 <= cc < 8
 
     # takes the color of the player whos possible moves will be returned
-    def moves_basic_checked(self, color: Color) -> List[BoardMove]:
+    def moves_basic(self, color: Color) -> List[BoardMove]:
         moves: List[BoardMove] = []
-        for i, row in enumerate(self.fields):
+        for pos, piece in self.iter_pieces():
-            for j, field in enumerate(row):
+            if piece.color != color:
-                if field.piece is None:
                 continue
-                if field.piece.color != color:
+            rays = self.moves_unchecked(piece, pos)
-                    continue
-                src_pos = BoardPos((i, j))
-                piece = field.piece
-                rays = self.moves_unchecked(piece, src_pos)
 
             for ray in rays:
                 # detect pawn capture ray and dont allow movement in that ray without capturing
                 is_pawn_capture_ray = (
                     piece.type == PieceType.PAWN
                     and len(ray) == 1
-                        and ray[0][1] != j
+                    and ray[0][1] != pos.y
-                        and abs(ray[0][0] - i) == 1
+                    and abs(ray[0][0] - pos.x) == 1
                 )
 
                 # non sliding rays get treated individually
@@ -246,7 +255,7 @@ class ChessBoard:
                         if is_pawn_capture_ray:
                             # pawn cant move diagonally into empty square
                             break
-                            moves.append(BoardMove(src_pos, BoardPos((tr, tc))))
+                        moves.append(BoardMove(pos, BoardPos((tr, tc))))
                         continue    
                     else:
                         # occupied
@@ -254,10 +263,108 @@ class ChessBoard:
                             break
                         else:
                             # opponent piece
-                                moves.append(BoardMove(src_pos, BoardPos((tr, tc))))
+                            moves.append(BoardMove(pos, BoardPos((tr, tc))))
                             break
+        
+        
+        # manually adding castling
+        row_height = 0
+        if color == Color.BLACK:
+            row_height = 7
+        expected_king_pos = BoardPos((row_height,4))
+
+        # if the king has moved, return as usual
+        if self.has_piece_moved(expected_king_pos):
             return moves
 
+        expected_lhs_rook_pos = BoardPos((row_height,0))
+        expected_rhs_rook_pos = BoardPos((row_height,7))
+        lhs_pieces_inbetween = [BoardPos((row_height, i)) for i in range(1, 3)]
+        rhs_pieces_inbetween = [BoardPos((row_height, i)) for i in range(5, 6)]
+
+        # make sure the pieces inbetween are gone and the rook has not moved
+        if not self.has_piece_moved(expected_lhs_rook_pos) and self.are_pieces_none(lhs_pieces_inbetween):
+            #todo: allow castling here
+        if not self.has_piece_moved(expected_rhs_rook_pos) and self.are_pieces_none(rhs_pieces_inbetween):
+            #todo: allow castling here
+        
+        return moves
+
+    def has_piece_moved(self, pos: BoardPos) -> bool:
+        for move in self.move_history:
+            if move[0].m_from == pos:
+                return True
+            elif move[0].m_to == pos:
+                return True 
+        return False
+
+    def are_pieces_none(self, positions: List[BoardPos]):
+        are_gone = True
+        for pos in positions:
+            if not self.is_piece_none(pos):
+                are_gone = False
+        return are_gone
+
+    def is_piece_none(self, pos: BoardPos) -> bool:
+        piece = self.fields[pos.x][pos.y].piece
+        if piece is None:
+            return True
+        else:
+            return False 
+
+    def is_piece_of_type(self, pos: BoardPos, piece_type: PieceType) -> bool:
+        piece = self.fields[pos.x][pos.y].piece
+        if piece is None:
+            return False
+        return piece.type == piece_type
+
+    def iter_pieces(self) -> Iterator[Tuple[BoardPos, Piece]]:
+        for row_idx, row in enumerate(self.fields):
+            for col_idx, field in enumerate(row):
+                if field.piece is not None:
+                    yield BoardPos((row_idx, col_idx)), field.piece
+
+    def pos_of_king(self, color: Color) -> BoardPos:
+        for pos, piece in self.iter_pieces():
+            if piece.type == PieceType.KING and piece.color == color:
+                return pos
+        raise ValueError("player has no king")
+            
+    # if a player does not have any moves, it has lost
+    def generate_moves(self, color: Color) -> List[BoardMove]:
+
+        # only moves after which the king is not in check are allowed
+        # todo: the current method of checking what move will resolve the check is based on bruteforcing
+        # maybe there is a better way of doing it, but for now this should suffice
+
+        us_moves_wo_check = []
+        for move in self.moves_basic(color):
+            # do a move and then check, if the king is still in check
+            # if it isnt, add the move to the possibles ones
+            if not self.make_move(move, color):
+                raise ValueError("self moves basic created a move, which cannot be done (hopefully unreachable)")
+            
+            all_basic_enemy_moves = self.moves_basic(color.opposite)
+            king_pos = self.pos_of_king(color)
+            king_in_check = False
+            for mv in all_basic_enemy_moves:
+                if mv.m_to == king_pos:
+                    king_in_check = True
+            if not king_in_check:
+                us_moves_wo_check.append(move) 
+
+            if not self.unmake_move():
+                raise ValueError("failed to unmake move, shouldnt be an issue here")
+
+        # check if our move lands directly on the enemies king, if it does, its illigal
+        us_moves_rule_compliant = []
+        king_pos = self.pos_of_king(color.opposite)
+        for move in us_moves_wo_check:
+            if move.m_to != king_pos:
+                us_moves_rule_compliant.append(move)
+        
+        return us_moves_wo_check
+
 
     def __str__(self) -> str:
         """
@@ -287,10 +394,9 @@ def play_random_game(board: Optional[ChessBoard] = None, max_moves: int = 400, v
     moves_played = 0
 
     while moves_played < max_moves:
-        legal = board.moves_basic_checked(current)
+        legal = board.generate_moves(current)
         if not legal:
             # no legal moves
-            winner = Color.BLACK if current == Color.WHITE else Color.WHITE
             if verbose:
                 print(f"no legal moves for {current.name} after {moves_played} moves")
             return  board, moves_played
@@ -309,14 +415,14 @@ def play_random_game(board: Optional[ChessBoard] = None, max_moves: int = 400, v
             print(f"{moves_played:03d}: {current.name} played {move}\nboard:\n{board}")
 
         # check whether the captured piece was a king
-        last_move, captured = board.move_history[-1]
+        _, captured = board.move_history[-1]
         if captured is not None and captured.type == PieceType.KING:
             if verbose:
                 print(f"king captured by {current.name} on move {moves_played}.")
+                raise RuntimeWarning("this should not be possible")
             return  board, moves_played
 
-        # switch side
+        current = current.opposite
-        current = Color.BLACK if current == Color.WHITE else Color.WHITE
 
     if verbose:
         print(f"reached move limit ({max_moves})")
@@ -334,6 +440,7 @@ def run_random_games(n: int = 100, max_moves: int = 400, verbose: bool = False):
 # used only for testing purposesa
 def main(): 
     default_brd = ChessBoard.init_default()
+    default_brd.generate_moves
 
     # standardized expected string for initial position
     expected_start = (
@@ -353,13 +460,13 @@ def main():
         raise AssertionError(f"initial board mismatch:\nexpected:\n{expected_start}\n\nactual:\n{actual_start}")
 
     # test num moves
-    mvs = default_brd.moves_basic_checked(Color.BLACK)
+    mvs = default_brd.moves_basic(Color.BLACK)
     expected_move_count = 20  # 16 pawn moves and 4 knight
     if len(mvs) != expected_move_count:
         raise AssertionError(f"initial move count for black mismatch: expected {expected_move_count}, got {len(mvs)}")
     
 
-    mvs = default_brd.moves_basic_checked(Color.WHITE)
+    mvs = default_brd.moves_basic(Color.WHITE)
     expected_move_count = 20  # 16 pawn moves and 4 knight
     if len(mvs) != expected_move_count:
         raise AssertionError(f"Initial move count for white mismatch: expected {expected_move_count}, got {len(mvs)}")