import chess
from itertools import permutations

def generate_chess_positions(pieces):
    board_size = 8
    all_squares = [(i, j) for i in range(board_size) for j in range(board_size)]
    for squares in permutations(all_squares, len(pieces)):
        board = chess.Board(None)
        for piece, (row, col) in zip(pieces, squares):
            square_index = row * 8 + col  # Convert from (row, col) to square index
            board.set_piece_at(square_index, chess.Piece.from_symbol(piece))
            yield board.fen()

def simplify_fen_string(fen):
    parts = fen.split(' ')
    simplified_fen = ' '.join(parts[:4])  # Keeping only the position information
    return simplified_fen

def evaluate_board(board):
    if board.is_checkmate():
        return 1000 if board.turn == chess.BLACK else -1000
    elif board.is_stalemate() or board.is_insufficient_material():
        return 0
    return 4  # Default heuristic

def minimax(board, depth, alpha, beta, maximizing_player, node_count, completed_depths):
    node_count[0] += 1
    if node_count[0] % 1000000 == 0:
        print(f"Nodes explored: {node_count[0]}")

    if board.is_game_over():
        if depth not in completed_depths:
            print(f"Depth {depth} completed for the first time.")
            completed_depths.add(depth)
        return [], evaluate_board(board)

    if depth > 3:  # Practical limit to prevent infinite recursion
        if depth not in completed_depths:
            print(f"Depth {depth} completed for the first time.")
            completed_depths.add(depth)
        return [], evaluate_board(board)

    best_eval = float('-inf') if maximizing_player else float('inf')
    best_sequence = []

    for move in board.legal_moves:
        board.push(move)
        sequence, eval = minimax(board, depth + 1, alpha, beta, not maximizing_player, node_count, completed_depths)
        board.pop()

        if maximizing_player and eval > best_eval:
            best_eval = eval
            best_sequence = [board.san(move)] + sequence
            alpha = max(alpha, eval)
            if beta <= alpha:
                break
        elif not maximizing_player and eval < best_eval:
            best_eval = eval
            best_sequence = [board.san(move)] + sequence
            beta = min(beta, eval)
            if beta <= alpha:
                break

    if depth not in completed_depths:
        print(f"Depth {depth} completed for the first time.")
        completed_depths.add(depth)

    return best_sequence, best_eval

# Main execution
initial_pieces = "KkQ"
positions = list(generate_chess_positions(initial_pieces))

if positions:
    # Use the first position from generated positions for minimax
    first_position_fen = "7k/2Q5/5K2/8/8/8/8/8 w - - 0 1" #positions[0]
    board = chess.Board(first_position_fen)

    node_count = [0]  # Track nodes across recursive calls
    completed_depths = set()  # Track completion of each depth level
    sequence, evaluation = minimax(board, 0, float('-inf'), float('inf'), True, node_count, completed_depths)

    # Output results
    simplified_fen = simplify_fen_string(first_position_fen)
    print("Simplified FEN for starting position:", simplified_fen)
    print("Best move sequence for first position:", sequence)
    print("Evaluation for first position:", evaluation)
    print("Total nodes explored:", node_count[0])
else:
    print("No positions were generated.")