# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
import numpy as np


def gather_nd(data_np, indices_np, batch_dims=0):
    """gather_nd implemented using numpy"""
    data_shape = data_np.shape
    indices_shape = indices_np.shape

    def gather_nd_batch_dims_1_ref(data, indices):
        res = []
        for i, row in enumerate(data):
            indices_tuple = tuple(indices[:, i])  # the indices for the i-th batch
            res.append(row[indices_tuple])
        # stack on the batch dim
        return np.stack(res, 0)

    if batch_dims > 1:
        data_np_reshape = np.reshape(data_np, (-1,) + data_shape[batch_dims:])
        indices_np_reshape = np.reshape(
            indices_np, (indices_shape[0], -1) + indices_shape[(batch_dims + 1) :]
        )

        ref_res = gather_nd_batch_dims_1_ref(data_np_reshape, indices_np_reshape)

        out_shape = indices_shape[1 : (batch_dims + 1)] + ref_res.shape[1:]
        ref_res = np.reshape(ref_res, out_shape)
    elif batch_dims == 1:
        ref_res = gather_nd_batch_dims_1_ref(data_np, indices_np)
    else:
        ref_res = data_np[tuple(indices_np)]

    return ref_res