import chess
import chess.engine


def initialize_engine():
    """Initialize and return the Stockfish engine."""
    return chess.engine.SimpleEngine.popen_uci("./stockfish")


def load_board_from_fen(fen: str):
    """Load a board from a given FEN string."""
    return chess.Board(fen)


def get_legal_moves(board: chess.Board):
    """Return legal moves for each piece on the board."""
    moves = {}
    for square in chess.SQUARES:
        piece = board.piece_at(square)
        if piece:
            legal_moves = [move for move in board.legal_moves if move.from_square == square]
            target_squares = [move.to_square for move in legal_moves]
            moves[chess.square_name(square)] = [chess.square_name(ts) for ts in target_squares]
    return moves


def aggregate_target_squares(moves: dict):
    """Count total attacks on each square."""
    aggregated = {}
    for _, target_squares in moves.items():
        for ts in target_squares:
            if ts in aggregated:
                aggregated[ts] += 1
            else:
                aggregated[ts] = 1
    return aggregated


if __name__ == "__main__":
    engine = initialize_engine()

    # Sample FEN string. Replace with reading from a file or any other source if needed.
    fen_string = "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1"
    board = load_board_from_fen(fen_string)

    moves = get_legal_moves(board)
    for square, target_squares in moves.items():
        print(f"Piece at {square} can move to: {target_squares}")

    aggregated_targets = aggregate_target_squares(moves)
    print("\nAggregated target squares:")
    for square, count in aggregated_targets.items():
        print(f"{square}: {count} attacks")

    engine.quit()