Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- import chess
- def simplify_fen(fen):
- """Simplifies a FEN string to only include position, turn, castling availability, and en passant target."""
- return ' '.join(fen.split(' ')[:4])
- def add_descendants_iteratively(game_tree, fen_to_node_id):
- """Expands the game tree by iteratively adding legal move descendants of each game state."""
- queue = [(1, 0)]
- while queue:
- node_id, _ = queue.pop(0)
- board = chess.Board(game_tree[node_id]['fen'] + " 0 1")
- for move in board.legal_moves:
- move = chess.Move(move.from_square, move.to_square, promotion=chess.QUEEN) if move.promotion else move
- board.push(move)
- simplified_fen = simplify_fen(board.fen())
- if simplified_fen not in fen_to_node_id:
- new_node_id = len(game_tree) + 1
- game_tree[new_node_id] = {
- 'fen': simplified_fen,
- 'moves_to_mate': None,
- 'parent': node_id,
- 'color': chess.WHITE if board.turn else chess.BLACK,
- 'result': None,
- 'processed': False,
- 'sequence': [],
- 'children': [],
- 'to_end': None,
- }
- fen_to_node_id[simplified_fen] = new_node_id
- game_tree[node_id]['children'].append(new_node_id)
- queue.append((new_node_id, 0))
- board.pop()
- def initialize_game_tree(initial_fen):
- """Initializes the game tree with the root node based on the initial FEN."""
- simplified_fen = simplify_fen(initial_fen)
- game_tree = {1: {'fen': simplified_fen, 'moves_to_mate': None, 'parent': None,
- 'color': chess.WHITE if 'w' in initial_fen else chess.BLACK,
- 'result': None, 'processed': False, 'sequence': [], 'children': [], 'to_end': None}}
- fen_to_node_id = {simplified_fen: 1}
- return game_tree, fen_to_node_id
- def update_game_outcomes(game_tree):
- """Updates game outcomes based on the current board state for each node."""
- for node in game_tree.values():
- board = chess.Board(node['fen'] + " 0 1")
- if board.is_game_over():
- outcome = board.outcome()
- node['processed'] = True
- node['result'] = 1 if outcome and outcome.winner == chess.WHITE else -1 if outcome and outcome.winner == chess.BLACK else 0
- node['moves_to_mate'] = 0 if node['result'] != 0 else None
- node['to_end'] = 0
- def update_parent_preferences(node_id, game_tree, stronger):
- node = game_tree[node_id]
- if node['processed']:
- return node['to_end'], node['moves_to_mate'], node['result']
- is_maximizing = node['color'] == stronger
- desired_result = 1 if stronger == chess.WHITE else -1
- # Initialize best_path_length based on whether we are maximizing or minimizing
- best_path_length = None # We start with None to be agnostic about the direction of comparison initially
- best_child_id = None
- best_result = None
- for child_id in node['children']:
- _, child_moves_to_mate, child_result = update_parent_preferences(child_id, game_tree, stronger)
- if child_moves_to_mate is None:
- continue
- # Decide when to update based on maximizing or minimizing player logic
- if is_maximizing:
- # For maximizing, we look for a smaller number than currently best (shortest path to checkmate)
- if best_path_length is None or (child_moves_to_mate < best_path_length):
- best_path_length, best_child_id, best_result = child_moves_to_mate, child_id, child_result
- else:
- # For minimizing, we look for a larger number than currently best (longest path to avoid checkmate)
- if best_path_length is None or (child_moves_to_mate > best_path_length):
- best_path_length, best_child_id, best_result = child_moves_to_mate, child_id, child_result
- if best_child_id is not None:
- node['sequence'] = [best_child_id]
- node['result'] = best_result
- node['to_end'] = best_path_length + 1 if best_path_length is not None else None
- node['moves_to_mate'] = best_path_length + 1 if best_result == desired_result else None
- else:
- node['to_end'], node['moves_to_mate'], node['sequence'] = None, None, []
- node['processed'] = True
- return node['to_end'], node['moves_to_mate'], node['result']
- def print_sequence_from_root(A):
- current_node_id = 1
- A[1]['sequence'] = [29]
- # A[1]['sequence'] = [13]
- while True:
- node = A[current_node_id]
- board = chess.Board(node['fen'] + " 0 1")
- print(f"Node ID: {current_node_id}, Moves to Mate: {node['moves_to_mate']}, Result: {node['result']}")
- print(board, "\n", node['fen'],"\n\n")
- if not node['sequence']:
- break
- current_node_id = node['sequence'][0]
- def main():
- initial_fen = "7Q/8/8/6k1/2K5/8/8/8 w - - 0 1"
- initial_fen = "8/5k1Q/3K4/8/8/8/8/8 b - - 0 1"
- initial_fen = "8/2k5/8/5Q2/8/8/2K5/8 w - - 0 1"
- initial_fen = "2k5/8/5Q2/3K4/8/8/8/8 w - - 0 1"
- game_tree, fen_to_node_id = initialize_game_tree(initial_fen)
- add_descendants_iteratively(game_tree, fen_to_node_id)
- update_game_outcomes(game_tree)
- update_parent_preferences(1, game_tree, chess.WHITE)
- print("Final Output:")
- print(game_tree[1])
- print(game_tree[2])
- print(game_tree[3])
- print(game_tree[4])
- print_sequence_from_root(game_tree)
- for key in range(1,20000):
- if game_tree[key]['to_end'] != None: # and A[key]['to_end'] > 9:
- print(f"{key}:: {game_tree[key]['to_end']} {game_tree[key]}\n{chess.Board(game_tree[key]['fen'])}<\n")
- if __name__ == "__main__":
- main()
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement