encapsulation

positions.py

@@ -1,28 +1,55 @@
 import chess
 import chess.engine
 
-# Initialize Stockfish engine
-stockfish = chess.engine.SimpleEngine.popen_uci("./stockfish")
 
-# Sample FEN for board position; you can replace this with reading from a file
-# This FEN represents the starting position of a chess game
-fen_string = "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1"
-board = chess.Board(fen_string)
+def initialize_engine():
+    """Initialize and return the Stockfish engine."""
+    return chess.engine.SimpleEngine.popen_uci("./stockfish")
 
-# Placeholder: Uncomment the lines below if you'd like to read the FEN from a file
-# with open("path_to_fen_file.txt", "r") as f:
-#     fen_string = f.readline().strip()
-# board = chess.Board(fen_string)
 
-# Iterate over all squares and find legal moves for each piece
-for square in chess.SQUARES:
-    piece = board.piece_at(square)
-    if piece:
-        # Get legal moves for the current piece
-        legal_moves = [move for move in board.legal_moves if move.from_square == square]
-        target_squares = [move.to_square for move in legal_moves]
-        
-        print(f"{piece.symbol()} at {chess.square_name(square)} can move to: {[chess.square_name(ts) for ts in target_squares]}")
+def load_board_from_fen(fen: str):
+    """Load a board from a given FEN string."""
+    return chess.Board(fen)
 
-# Don't forget to quit the engine when done
-stockfish.quit()
+
+def get_legal_moves(board: chess.Board):
+    """Return legal moves for each piece on the board."""
+    moves = {}
+    for square in chess.SQUARES:
+        piece = board.piece_at(square)
+        if piece:
+            legal_moves = [move for move in board.legal_moves if move.from_square == square]
+            target_squares = [move.to_square for move in legal_moves]
+            moves[chess.square_name(square)] = [chess.square_name(ts) for ts in target_squares]
+    return moves
+
+
+def aggregate_target_squares(moves: dict):
+    """Count total attacks on each square."""
+    aggregated = {}
+    for _, target_squares in moves.items():
+        for ts in target_squares:
+            if ts in aggregated:
+                aggregated[ts] += 1
+            else:
+                aggregated[ts] = 1
+    return aggregated
+
+
+if __name__ == "__main__":
+    engine = initialize_engine()
+
+    # Sample FEN string. Replace with reading from a file or any other source if needed.
+    fen_string = "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1"
+    board = load_board_from_fen(fen_string)
+
+    moves = get_legal_moves(board)
+    for square, target_squares in moves.items():
+        print(f"Piece at {square} can move to: {target_squares}")
+
+    aggregated_targets = aggregate_target_squares(moves)
+    print("\nAggregated target squares:")
+    for square, count in aggregated_targets.items():
+        print(f"{square}: {count} attacks")
+
+    engine.quit()