/*
 * Copyright (c) 2019-2020, 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 <cuml/matrix/kernelparams.h>
#include <cuml/common/logger.hpp>
#include <cuml/cuml.hpp>
#include "svm_model.h"
#include "svm_parameter.h"

namespace ML {
namespace SVM {

// Forward declarations of the stateless API
/**
 * @brief Fit a support vector classifier to the training data.
 *
 * Each row of the input data stores a feature vector.
 * We use the SMO method to fit the SVM.
 *
 * The output device buffers in model shall be unallocated on entry.
 *
 * @tparam math_t floating point type
 * @param [in] handle the cuML handle
 * @param [in] input device pointer for the input data in column major format.
 *   Size n_rows x n_cols.
 * @param [in] n_rows number of rows
 * @param [in] n_cols number of columns
 * @param [in] labels device pointer for the labels. Size [n_rows].
 * @param [in] param parameters for training
 * @param [in] kernel_params parameters for the kernel function
 * @param [out] model parameters of the trained model
 * @param [in] sample_weight optional sample weights, size [n_rows]
 */
template <typename math_t>
void svcFit(const cumlHandle &handle, math_t *input, int n_rows, int n_cols,
            math_t *labels, const svmParameter &param,
            MLCommon::Matrix::KernelParams &kernel_params,
            svmModel<math_t> &model, const math_t *sample_weight = nullptr);

/**
 * @brief Predict classes or decision function value for samples in input.
 *
 * We evaluate the decision function f(x_i). Depending on the parameter
 * predict_class, we either return f(x_i) or the label corresponding to
 * sign(f(x_i)).
 *
 * The predictions are calculated according to the following formulas:
 * \f[
 *    f(x_i) = \sum_{j=1}^n_support K(x_i, x_j) * dual_coefs[j] + b)
 * \f]
 *
 * pred(x_i) = label[sign(f(x_i))], if predict_class==true, or
 * pred(x_i) = f(x_i),       if predict_class==falsee.
 *
 * @tparam math_t floating point type
 * @param handle the cuML handle
 * @param [in] input device pointer for the input data in column major format,
 *   size [n_rows x n_cols].
 * @param [in] n_rows number of rows (input vectors)
 * @param [in] n_cols number of columns (features)
 * @param [in] kernel_params parameters for the kernel function
 * @param [in] model SVM model parameters
 * @param [out] preds device pointer to store the predicted class labels.
 *    Size [n_rows]. Should be allocated on entry.
 * @param [in] buffer_size size of temporary buffer in MiB
 * @param [in] predict_class whether to predict class label (true), or just
 *     return the decision function value (false)
 */
template <typename math_t>
void svcPredict(const cumlHandle &handle, math_t *input, int n_rows, int n_cols,
                MLCommon::Matrix::KernelParams &kernel_params,
                const svmModel<math_t> &model, math_t *preds,
                math_t buffer_size, bool predict_class = true);

/**
 * Deallocate device buffers in the svmModel struct.
 *
 * @param [in] handle cuML handle
 * @param [inout] m SVM model parameters
 */
template <typename math_t>
void svmFreeBuffers(const cumlHandle &handle, svmModel<math_t> &m);

/**
 * @brief C-Support Vector Classification
 *
 * This is a Scikit-Learn like wrapper around the stateless C++ functions.
 * See Issue #456 for general discussion about stateful Sklearn like wrappers.
 *
 * The classifier will be fitted using the SMO algorithm in dual space.
 *
 * The decision function takes the following form
 * \f[
 *    sign\left( \sum_{i=1}^{N_{support}} y_i \alpha_i K(x_i,x) + b \right),
 * \f]
 * where \f$x_i\f$ are the support vectors, and \f$ y_i \alpha_i \f$ are the dual
 * coordinates.
 *
 * The penalty parameter C limits the values of the dual coefficients
 * \f[ 0 <= \alpha <= C \f]
 *
 */
template <typename math_t>
class SVC {
 public:
  // Public members for easier access during testing from Python.

  MLCommon::Matrix::KernelParams kernel_params;
  svmParameter param;
  svmModel<math_t> model;
  /**
   * @brief Constructs a support vector classifier
   * @param handle cuML handle
   * @param C penalty term
   * @param tol tolerance to stop fitting
   * @param kernel_params parameters for kernels
   * @param cache_size size of kernel cache in device memory (MiB)
   * @param max_iter maximum number of outer iterations in SmoSolver
   * @param nochange_steps number of steps with no change wrt convergence
   * @param verbosity verbosity level for logging messages during execution
   */
  SVC(cumlHandle &handle, math_t C = 1, math_t tol = 1.0e-3,
      MLCommon::Matrix::KernelParams kernel_params =
        MLCommon::Matrix::KernelParams{MLCommon::Matrix::LINEAR, 3, 1, 0},
      math_t cache_size = 200, int max_iter = -1, int nochange_steps = 1000,
      int verbosity = CUML_LEVEL_INFO);

  ~SVC();

  /**
   * @brief Fit a support vector classifier to the training data.
   *
   * Each row of the input data stores a feature vector.
   * We use the SMO method to fit the SVM.
   *
   * @param input device pointer for the input data in column major format. Size n_rows x n_cols.
   * @param n_rows number of rows
   * @param n_cols number of columns
   * @param labels device pointer for the labels. Size n_rows.
   * @param [in] sample_weight optional sample weights, size [n_rows]
   */
  void fit(math_t *input, int n_rows, int n_cols, math_t *labels,
           const math_t *sample_weight = nullptr);

  /**
   * @brief Predict classes for samples in input.
   * @param [in]  input device pointer for the input data in column major format,
   *   size [n_rows x n_cols].
   * @param [in] n_rows number of vectors
   * @param [in] n_cols number of features
   * @param [out] preds device pointer to store the predicted class labels.
   *    Size [n_rows]. Should be allocated on entry.
   */
  void predict(math_t *input, int n_rows, int n_cols, math_t *preds);

  /**
   * @brief Calculate decision function value for samples in input.
   * @param [in] input device pointer for the input data in column major format,
   *   size [n_rows x n_cols].
   * @param [in] n_rows number of vectors
   * @param [in] n_cols number of features
   * @param [out] preds device pointer to store the decision function value
   *    Size [n_rows]. Should be allocated on entry.
   */
  void decisionFunction(math_t *input, int n_rows, int n_cols, math_t *preds);

 private:
  const cumlHandle &handle;
};

};  // end namespace SVM
};  // end namespace ML
