good but is 2 but should be 4

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):
    def recurse(node_id):
        node = game_tree[node_id]
        if 'result' in node:
            # Node has a direct outcome (win, loss, or draw)
            moves_to_mate = 0 if 'moves_to_mate' in node and node['result'] in [-1, 1] else None
            return node['result'], [node_id], moves_to_mate

        best_result = None  # Track the best outcome (initially undefined)
        best_path = []  # Best path leading to this outcome
        best_moves_to_mate = float('inf')  # Use infinity as a comparison basis for finding the minimum

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

            # Skip draws or paths not leading to a decisive outcome
            if child_result == 0 or child_moves_to_mate is None:
                continue

            # Identify and select the best outcome for the current player
            if best_result is None or (child_result == -1 and node['color'] == chess.WHITE and child_moves_to_mate < best_moves_to_mate):
                # For White, a result of -1 (Black win) with fewer moves is worse, so it's avoided unless it's the only outcome
                best_result, best_path, best_moves_to_mate = child_result, child_path, child_moves_to_mate
            elif child_result == 1 and node['color'] == chess.BLACK and child_moves_to_mate < best_moves_to_mate:
                # For Black, similar logic applies but reversed
                best_result, best_path, best_moves_to_mate = child_result, child_path, child_moves_to_mate

        # If a winning path was found, update the node
        if best_result in [-1, 1]:
            node['result'] = best_result
            node['moves_to_mate'] = best_moves_to_mate + 1  # Adjust for the current move
        else:
            # No winning path found, consider it a draw or maintain undefined result
            node['result'] = 0 if best_result is None else best_result
            node['moves_to_mate'] = None

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

    _, path, _ = recurse(1)  # Starting from the root of the game tree
    return path



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,node,"<<\n")
        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"
    initial_fen = "8/8/8/3k4/1K6/8/5q2/8 b - - 0 1"
    initial_fen = "5K2/q7/6k1/8/8/8/8/8 w - - 0 1"
  #  initial_fen = "6K1/q7/6k1/8/8/8/8/8 b - - 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()