/*
 * Copyright (c) 2019-2021, NVIDIA CORPORATION.
 *
 * Licensed 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.
 */

#pragma once

#include <raft/handle.hpp>
#include <vector>

#include <opg/matrix/data.hpp>
#include <opg/matrix/part_descriptor.hpp>

namespace ML {
namespace KNN {
namespace opg {

using namespace MLCommon;

/**
 * Performs a multi-node multi-GPU KNN.
 * @param[in] handle the raft::handle_t to use for managing resources
 * @param[out] out_I vector of output index partitions. size should match the
 *        number of local input partitions.
 * @param[out] out_D vector of output distance partitions. size should match
 *        the number of local input partitions.
 * @param[in] idx_data vector of local indices to query
 * @param[in] idx_desc describes how the index partitions are distributed
 *        across the ranks.
 * @param[in] query_data vector of local query partitions
 * @param[in] query_desc describes how the query partitions are distributed
 *        across the cluster.
 * @param[in] rowMajorIndex boolean indicating whether the index is row major.
 * @param[in] rowMajorQuery boolean indicating whether the query is row major.
 * @param[in] k the number of neighbors to query
 * @param[in] batch_size the max number of rows to broadcast at a time
 * @param[in] verbose print extra logging info
 */
void knn(raft::handle_t &handle, std::vector<Matrix::Data<int64_t> *> *out_I,
         std::vector<Matrix::floatData_t *> *out_D,
         std::vector<Matrix::floatData_t *> &idx_data,
         Matrix::PartDescriptor &idx_desc,
         std::vector<Matrix::floatData_t *> &query_data,
         Matrix::PartDescriptor &query_desc, bool rowMajorIndex,
         bool rowMajorQuery, int k, size_t batch_size, bool verbose);

/**
 * Performs a multi-node multi-GPU KNN classify.
 * @param[in] handle the raft::handle_t to use for managing resources
 * @param[out] out vector of output labels partitions. size should match the
 *        number of local input partitions.
 * @param[in] probas (optional) pointer to a vector containing arrays of probabilities
 * @param[in] idx_data vector of local indices to query
 * @param[in] idx_desc describes how the index partitions are distributed
 *        across the ranks.
 * @param[in] query_data vector of local query partitions
 * @param[in] query_desc describes how the query partitions are distributed
 *        across the cluster.
 * @param[in] y vector of vector of label arrays. for multilabel classification, each
 *          element in the vector is a different "output" array of labels corresponding
 *          to the i'th output. size should match the number of local input partitions.
 * @param[in] uniq_labels vector of the sorted unique labels for each array in y
 * @param[in] n_unique vector of sizes for each array in uniq_labels
 * @param[in] rowMajorIndex boolean indicating whether the index is row major. 
 * @param[in] rowMajorQuery boolean indicating whether the query is row major. 
 * @param[in] probas_only return probas instead of performing complete knn_classify
 * @param[in] k the number of neighbors to query
 * @param[in] batch_size the max number of rows to broadcast at a time
 * @param[in] verbose print extra logging info
 */
void knn_classify(raft::handle_t &handle, std::vector<Matrix::Data<int> *> *out,
                  std::vector<std::vector<float *>> *probas,
                  std::vector<Matrix::floatData_t *> &idx_data,
                  Matrix::PartDescriptor &idx_desc,
                  std::vector<Matrix::floatData_t *> &query_data,
                  Matrix::PartDescriptor &query_desc,
                  std::vector<std::vector<int *>> &y,
                  std::vector<int *> &uniq_labels, std::vector<int> &n_unique,
                  bool rowMajorIndex = false, bool rowMajorQuery = false,
                  bool probas_only = false, int k = 10,
                  size_t batch_size = 1 << 15, bool verbose = false);

/**
 * Performs a multi-node multi-GPU KNN regress.
 * @param[in] handle the raft::handle_t to use for managing resources
 * @param[out] out vector of output partitions. size should match the
 *        number of local input partitions.
 * @param[in] idx_data vector of local indices to query
 * @param[in] idx_desc describes how the index partitions are distributed
 *        across the ranks.
 * @param[in] query_data vector of local query partitions
 * @param[in] query_desc describes how the query partitions are distributed
 *        across the cluster.
 * @param[in] y vector of vector of output arrays. for multi-output regression, each
 *          element in the vector is a different "output" array corresponding
 *          to the i'th output. size should match the number of local input partitions.
 * @param[in] rowMajorIndex boolean indicating whether the index is row major.
 * @param[in] rowMajorQuery boolean indicating whether the query is row major.
 * @param[in] k the number of neighbors to query
 * @param[in] n_outputs number of outputs
 * @param[in] batch_size the max number of rows to broadcast at a time
 * @param[in] verbose print extra logging info
 */
void knn_regress(raft::handle_t &handle,
                 std::vector<Matrix::Data<float> *> *out,
                 std::vector<Matrix::floatData_t *> &idx_data,
                 Matrix::PartDescriptor &idx_desc,
                 std::vector<Matrix::floatData_t *> &query_data,
                 Matrix::PartDescriptor &query_desc,
                 std::vector<std::vector<float *>> &y, bool rowMajorIndex,
                 bool rowMajorQuery, int k, int n_outputs, size_t batch_size,
                 bool verbose);

};  // END namespace opg
};  // namespace KNN
};  // namespace ML
