gives 2 plies instead of 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 a direct outcome is already determined at this node
#         if 'result' in node:
#             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

#         # Initialize variables to track the best outcome
#         best_result = 0  # Assume draw as default
#         best_path = []
#         best_moves_to_mate = None

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

#             # Skip paths that don't lead to a decisive outcome
#             if child_result == 0 or child_moves_to_mate is None:
#                 continue

#             # Logic to select the optimal path based on the current node's perspective
#             if best_moves_to_mate is None:
#                 # First child with a decisive outcome sets the initial values
#                 best_result, best_path, best_moves_to_mate = child_result, child_path, child_moves_to_mate
#             else:
#                 if node['color'] == chess.WHITE:
#                     # White seeks to minimize moves to mate when winning
#                     if 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
#                 else:
#                     # Black seeks to minimize moves to mate when winning
#                     if 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

#         # Update the current node with the optimal outcome found
#         if best_moves_to_mate is not None:
#             node['result'] = best_result
#             node['moves_to_mate'] = 1 + best_moves_to_mate  # Increment for the move to this node
#         else:
#             node['result'] = best_result  # Could be 0 if all children result in a draw
#             node['moves_to_mate'] = None

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

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

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



# This function assumes that 'game_tree' is a dictionary representing the nodes of your game tree,
# where each node ID maps to a dictionary containing:
# 'fen': The FEN string of the board position,
# 'parent': The parent node ID,
# 'color': The color to move (True for White, False for Black),
# 'children': A list of child node IDs,
# 'result': The game outcome from this position (-1 for Black win, 1 for White win, 0 for draw),
# 'moves_to_mate': The shortest number of moves to mate from this position (if applicable).

# To execute this code, you'll need to define your 'game_tree' variable with the initial game state and
# have the chess positions added to the tree using the 'add_descendants_iteratively' function you've defined.



# def update_parent_preferences(game_tree):
#     def recurse(node_id):
#         node = game_tree[node_id]
#         if 'result' in node:
#             # Direct outcome detected, set moves_to_mate accordingly
#             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 = 0  # Default to draw
#         best_path = []
#         best_moves_to_mate = float('inf')  # Initialize with a large number for comparison

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

#             if child_result == 0 or child_moves_to_mate is None:
#                 continue  # Skip non-result or draw paths

#             # Logic for selecting the shortest path to victory
#             if ((node['color'] == chess.WHITE and child_result == 1) 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_moves_to_mate != float('inf'):
#             node['result'] = best_result
#             node['moves_to_mate'] = best_moves_to_mate + 1  # Correctly increment for the move to this node
#         else:
#             node['result'] = best_result  # Draw if no better path found
#             node['moves_to_mate'] = None  # No valid path to victory

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

#     result, path, moves_to_mate = recurse(1)  # Assuming 1 is the root node's ID
#     return path


# def update_parent_preferences(game_tree):
#     def recurse(node_id):
#         node = game_tree[node_id]
#         if 'result' in node:  # Check if the node has a direct outcome
#             return node['result'], [node_id], 0 if 'moves_to_mate' in node else None

#         best_result = 0  # Assume draw as the default outcome
#         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_result == 0 or child_moves_to_mate is None:  # Skip draws or paths not leading to checkmate
#                 continue

#             # Ensure we have a valid best_moves_to_mate for comparison
#             if best_moves_to_mate is None:
#                 best_moves_to_mate = child_moves_to_mate
#                 best_result, best_path = child_result, child_path
#             else:
#                 # Compare only if child_moves_to_mate is not None and update accordingly
#                 if node['color'] == chess.WHITE and child_result == 1:
#                     if child_moves_to_mate < best_moves_to_mate:
#                         best_result, best_path, best_moves_to_mate = child_result, child_path, child_moves_to_mate
#                 elif node['color'] == chess.BLACK and child_result == -1:
#                     if 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_moves_to_mate is not None:
#             node['result'] = best_result
#             node['moves_to_mate'] = 1 + best_moves_to_mate
#         else:
#             node['result'] = best_result  # could still be 0 if no child paths lead to a win
#             node['moves_to_mate'] = None

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

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


# Make sure to replace or integrate this function into your main execution logic and re-run your analysis.

# def update_parent_preferences(game_tree):
#     def recurse(node_id):
#         node = game_tree[node_id]
#         if 'result' in node:
#             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 = 0  # Default outcome
#         best_path = []
#         best_moves_to_mate = float('inf')  # Use infinity as the initial comparison value

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

#             if child_result == 0 or child_moves_to_mate is None:
#                 continue  # Skip non-result paths

#             # Looking for the shortest path to checkmate for both colors
#             if (child_result == 1 and node['color'] == chess.WHITE) or (child_result == -1 and node['color'] == chess.BLACK):
#                 if 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_moves_to_mate != float('inf'):
#             node['result'] = best_result
#             node['moves_to_mate'] = 1 + best_moves_to_mate  # Incrementing for the current move
#         else:
#             node['result'] = 0  # Draw or no path to checkmate found
#             node['moves_to_mate'] = None

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

#     _, path, _ = recurse(1)
#     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()