quite good moves_to_mate 14 but is 10

import chess

def simplify_fen(fen):
    """Simplifies a FEN string to include only position, turn, castling availability, and en passant target."""
    return ' '.join(fen.split(' ')[:4])

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, 'parent': None, 'color': chess.WHITE if 'w' in initial_fen else chess.BLACK, 'children': [], 'moves_to_mate': None}}
    fen_to_node_id = {simplified_fen: 1}
    return game_tree, fen_to_node_id

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:
            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, 'parent': node_id, 'color': chess.WHITE if board.turn else chess.BLACK, 'children': [], 'moves_to_mate': 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 update_game_outcomes(game_tree):
    """Updates game outcomes focusing only on checkmates."""
    for node_id, node in game_tree.items():
        board = chess.Board(node['fen'] + " 0 1")
        if board.is_checkmate():
            node['result'] = 1 if board.turn == chess.BLACK else -1  # Black's turn but checkmate means White wins, and vice versa
            node['moves_to_mate'] = 0  # Checkmate is immediate, no more moves required.
        elif board.is_game_over():
            node['result'] = 0  # For draws or stalemates, we consider the result as 0 (this will be ignored in propagation)

# def update_parent_preferences(game_tree):
#     """Propagates the best outcome for each node up the game tree, prioritizing checkmates and updating moves to mate."""
#     def recurse(node_id):
#         node = game_tree[node_id]
#         if 'result' in node and node['result'] in [-1, 1]:  # Checkmate detected
#             return node['result'], [node_id], 0 if 'moves_to_mate' in node else None

#         best_result = None
#         best_path = []
#         best_moves_to_mate = None
#         for child_id in node['children']:
#             child_result, child_path, child_moves_to_mate = recurse(child_id)

#             if child_moves_to_mate is not None:
#                 child_moves_to_mate += 1  # Increment moves to mate since we move up the tree

#                 if best_moves_to_mate is None or (node['color'] == chess.WHITE and child_result == 1 and child_moves_to_mate < best_moves_to_mate) or (node['color'] == chess.BLACK and child_result == -1 and child_moves_to_mate > best_moves_to_mate):
#                     best_result, best_path, best_moves_to_mate = child_result, child_path, child_moves_to_mate

#         if best_result in [-1, 1]:
#             node['result'] = best_result
#             node['moves_to_mate'] = best_moves_to_mate
#         else:
#             node['result'] = 0  # Default to draw if no win/loss is found

#         return node.get('result', 0), [node_id] + best_path, node.get('moves_to_mate', None)

#     result, path, moves_to_mate = recurse(1)  # Start the recursion from the root node.
#     return path

def update_parent_preferences(game_tree):
    """Updates the parent preferences for each node up the game tree, prioritizing checkmates."""
    def recurse(node_id):
        node = game_tree[node_id]
        if 'result' in node and node['result'] in [-1, 1]:  # Direct checkmate detected
            return node['result'], [node_id], 0

        best_result = None
        best_path = []
        # Initialize moves_to_mate with a large number for comparison
        best_moves_to_mate = float('inf') if node['color'] == chess.WHITE else -float('inf')

        for child_id in node['children']:
            child_result, child_path, child_moves_to_mate = recurse(child_id)

            # Update for WHITE (minimizing moves to mate) or BLACK (maximizing moves to mate)
            if child_moves_to_mate is not None:
                updated = False
                if node['color'] == chess.WHITE and child_result == 1:
                    if best_moves_to_mate > child_moves_to_mate:
                        best_moves_to_mate = child_moves_to_mate
                        updated = True
                elif node['color'] == chess.BLACK and child_result == -1:
                    if best_moves_to_mate < child_moves_to_mate:
                        best_moves_to_mate = child_moves_to_mate
                        updated = True

                if updated:
                    best_result = child_result
                    best_path = child_path

        # If no checkmate path was found, revert moves_to_mate to None and set result to draw
        if best_moves_to_mate == float('inf') or best_moves_to_mate == -float('inf'):
            best_moves_to_mate = None
            best_result = 0

        if best_result in [-1, 1]:
            node['moves_to_mate'] = best_moves_to_mate + 1 if best_moves_to_mate is not None else None
        else:
            node['result'] = 0  # Default to draw if no win/loss is found
            node['moves_to_mate'] = None

        return best_result, [node_id] + best_path, node.get('moves_to_mate')

    # Execute the recursive function starting from the root node.
    result, path, moves_to_mate = recurse(1)
    return path

# Ensure you call this function in your main to see the corrected output.


def print_path(game_tree, path):
    """Prints the board positions along the path."""
    for node_id in path:
        node = game_tree[node_id]
        board = chess.Board(node['fen'] + " 0 1")
        moves_to_mate = "N/A" if node.get('moves_to_mate') is None else node.get('moves_to_mate')
        print(f"Node ID: {node_id}, Result: {node.get('result', 'N/A')}, Moves to Mate: {moves_to_mate}")
        print(board)
        print("---")

def main():
    initial_fen = "8/8/8/3k4/8/8/2K2Q2/8 w - - 0 1"  # Example FEN
    initial_fen = "8/8/8/3k4/8/2K5/5q2/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)
    path = update_parent_preferences(game_tree)
    print(game_tree[1])
    print("Path to the outcome:")
    print_path(game_tree, path)

if __name__ == "__main__":
    main()