"""
IBM Confidential
OCO Source Materials
5737-H76, 5725-W78, 5900-A1R
(c) Copyright IBM Corp. 2020 All Rights Reserved.
The source code for this program is not published or otherwise divested of its trade secrets,
irrespective of what has been deposited with the U.S. Copyright Office.
"""
from __future__ import print_function
import joblib
import logging

from sklearn.base import is_classifier
from sklearn.experimental import enable_hist_gradient_boosting
from sklearn.ensemble import HistGradientBoostingRegressor
from sklearn.ensemble import HistGradientBoostingClassifier
from sklearn.ensemble._hist_gradient_boosting.predictor import TreePredictor
from sklearn.ensemble._hist_gradient_boosting.common import PREDICTOR_RECORD_DTYPE

from sklearn.ensemble._hist_gradient_boosting.loss import BinaryCrossEntropy, \
    CategoricalCrossEntropy, LeastSquares

logger = logging.getLogger(__name__)


def export_sklearn(model, full_path=None):
    """
    Auxillary function to first convert the FL XGBoost Model to Scikit Learn's
    Histogram Gradient Boosting model object, then persist the object as a
    pickled model object. This enables standalone use of the trained model
    without dependency to FL's libraries.

    Parameters/Features that are Not Supported or Transferred:
    - monotonic_cst (in v0.21)
    - warm_start
    - early_stopping
    - scoring
    - validation_fraction
    - n_iter_no_change
    - tol

    :param model: The FL XGBoost model object to convert from.
    :type  model: `XGBFLModel`
    :param full_path: The full absolute path of where the model will persist.
    :type  full_path: `str`
    :return: `None`
    """
    # Initialize Model Object (Parameter Transfer from FL Model to Sklearn)
    if model.loss == 'least_squares':
        export_model = HistGradientBoostingRegressor(
            loss=model.loss,
            learning_rate=model.learning_rate,
            max_iter=model.max_iter,
            max_leaf_nodes=model.max_leaf_nodes,
            max_depth=model.max_depth,
            min_samples_leaf=model.min_samples_leaf,
            l2_regularization=model.l2_regularization,
            max_bins=model.max_bins,
            verbose=model.verbose,
            random_state=model.random_state)
    elif model.loss == 'categorical_crossentropy' or \
            model.loss == 'binary_crossentropy' or \
            model.loss == 'auto':
        export_model = HistGradientBoostingClassifier(
            loss=model.loss,
            learning_rate=model.learning_rate,
            max_iter=model.max_iter,
            max_leaf_nodes=model.max_leaf_nodes,
            max_depth=model.max_depth,
            min_samples_leaf=model.min_samples_leaf,
            l2_regularization=model.l2_regularization,
            max_bins=model.max_bins,
            verbose=model.verbose,
            random_state=model.random_state)

    # Attribute Transfer
    export_model._baseline_prediction = model._baseline_prediction
    export_model.n_trees_per_iteration_ = model.n_trees
    export_model.n_features_ = model.n_features_

    # Model Predictor Object (Deep Object Rebuild)
    export_model._predictors = []
    for i, pred in enumerate(model._predictors):
        export_model._predictors.append([])
        for j, p in enumerate(pred):
            nodes = model._predictors[i][j].nodes.astype(
                PREDICTOR_RECORD_DTYPE).copy()
            export_model._predictors[i].append(TreePredictor(nodes))

    # Model Loss Function
    if hasattr(model, 'classes_'):
        if model.loss == 'binary_crossentropy' or len(model.classes_) == 2:
            export_model.loss_ = BinaryCrossEntropy(None)
        elif model.loss == 'categorical_crossentropy':
            export_model.loss_ = CategoricalCrossEntropy(None)
    elif model.loss == 'least_squares':
        export_model.loss_ = LeastSquares(None)

    # Target Class Encoding (For Classification Models)
    if hasattr(model, 'classes_'):
        export_model.classes_ = model.classes_

    # Perform Model Pickling
    with open(full_path, 'wb') as f:
        joblib.dump(export_model, f)
