/*!
 * Copyright (c) 2017-2023 by Contributors
 * \file model_loader.h
 * \brief Functions for loading tree models from disk or file stream
 * \author Hyunsu Cho
 */
#ifndef TREELITE_MODEL_LOADER_H_
#define TREELITE_MODEL_LOADER_H_

#include <cstddef>
#include <memory>
#include <string>

namespace treelite {

class Model;  // forward declaration

namespace model_loader {

//--------------------------------------------------------------------------
// Model loaders for XGBoost and LightGBM
//--------------------------------------------------------------------------
/*!
 * \brief Load a model file generated by XGBoost (dmlc/xgboost), stored in the legacy binary
 *        format.
 * \param filename Name of model file
 * \return Loaded model
 */
std::unique_ptr<treelite::Model> LoadXGBoostModelLegacyBinary(std::string const& filename);
/*!
 * \brief Load an XGBoost model from a memory buffer.
 * \param buf memory buffer
 * \param len size of memory buffer
 * \return Loaded model
 */
std::unique_ptr<treelite::Model> LoadXGBoostModelLegacyBinary(void const* buf, std::size_t len);
/*!
 * \brief Load a model file generated by XGBoost (dmlc/xgboost), stored in the JSON format.
 * \param filename Name of model file
 * \param config_json JSON string consisting of key-value pairs; used for configuring the model
 *                    parser
 * \return Loaded model
 */
std::unique_ptr<treelite::Model> LoadXGBoostModel(
    std::string const& filename, char const* config_json);
/*!
 * \brief Load an XGBoost model from a JSON string
 * \param json_str JSON char array
 * \param length Length of JSON char array
 * \param config_json JSON string consisting of key-value pairs; used for configuring the model
 *                    parser
 * \return Loaded model
 */
std::unique_ptr<treelite::Model> LoadXGBoostModelFromString(
    char const* json_str, std::size_t length, char const* config_json);
/*!
 * \brief load a model file generated by LightGBM (Microsoft/LightGBM). The
 *        model file must contain a decision tree ensemble.
 * \param filename name of model file
 * \return loaded model
 */
std::unique_ptr<treelite::Model> LoadLightGBMModel(std::string const& filename);
/*!
 * \brief Load a LightGBM model from a string. The string should be created with the
 *        model_to_string() method in LightGBM.
 * \param model_str the model string
 * \return loaded model
 */
std::unique_ptr<treelite::Model> LoadLightGBMModelFromString(char const* model_str);

//--------------------------------------------------------------------------
// Model loaders for Scikit-Learn models
//--------------------------------------------------------------------------
namespace sklearn {
/*!
 * \brief Load a scikit-learn RandomForestRegressor model from a collection of arrays. Refer to
 *        https://scikit-learn.org/stable/auto_examples/tree/plot_unveil_tree_structure.html to
 *        learn the meaning of the arrays in detail. Note that this function can also be used to
 *        load an ensemble of extremely randomized trees (sklearn.ensemble.ExtraTreesRegressor).
 * \param n_estimators Number of trees in the random forest
 * \param n_features Number of features in the training data
 * \param n_targets Number of targets (outputs)
 * \param node_count node_count[i] stores the number of nodes in the i-th tree
 * \param children_left children_left[i][k] stores the ID of the left child node of node k of the
 *                      i-th tree. This is only defined if node k is an internal (non-leaf) node.
 * \param children_right children_right[i][k] stores the ID of the right child node of node k of the
 *                       i-th tree. This is only defined if node k is an internal (non-leaf) node.
 * \param feature feature[i][k] stores the ID of the feature used in the binary tree split at node k
 *                of the i-th tree. This is only defined if node k is an internal (non-leaf) node.
 * \param threshold threshold[i][k] stores the threshold used in the binary tree split at node k of
 *                  the i-th tree. This is only defined if node k is an internal (non-leaf) node.
 * \param value value[i][k] stores the leaf output of node k of the i-th tree. This is only defined
 *              if node k is a leaf node.
 * \param n_node_samples n_node_samples[i][k] stores the number of data samples associated with
 *                       node k of the i-th tree.
 * \param weighted_n_node_samples weighted_n_node_samples[i][k] stores the sum of weighted data
 *                                samples associated with node k of the i-th tree.
 * \param impurity impurity[i][k] stores the impurity measure (gini, entropy etc) associated with
 *                 node k of the i-th tree.
 * \return Loaded model
 */
std::unique_ptr<treelite::Model> LoadRandomForestRegressor(int n_estimators, int n_features,
    int n_targets, std::int64_t const* node_count, std::int64_t const** children_left,
    std::int64_t const** children_right, std::int64_t const** feature, double const** threshold,
    double const** value, std::int64_t const** n_node_samples,
    double const** weighted_n_node_samples, double const** impurity);
/*!
 * \brief Load a scikit-learn IsolationForest model from a collection of arrays. Refer to
 *        https://scikit-learn.org/stable/auto_examples/tree/plot_unveil_tree_structure.html to
 *        learn the meaning of the arrays in detail.
 * \param n_estimators Number of trees in the isolation forest
 * \param n_features Number of features in the training data
 * \param node_count node_count[i] stores the number of nodes in the i-th tree
 * \param children_left children_left[i][k] stores the ID of the left child node of node k of the
 *                      i-th tree. This is only defined if node k is an internal (non-leaf) node.
 * \param children_right children_right[i][k] stores the ID of the right child node of node k of the
 *                       i-th tree. This is only defined if node k is an internal (non-leaf) node.
 * \param feature feature[i][k] stores the ID of the feature used in the binary tree split at node k
 *                of the i-th tree. This is only defined if node k is an internal (non-leaf) node.
 * \param threshold threshold[i][k] stores the threshold used in the binary tree split at node k of
 *                  the i-th tree. This is only defined if node k is an internal (non-leaf) node.
 * \param value value[i][k] stores the expected isolation depth of node k of the i-th tree. This is
 *              only defined if node k is a leaf node.
 * \param n_node_samples n_node_samples[i][k] stores the number of data samples associated with
 *                       node k of the i-th tree.
 * \param weighted_n_node_samples weighted_n_node_samples[i][k] stores the sum of weighted data
 *                                samples associated with node k of the i-th tree.
 * \param impurity Not used, but must be passed as array of arrays for each tree and node.
 * \param ratio_c Standardizing constant to use for calculation of the anomaly score.
 * \return Loaded model
 */
std::unique_ptr<treelite::Model> LoadIsolationForest(int n_estimators, int n_features,
    std::int64_t const* node_count, std::int64_t const** children_left,
    std::int64_t const** children_right, std::int64_t const** feature, double const** threshold,
    double const** value, std::int64_t const** n_node_samples,
    double const** weighted_n_node_samples, double const** impurity, double ratio_c);
/*!
 * \brief Load a scikit-learn RandomForestClassifier model from a collection of arrays. Refer to
 *        https://scikit-learn.org/stable/auto_examples/tree/plot_unveil_tree_structure.html to
 *        learn the meaning of the arrays in detail. Note that this function can also be used to
 *        load an ensemble of extremely randomized trees (sklearn.ensemble.ExtraTreesClassifier).
 * \param n_estimators Number of trees in the random forest
 * \param n_features Number of features in the training data
 * \param n_targets Number of targets (outputs)
 * \param n_classes n_classes[i] stores the number of classes in the i-th target
 * \param node_count node_count[i] stores the number of nodes in the i-th tree
 * \param children_left children_left[i][k] stores the ID of the left child node of node k of the
 *                      i-th tree. This is only defined if node k is an internal (non-leaf) node.
 * \param children_right children_right[i][k] stores the ID of the right child node of node k of the
 *                       i-th tree. This is only defined if node k is an internal (non-leaf) node.
 * \param feature feature[i][k] stores the ID of the feature used in the binary tree split at node k
 *                of the i-th tree. This is only defined if node k is an internal (non-leaf) node.
 * \param threshold threshold[i][k] stores the threshold used in the binary tree split at node k of
 *                  the i-th tree. This is only defined if node k is an internal (non-leaf) node.
 * \param value value[i][k] stores the leaf output of node k of the i-th tree. This is only defined
 *              if node k is a leaf node.
 * \param n_node_samples n_node_samples[i][k] stores the number of data samples associated with
 *                       node k of the i-th tree.
 * \param weighted_n_node_samples weighted_n_node_samples[i][k] stores the sum of weighted data
 *                                samples associated with node k of the i-th tree.
 * \param impurity impurity[i][k] stores the impurity measure (gini, entropy etc) associated with
 *                 node k of the i-th tree.
 * \return Loaded model
 */
std::unique_ptr<treelite::Model> LoadRandomForestClassifier(int n_estimators, int n_features,
    int n_targets, int32_t const* n_classes, std::int64_t const* node_count,
    std::int64_t const** children_left, std::int64_t const** children_right,
    std::int64_t const** feature, double const** threshold, double const** value,
    std::int64_t const** n_node_samples, double const** weighted_n_node_samples,
    double const** impurity);
/*!
 * \brief Load a scikit-learn GradientBoostingRegressor model from a collection of arrays. Refer
 *        to https://scikit-learn.org/stable/auto_examples/tree/plot_unveil_tree_structure.html to
 *        learn the meaning of the arrays in detail.
 *        Note: GradientBoostingRegressor does not support multiple targets (outputs).
 * \param n_iter Number of boosting iterations
 * \param n_features Number of features in the training data
 * \param node_count node_count[i] stores the number of nodes in the i-th tree
 * \param children_left children_left[i][k] stores the ID of the left child node of node k of the
 *                      i-th tree. This is only defined if node k is an internal (non-leaf) node.
 * \param children_right children_right[i][k] stores the ID of the right child node of node k of the
 *                       i-th tree. This is only defined if node k is an internal (non-leaf) node.
 * \param feature feature[i][k] stores the ID of the feature used in the binary tree split at node k
 *                of the i-th tree. This is only defined if node k is an internal (non-leaf) node.
 * \param threshold threshold[i][k] stores the threshold used in the binary tree split at node k of
 *                  the i-th tree. This is only defined if node k is an internal (non-leaf) node.
 * \param value value[i][k] stores the leaf output of node k of the i-th tree. This is only defined
 *              if node k is a leaf node.
 * \param n_node_samples n_node_samples[i][k] stores the number of data samples associated with
 *                       node k of the i-th tree.
 * \param weighted_n_node_samples weighted_n_node_samples[i][k] stores the sum of weighted data
 *                                samples associated with node k of the i-th tree.
 * \param impurity impurity[i][k] stores the impurity measure (gini, entropy etc) associated with
 *                 node k of the i-th tree.
 * \param baseline_prediction Baseline predictions for outputs. At prediction, margin scores will be
 *                            adjusted by this amount before applying the post-processing (link)
 *                            function. Required shape: (1,)
 * \return Loaded model
 */
std::unique_ptr<treelite::Model> LoadGradientBoostingRegressor(int n_iter, int n_features,
    std::int64_t const* node_count, std::int64_t const** children_left,
    std::int64_t const** children_right, std::int64_t const** feature, double const** threshold,
    double const** value, std::int64_t const** n_node_samples,
    double const** weighted_n_node_samples, double const** impurity,
    double const* baseline_prediction);
/*!
 * \brief Load a scikit-learn GradientBoostingClassifier model from a collection of arrays. Refer
 *        to https://scikit-learn.org/stable/auto_examples/tree/plot_unveil_tree_structure.html to
 *        learn the meaning of the arrays in detail.
 *        Note: GradientBoostingClassifier does not support multiple targets (outputs).
 * \param n_iter Number of boosting iterations
 * \param n_features Number of features in the training data
 * \param n_classes Number of classes in the target variable
 * \param node_count node_count[i] stores the number of nodes in the i-th tree
 * \param children_left children_left[i][k] stores the ID of the left child node of node k of the
 *                      i-th tree. This is only defined if node k is an internal (non-leaf) node.
 * \param children_right children_right[i][k] stores the ID of the right child node of node k of the
 *                       i-th tree. This is only defined if node k is an internal (non-leaf) node.
 * \param feature feature[i][k] stores the ID of the feature used in the binary tree split at node k
 *                of the i-th tree. This is only defined if node k is an internal (non-leaf) node.
 * \param threshold threshold[i][k] stores the threshold used in the binary tree split at node k of
 *                  the i-th tree. This is only defined if node k is an internal (non-leaf) node.
 * \param value value[i][k] stores the leaf output of node k of the i-th tree. This is only defined
 *              if node k is a leaf node.
 * \param n_node_samples n_node_samples[i][k] stores the number of data samples associated with
 *                       node k of the i-th tree.
 * \param weighted_n_node_samples weighted_n_node_samples[i][k] stores the sum of weighted data
 *                                samples associated with node k of the i-th tree.
 * \param impurity impurity[i][k] stores the impurity measure (gini, entropy etc) associated with
 *                 node k of the i-th tree.
 * \param baseline_prediction Baseline predictions for outputs. At prediction, margin scores will be
 *                            adjusted by this amount before applying the post-processing (link)
 *                            function. Required shape: (n_classes,)
 * \return Loaded model
 */
std::unique_ptr<treelite::Model> LoadGradientBoostingClassifier(int n_iter, int n_features,
    int n_classes, std::int64_t const* node_count, std::int64_t const** children_left,
    std::int64_t const** children_right, std::int64_t const** feature, double const** threshold,
    double const** value, std::int64_t const** n_node_samples,
    double const** weighted_n_node_samples, double const** impurity,
    double const* baseline_prediction);

/*!
 * \brief Load a scikit-learn HistGradientBoostingRegressor model from a collection of arrays.
 *        Note: HistGradientBoostingRegressor does not support multiple targets (outputs).
 * \param n_iter Number of boosting iterations
 * \param n_features Number of features in the training data
 * \param node_count node_count[i] stores the number of nodes in the i-th tree
 * \param nodes nodes[i][k] stores the k-th node of the i-th tree.
 * \param expected_sizeof_node_struct Expected size of Node struct, in bytes
 * \param n_categorical_splits n_categorical_splits[i] stores the number of categorical splits
 *                             in the i-th tree.
 * \param raw_left_cat_bitsets raw_left_cat_bitsets[i][k] stores the bitmaps for node k of tree i.
 *                             The bitmaps are used to represent categorical tests.
 *                             Shape of raw_left_cat_bitsets[i]: (n_categorical_splits, 8)
 * \param known_cat_bitsets Bitsets representing the list of known categories per categorical
 *                          feature. Shape: (n_categorical_features, 8)
 * \param known_cat_bitsets_offset_map Map from an original feature index to the corresponding
 *                                     index in the known_cat_bitsets array. Shape: (n_features,)
 * \param features_map Mapping to re-order features. This is needed because HistGradientBoosting
 *                     estimator internally re-orders features using ColumnTransformer so that
 *                     the categorical features come before the numerical features.
 * \param categories_map Mapping to transform categorical features. This is needed because
 *                       HistGradientBoosting estimator embeds an OrdinalEncoder.
 *                       categories_map[i] represents the mapping for i-th categorical feature.
 * \param base_scores Baseline predictions for outputs. At prediction, margin scores will be
 *                    adjusted by this amount before applying the post-processing (link)
 *                    function. Required shape: (1,)
 * \return Loaded model
 */
std::unique_ptr<treelite::Model> LoadHistGradientBoostingRegressor(int n_iter, int n_features,
    std::int64_t const* node_count, void const** nodes, int expected_sizeof_node_struct,
    std::uint32_t n_categorical_splits, std::uint32_t const** raw_left_cat_bitsets,
    std::uint32_t const* known_cat_bitsets, std::uint32_t const* known_cat_bitsets_offset_map,
    std::int32_t const* features_map, std::int64_t const** categories_map,
    double const* base_scores);

/*!
 * \brief Load a scikit-learn HistGradientBoostingClassifier model from a collection of arrays.
 *        Note: HistGradientBoostingClassifier does not support multiple targets (outputs).
 * \param n_iter Number of boosting iterations
 * \param n_features Number of features in the training data
 * \param n_classes Number of classes in the target variable
 * \param node_count node_count[i] stores the number of nodes in the i-th tree
 * \param nodes nodes[i][k] stores the k-th node of the i-th tree.
 * \param expected_sizeof_node_struct Expected size of Node struct, in bytes
 * \param n_categorical_splits n_categorical_splits[i] stores the number of categorical splits
 *                             in the i-th tree.
 * \param raw_left_cat_bitsets raw_left_cat_bitsets[i][k] stores the bitmaps for node k of tree i.
 *                             The bitmaps are used to represent categorical tests.
 *                             Shape of raw_left_cat_bitsets[i]: (n_categorical_splits, 8)
 *
 * \param known_cat_bitsets Bitsets representing the list of known categories per categorical
 *                          feature. Shape: (n_categorical_features, 8)
 * \param known_cat_bitsets_offset_map Map from an original feature index to the corresponding
 *                                     index in the known_cat_bitsets array. Shape: (n_features,)
 * \param features_map Mapping to re-order features. This is needed because HistGradientBoosting
 *                     estimator internally re-orders features using ColumnTransformer so that
 *                     the categorical features come before the numerical features.
 * \param categories_map Mapping to transform categorical features. This is needed because
 *                       HistGradientBoosting estimator embeds an OrdinalEncoder.
 *                       categories_map[i] represents the mapping for i-th categorical feature.
 * \param base_scores Baseline predictions for outputs. At prediction, margin scores will be
 *                    adjusted by this amount before applying the post-processing (link)
 *                    function. Required shape: (1,) for binary classification;
 *                    (n_classes,) for multi-class classification
 * \return Loaded model
 */
std::unique_ptr<treelite::Model> LoadHistGradientBoostingClassifier(int n_iter, int n_features,
    int n_classes, int64_t const* node_count, void const** nodes, int expected_sizeof_node_struct,
    std::uint32_t n_categorical_splits, std::uint32_t const** raw_left_cat_bitsets,
    std::uint32_t const* known_cat_bitsets, std::uint32_t const* known_cat_bitsets_offset_map,
    std::int32_t const* features_map, std::int64_t const** categories_map,
    double const* base_scores);

}  // namespace sklearn

}  // namespace model_loader
}  // namespace treelite

#endif  // TREELITE_MODEL_LOADER_H_
