/**
 * Copyright (c) 2026 Huawei Technologies Co., Ltd.
 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
 * CANN Open Software License Agreement Version 2.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */

/* Generated By CANNBot */

/*!
 * \file ball_query_proto.h
 * \brief Ball query operator proto declaration.
 */

#ifndef OPS_OP_PROTO_INC_BALL_QUERY_H_
#define OPS_OP_PROTO_INC_BALL_QUERY_H_

#include "graph/operator_reg.h"
#include "graph/types.h"

namespace ge {

/**
 *@brief Ball query: for each query center point center_xyz[m, b],
 *       search the point set xyz[b] (coordinates stored in the middle
 *       dimension) and collect the indices of points whose squared
 *       distance d2 to the center satisfies d2 == 0 ||
 *       (min_radius^2 <= d2 < max_radius^2). Up to sample_num points
 *       are kept in order; when fewer points are found, the remaining
 *       slots are filled with the index of the first selected point.
 *@par Inputs:
 * Two inputs, including:
 * @li xyz: A ND Tensor of shape (B, 3, N) where coordinates lie in the
 *         middle dimension. Must be one of: float16, float32.
 * @li center_xyz: A ND Tensor of shape (M, B, 3) holding the query
 *                center coordinates. Must be same dtype as xyz.
 *
 *@par Outputs:
 * idx: A ND Tensor of shape (M, B, sample_num) holding the selected
 *     point indices. Type int32.
 *
 *@par Attributes:
 * @li min_radius: Float, inner radius of the query ball, must be >= 0.
 * @li max_radius: Float, outer radius of the query ball, must be > 0.
 * @li sample_num: Int, max number of sampled points per ball, must be >= 1.
 */
REG_OP(BallQuery)
    .INPUT(xyz, TensorType({DT_FLOAT16, DT_FLOAT}))
    .INPUT(center_xyz, TensorType({DT_FLOAT16, DT_FLOAT}))
    .OUTPUT(idx, TensorType({DT_INT32}))
    .REQUIRED_ATTR(min_radius, Float)
    .REQUIRED_ATTR(max_radius, Float)
    .REQUIRED_ATTR(sample_num, Int)
    .OP_END_FACTORY_REG(BallQuery)

} // namespace ge

#endif // OPS_OP_PROTO_INC_BALL_QUERY_H_