/*******************************************************************************
* Copyright 2020 Intel 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.
*******************************************************************************/

#include <daal/include/services/error_handling.h>
#include <daal/src/algorithms/dtrees/forest/classification/df_classification_model_impl.h>
#include <daal/src/services/service_algo_utils.h>
#include <daal/include/algorithms/decision_forest/decision_forest_classification_training_batch.h>
#include <daal/include/algorithms/decision_forest/decision_forest_classification_training_types.h>
#include <daal/src/algorithms/dtrees/forest/classification/df_classification_train_kernel.h>

#include "oneapi/dal/algo/decision_forest/backend/cpu/train_kernel.hpp"

#include "oneapi/dal/table/row_accessor.hpp"
#include "oneapi/dal/backend/interop/common.hpp"
#include "oneapi/dal/backend/interop/error_converter.hpp"
#include "oneapi/dal/backend/interop/table_conversion.hpp"
#include "oneapi/dal/algo/decision_forest/backend/model_impl.hpp"

namespace oneapi::dal::decision_forest::backend {

using dal::backend::context_cpu;
using model_t = model<task::classification>;
using input_t = train_input<task::classification>;
using result_t = train_result<task::classification>;
using descriptor_t = detail::descriptor_base<task::classification>;

namespace daal_df = daal::algorithms::decision_forest;
namespace daal_df_cls_train = daal_df::classification::training;
namespace interop = dal::backend::interop;

template <typename Float, daal::CpuType Cpu>
using cls_dense_kernel_t = daal_df_cls_train::internal::
    ClassificationTrainBatchKernel<Float, daal_df_cls_train::defaultDense, Cpu>;

template <typename Float, daal::CpuType Cpu>
using cls_hist_kernel_t = daal_df_cls_train::internal::
    ClassificationTrainBatchKernel<Float, daal_df_cls_train::hist, Cpu>;

template <typename Float, template <typename, daal::CpuType> typename CpuKernel>
static result_t call_daal_kernel(const context_cpu& ctx,
                                 const descriptor_t& desc,
                                 const table& data,
                                 const table& responses,
                                 const table& weights) {
    const std::int64_t row_count = data.get_row_count();
    const std::int64_t column_count = data.get_column_count();

    const auto daal_data = interop::convert_to_daal_table<Float>(data);
    const auto daal_responses = interop::convert_to_daal_table<Float>(responses);
    const auto daal_weights = interop::convert_to_daal_table<Float>(weights);

    /* init param for daal kernel */
    auto daal_input = daal::algorithms::classifier::training::Input();
    daal_input.set(daal::algorithms::classifier::training::data, daal_data);
    daal_input.set(daal::algorithms::classifier::training::labels, daal_responses);
    daal_input.set(daal::algorithms::classifier::training::weights, daal_weights);

    auto daal_parameter = daal_df_cls_train::Parameter(
        dal::detail::integral_cast<std::size_t>(desc.get_class_count()));
    daal_parameter.nTrees = dal::detail::integral_cast<std::size_t>(desc.get_tree_count());
    daal_parameter.observationsPerTreeFraction = desc.get_observations_per_tree_fraction();
    daal_parameter.featuresPerNode =
        dal::detail::integral_cast<std::size_t>(desc.get_features_per_node());
    daal_parameter.maxTreeDepth =
        dal::detail::integral_cast<std::size_t>(desc.get_max_tree_depth());
    daal_parameter.minObservationsInLeafNode =
        dal::detail::integral_cast<std::size_t>(desc.get_min_observations_in_leaf_node());
    // TODO take engines from desc
    daal_parameter.engine = daal::algorithms::engines::mt2203::Batch<>::create(desc.get_seed());
    daal_parameter.impurityThreshold = desc.get_impurity_threshold();
    daal_parameter.memorySavingMode = desc.get_memory_saving_mode();
    daal_parameter.bootstrap = desc.get_bootstrap();
    daal_parameter.minObservationsInSplitNode =
        dal::detail::integral_cast<std::size_t>(desc.get_min_observations_in_split_node());
    daal_parameter.minWeightFractionInLeafNode = desc.get_min_weight_fraction_in_leaf_node();
    daal_parameter.minImpurityDecreaseInSplitNode = desc.get_min_impurity_decrease_in_split_node();
    daal_parameter.maxLeafNodes =
        dal::detail::integral_cast<std::size_t>(desc.get_max_leaf_nodes());
    daal_parameter.maxBins = dal::detail::integral_cast<std::size_t>(desc.get_max_bins());
    daal_parameter.minBinSize = dal::detail::integral_cast<std::size_t>(desc.get_min_bin_size());

    daal_parameter.resultsToCompute = static_cast<std::uint64_t>(desc.get_error_metric_mode());

    auto vimp = desc.get_variable_importance_mode();

    daal_parameter.varImportance = convert_to_daal_variable_importance_mode(vimp);

    auto splitter = desc.get_splitter_mode();

    daal_parameter.splitter = convert_to_daal_splitter_mode(splitter);

    result_t res;

    auto daal_result = daal_df_cls_train::Result();

    /* init daal result's objects */
    if (check_mask_flag(desc.get_error_metric_mode(), error_metric_mode::out_of_bag_error)) {
        auto arr_oob_err = array<Float>::empty(1 * 1);
        res.set_oob_err(dal::detail::homogen_table_builder{}.reset(arr_oob_err, 1, 1).build());

        const auto res_oob_err = interop::convert_to_daal_homogen_table(arr_oob_err, 1, 1);
        daal_result.set(daal_df_cls_train::outOfBagError, res_oob_err);
    }

    if (check_mask_flag(desc.get_error_metric_mode(),
                        error_metric_mode::out_of_bag_error_per_observation)) {
        auto arr_oob_per_obs_err = array<Float>::empty(row_count * 1);
        res.set_oob_err_per_observation(
            dal::detail::homogen_table_builder{}.reset(arr_oob_per_obs_err, row_count, 1).build());

        const auto res_oob_per_obs_err =
            interop::convert_to_daal_homogen_table(arr_oob_per_obs_err, row_count, 1);
        daal_result.set(daal_df_cls_train::outOfBagErrorPerObservation, res_oob_per_obs_err);
    }

    if (check_mask_flag(desc.get_error_metric_mode(),
                        error_metric_mode::out_of_bag_error_accuracy)) {
        auto arr_oob_accuracy_err = array<Float>::empty(1 * 1);
        res.set_oob_err_accuracy(
            dal::detail::homogen_table_builder{}.reset(arr_oob_accuracy_err, 1, 1).build());

        const auto res_oob_accuracy_err =
            interop::convert_to_daal_homogen_table(arr_oob_accuracy_err, 1, 1);
        daal_result.set(daal_df_cls_train::outOfBagErrorAccuracy, res_oob_accuracy_err);
    }

    if (check_mask_flag(desc.get_error_metric_mode(),
                        error_metric_mode::out_of_bag_error_decision_function)) {
        auto n_classes = desc.get_class_count();
        ONEDAL_ASSERT_MUL_OVERFLOW(std::int64_t, row_count, n_classes);
        auto arr_oob_decision_function_err = array<Float>::empty(row_count * n_classes);
        res.set_oob_err_decision_function(
            dal::detail::homogen_table_builder{}
                .reset(arr_oob_decision_function_err, row_count, desc.get_class_count())
                .build());

        const auto res_oob_decision_function_err =
            interop::convert_to_daal_homogen_table(arr_oob_decision_function_err,
                                                   row_count,
                                                   desc.get_class_count());
        daal_result.set(daal_df_cls_train::outOfBagErrorDecisionFunction,
                        res_oob_decision_function_err);
    }

    if (variable_importance_mode::none != vimp) {
        auto arr_var_imp = array<Float>::empty(1 * column_count);
        res.set_var_importance(
            dal::detail::homogen_table_builder{}.reset(arr_var_imp, 1, column_count).build());

        const auto res_var_imp =
            interop::convert_to_daal_homogen_table(arr_var_imp, 1, column_count);
        daal_result.set(daal_df_cls_train::variableImportance, res_var_imp);
    }

    daal_df::classification::ModelPtr mptr = daal_df::classification::ModelPtr(
        new daal_df::classification::internal::ModelImpl(column_count));

    interop::status_to_exception(
        interop::call_daal_kernel<Float, CpuKernel>(ctx,
                                                    daal::services::internal::hostApp(daal_input),
                                                    daal_data.get(),
                                                    daal_responses.get(),
                                                    daal_weights.get(),
                                                    *mptr,
                                                    daal_result,
                                                    daal_parameter));

    /* extract results from daal objects */
    if (check_mask_flag(desc.get_error_metric_mode(), error_metric_mode::out_of_bag_error)) {
        auto table_oob_err = interop::convert_from_daal_homogen_table<Float>(
            daal_result.get(daal_df_cls_train::outOfBagError));
        res.set_oob_err(table_oob_err);
    }

    if (check_mask_flag(desc.get_error_metric_mode(),
                        error_metric_mode::out_of_bag_error_per_observation)) {
        auto table_oob_per_obs_err = interop::convert_from_daal_homogen_table<Float>(
            daal_result.get(daal_df_cls_train::outOfBagErrorPerObservation));
        res.set_oob_err_per_observation(table_oob_per_obs_err);
    }

    if (check_mask_flag(desc.get_error_metric_mode(),
                        error_metric_mode::out_of_bag_error_accuracy)) {
        auto table_oob_err_accuracy = interop::convert_from_daal_homogen_table<Float>(
            daal_result.get(daal_df_cls_train::outOfBagErrorAccuracy));
        res.set_oob_err_accuracy(table_oob_err_accuracy);
    }

    if (check_mask_flag(desc.get_error_metric_mode(),
                        error_metric_mode::out_of_bag_error_decision_function)) {
        auto table_oob_err_decision_function = interop::convert_from_daal_homogen_table<Float>(
            daal_result.get(daal_df_cls_train::outOfBagErrorDecisionFunction));
        res.set_oob_err_decision_function(table_oob_err_decision_function);
    }

    if (variable_importance_mode::none != vimp) {
        auto table_var_imp = interop::convert_from_daal_homogen_table<Float>(
            daal_result.get(daal_df_cls_train::variableImportance));
        res.set_var_importance(table_var_imp);
    }

    auto model_impl = std::make_shared<model_impl_cls>(new model_interop_cls{ mptr });
    model_impl->tree_count = mptr->getNumberOfTrees();
    model_impl->class_count = mptr->getNumberOfClasses();

    return res.set_model(dal::detail::make_private<model_t>(model_impl));
}

template <typename Float, template <typename, daal::CpuType> typename CpuKernel>
static result_t train(const context_cpu& ctx, const descriptor_t& desc, const input_t& input) {
    return call_daal_kernel<Float, CpuKernel>(ctx,
                                              desc,
                                              input.get_data(),
                                              input.get_responses(),
                                              input.get_weights());
}

template <typename Float, typename Task>
struct train_kernel_cpu<Float, method::dense, Task> {
    result_t operator()(const context_cpu& ctx,
                        const descriptor_t& desc,
                        const input_t& input) const {
        return train<Float, cls_dense_kernel_t>(ctx, desc, input);
    }
};

template <typename Float, typename Task>
struct train_kernel_cpu<Float, method::hist, Task> {
    result_t operator()(const context_cpu& ctx,
                        const descriptor_t& desc,
                        const input_t& input) const {
        return train<Float, cls_hist_kernel_t>(ctx, desc, input);
    }
};

template struct train_kernel_cpu<float, method::dense, task::classification>;
template struct train_kernel_cpu<float, method::hist, task::classification>;
template struct train_kernel_cpu<double, method::dense, task::classification>;
template struct train_kernel_cpu<double, method::hist, task::classification>;

} // namespace oneapi::dal::decision_forest::backend
