# Copyright (c) 2017, Apple Inc. All rights reserved.
#
# Use of this source code is governed by a BSD-3-clause license that can be
# found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause

import coremltools as ct
from coremltools import __version__ as ct_version

from ._tree_ensemble import convert_tree_ensemble as _convert_tree_ensemble


def convert(
    model,
    feature_names=None,
    target="target",
    force_32bit_float=True,
    mode="regressor",
    class_labels=None,
    n_classes=None,
):
    """
    Convert a trained XGBoost model to Core ML format.

    Parameters
    ----------
    decision_tree : Booster
        A trained XGboost tree model.

    feature_names: [str] | str
        Names of input features that will be exposed in the Core ML model
        interface.

        Can be set to one of the following:

        - ``None`` for using the feature names from the model.
        - List of names of the input features that should be exposed in the
          interface to the Core ML model. These input features are in the same
          order as the XGboost model.

    target: str
        Name of the output feature name exposed to the Core ML model.

    force_32bit_float: bool
        If ``True``, then the resulting CoreML model will use 32 bit floats internally.

    mode: str in ['regressor', 'classifier']
        Mode of the tree model.

    class_labels: list[int] or None
        List of classes. When set to None, the class labels are just the range from
        0 to ``n_classes - 1``.

    n_classes: int or None
        Number of classes in classification. When set to ``None``, the number of
        classes is expected from the model or ``class_labels`` should be provided.

    Returns
    -------
    model:MLModel
        Returns an MLModel instance representing a Core ML model.

    Examples
    --------
    .. sourcecode:: python

		# Convert it with default input and output names
   		>>> import coremltools
		>>> coreml_model = coremltools.converters.xgboost.convert(model)

		# Saving the Core ML model to a file.
		>>> coreml_model.save('my_model.mlmodel')
    """
    model = ct.models.MLModel(
        _convert_tree_ensemble(
            model,
            feature_names,
            target,
            force_32bit_float=force_32bit_float,
            mode=mode,
            class_labels=class_labels,
            n_classes=n_classes,
        )
    )

    from xgboost import __version__ as xgboost_version

    model.user_defined_metadata[ct.models._METADATA_VERSION] = ct_version
    model.user_defined_metadata[ct.models._METADATA_SOURCE] = "xgboost=={0}".format(xgboost_version)

    return model