#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# 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

import numpy as np

__golden__ = {
    "kernel": {"ball_query": "ball_query_golden"},
}


def ball_query_golden(xyz, center_xyz, *, min_radius, max_radius, sample_num, **kwargs):
    """
    Golden function for ball_query.
    All the parameters (names and order) follow SE doc prototype definition without outputs.
    All the input Tensors are numpy.ndarray.

    布局与距离判定：
    - xyz 布局 (B, 3, N)：坐标在中间维
    - center_xyz 布局 (M, B, 3)
    - idx 输出 (M, B, sample_num)
    - 距离判定：d2 == 0 || (min_radius^2 <= d2 < max_radius^2)  环形区间+重合点
    - 首点填充：命中不足时用 first_num 填充剩余位置

    Args:
        xyz (np.ndarray): (B, 3, N) all point coordinates, float16 or float32.
            坐标在中间维。
        center_xyz (np.ndarray): (M, B, 3) query center coordinates, same dtype as xyz.
            布局 (M, B, 3)。
        min_radius (float): 最小查询半径（环形内边界），>= 0。
        max_radius (float): 最大查询半径（球外边界），> min_radius。
        sample_num (int): 每个查询球最大采样点数，>= 1。
        **kwargs: {input,output}_{dtypes,ori_shapes,formats,ori_formats},
                  full_soc_version, short_soc_version, testcase_name

    Returns:
        np.ndarray: (M, B, sample_num) int32, sampled point indices。
    """
    # 统一用 float32 计算距离，避免 float16 溢出
    xyz_f = xyz.astype(np.float32)
    center_xyz_f = center_xyz.astype(np.float32)

    M, B, _ = center_xyz_f.shape
    _, _, N = xyz_f.shape

    # N=0 边界处理：空点集时无候选点，first_num 初始值 0，未命中位置填充 0。
    # 短路返回可避免下方 cumsum[:, -1] 对空数组索引越界
    # （index -1 is out of bounds for axis 1 with size 0）。
    if N == 0:
        return np.zeros((M, B, sample_num), dtype=np.int32)

    # 输出初始化为 0：未命中时 first_num=0，填充 0
    idx = np.zeros((M, B, sample_num), dtype=np.int32)

    # 半径平方化
    min_radius2 = float(min_radius) * float(min_radius)
    max_radius2 = float(max_radius) * float(max_radius)

    # 输出位置网格（用于向量化首点填充掩码）
    pos_grid = np.arange(sample_num, dtype=np.int32)  # (sample_num,)

    # 向量化优化：按 m 维度外层循环，B×N 维度用 numpy 向量化计算
    # （替代原 M×B×N 三重纯 Python 标量循环，大 shape 性能提升数千倍）
    # 计算逻辑与原实现完全一致：距离判定、前 sample_num 个命中点收集、首点填充
    for m in range(M):
        # 距离计算（向量化）：d2[b, k] = ||center_xyz[m,b,:] - xyz[b,:,k]||^2
        # center_xyz_f[m]: (B, 3) → 广播为 (B, 3, 1)
        # xyz_f: (B, 3, N)
        diff = center_xyz_f[m, :, :, None] - xyz_f  # (B, 3, N)
        d2 = np.sum(diff * diff, axis=1)  # (B, N)

        # 距离判定：d2 == 0 || (min_radius2 <= d2 < max_radius2)
        hit = (d2 == 0.0) | ((d2 >= min_radius2) & (d2 < max_radius2))  # (B, N) bool

        # 命中点累积计数（等价于原循环提前终止：取前 sample_num 个命中点）
        cumsum = np.cumsum(hit, axis=1)  # (B, N) 命中点从1开始编号
        cnt = np.minimum(
            cumsum[:, -1], sample_num
        )  # (B,) 每批次命中数（截断到 sample_num）

        # first_num：每个批次首个命中点索引
        # 无命中时 first_num=0
        has_hit = cnt > 0
        first_idx = np.argmax(
            hit, axis=1
        )  # (B,) 每行首个 True 索引（全 False 行返回0）
        first_num = np.where(has_hit, first_idx, 0).astype(np.int32)  # (B,)

        # 收集前 sample_num 个命中点的索引（等价于原循环 cnt < sample_num 时收集）
        idx_m = np.zeros((B, sample_num), dtype=np.int32)
        selected = hit & (cumsum <= sample_num)  # (B, N) 前 sample_num 个命中点
        b_sel, k_sel = np.nonzero(selected)  # 选中的 (b, k) 坐标
        pos_sel = cumsum[b_sel, k_sel] - 1  # 对应输出位置（0-based）
        idx_m[b_sel, pos_sel] = k_sel.astype(np.int32)

        # 首点填充：命中不足时用 first_num 填充剩余位置（pos >= cnt 的位置）
        fill_mask = pos_grid[None, :] >= cnt[:, None]  # (B, sample_num) 需填充位置
        idx_m = np.where(fill_mask, first_num[:, None], idx_m)

        idx[m] = idx_m

    return idx