# SPDX-License-Identifier: Apache-2.0

import numpy

from ..common._container import LightGbmModelContainer
from ..common._topology import Topology
from ..common.data_types import (
    FloatTensorType,
    SequenceType,
    DictionaryType,
    StringType,
    Int64Type,
)

from lightgbm import LGBMClassifier, LGBMRegressor

lightgbm_classifier_list = [LGBMClassifier]

# Associate scikit-learn types with our operator names.
# If two scikit-learn models share a single name, it means their
# are equivalent in terms of conversion.
lightgbm_operator_name_map = {
    LGBMClassifier: "LgbmClassifier",
    LGBMRegressor: "LgbmRegressor",
}


class WrappedBooster:
    def __init__(self, booster):
        self.booster_ = booster
        self.n_features_ = self.booster_.feature_name()
        self.objective_ = self.get_objective()
        if self.objective_.startswith("binary"):
            self.operator_name = "LgbmClassifier"
            self.classes_ = self._generate_classes(booster)
        elif self.objective_.startswith("multiclass"):
            self.operator_name = "LgbmClassifier"
            self.classes_ = self._generate_classes(booster)
        elif self.objective_.startswith("regression"):
            self.operator_name = "LgbmRegressor"
        else:
            raise NotImplementedError(
                "Unsupported LightGbm objective: %r." % self.objective_
            )
        try:
            average_output = self.booster_.attr("average_output")
        except AttributeError:
            average_output = self.booster_.params.get("average_output", None)
        if average_output:
            self.boosting_type = "rf"
        else:
            # Other than random forest, other boosting types do not affect later conversion.
            # Here `gbdt` is chosen for no reason.
            self.boosting_type = "gbdt"

    @staticmethod
    def _generate_classes(booster):
        if isinstance(booster, dict):
            num_class = booster["num_class"]
        else:
            try:
                num_class = booster.attr("num_class")
            except AttributeError:
                num_class = booster.params.get("num_class", None)
        if num_class is None:
            dp = booster.dump_model(num_iteration=1)
            num_class = dp["num_class"]
        if num_class == 1:
            return numpy.asarray([0, 1])
        return numpy.arange(num_class)

    def get_objective(self):
        "Returns the objective."
        if hasattr(self, "objective_") and self.objective_ is not None:
            return self.objective_
        try:
            objective = self.booster_.attr("objective")
        except AttributeError:
            objective = self.booster_.params.get("objective", None)
        if objective is not None:
            return objective
        dp = self.booster_.dump_model(num_iteration=1)
        return dp["objective"]


def _get_lightgbm_operator_name(model):
    """
    Get operator name of the input argument

    :param model:  A lightgbm object.
    :return: A string which stands for the type of the input model in our conversion framework
    """
    if isinstance(model, WrappedBooster):
        return model.operator_name
    model_type = type(model)
    if model_type not in lightgbm_operator_name_map:
        raise ValueError("No proper operator name found for '%s'" % model_type)
    return lightgbm_operator_name_map[model_type]


def _parse_lightgbm_simple_model(scope, model, inputs, split=None):
    """
    This function handles all non-pipeline models.

    :param scope: Scope object
    :param model: A lightgbm object
    :param inputs: A list of variables
    :param split: split TreeEnsembleRegressor into multiple node to reduce
        discrepancies
    :return: A list of output variables which will be passed to next stage
    """
    operator_name = _get_lightgbm_operator_name(model)
    this_operator = scope.declare_local_operator(operator_name, model)
    this_operator.split = split
    this_operator.inputs = inputs

    if operator_name == "LgbmClassifier":
        # For classifiers, we may have two outputs, one for
        # label and the other one for probabilities of all classes.
        # Notice that their types here are not necessarily correct
        # and they will be fixed in shape inference phase
        label_variable = scope.declare_local_variable("lgbmlabel", FloatTensorType())
        probability_map_variable = scope.declare_local_variable(
            "lgbmprobabilities", FloatTensorType()
        )
        this_operator.outputs.append(label_variable)
        this_operator.outputs.append(probability_map_variable)
    else:
        # We assume that all scikit-learn operator can only produce a single float tensor.
        variable = scope.declare_local_variable("variable", FloatTensorType())
        this_operator.outputs.append(variable)
    return this_operator.outputs


def _parse_sklearn_classifier(scope, model, inputs, zipmap=True):
    probability_tensor = _parse_lightgbm_simple_model(scope, model, inputs)
    this_operator = scope.declare_local_operator("LgbmZipMap")
    this_operator.inputs = probability_tensor
    this_operator.zipmap = zipmap

    classes = model.classes_
    label_type = Int64Type()

    if isinstance(model.classes_, list) and isinstance(
        model.classes_[0], numpy.ndarray
    ):
        # multi-label problem
        # this_operator.classlabels_int64s = list(range(0, len(classes)))
        raise NotImplementedError("multi-label is not supported")
    elif numpy.issubdtype(model.classes_.dtype, numpy.floating):
        classes = numpy.array(list(map(lambda x: int(x), classes)))
        if set(map(lambda x: float(x), classes)) != set(model.classes_):
            raise RuntimeError(
                "skl2onnx implicitly converts float class "
                "labels into integers but at least one label "
                "is not an integer. Class labels should "
                "be integers or strings."
            )
        this_operator.classlabels_int64s = classes
    elif numpy.issubdtype(model.classes_.dtype, numpy.signedinteger):
        this_operator.classlabels_int64s = classes
    else:
        classes = numpy.array([s.encode("utf-8") for s in classes])
        this_operator.classlabels_strings = classes
        label_type = StringType()

    output_label = scope.declare_local_variable("label", label_type)
    if zipmap:
        output_probability = scope.declare_local_variable(
            "probabilities", SequenceType(DictionaryType(label_type, FloatTensorType()))
        )
    else:
        output_probability = scope.declare_local_variable(
            "probabilities", FloatTensorType(shape=[None, len(classes)])
        )
    this_operator.outputs.append(output_label)
    this_operator.outputs.append(output_probability)
    return this_operator.outputs


def _parse_lightgbm(scope, model, inputs, zipmap=True, split=None):
    """
    This is a delegate function. It doesn't nothing but
    invoke the correct parsing function according to the input
    model's type.

    :param scope: Scope object
    :param model: A lightgbm object
    :param inputs: A list of variables
    :param zipmap: add operator ZipMap after operator TreeEnsembleClassifier
    :param split: split TreeEnsembleRegressor into multiple node to reduce
        discrepancies
    :return: The output variables produced by the input model
    """
    if isinstance(model, LGBMClassifier):
        return _parse_sklearn_classifier(scope, model, inputs, zipmap=zipmap)
    if isinstance(model, WrappedBooster) and model.operator_name == "LgbmClassifier":
        return _parse_sklearn_classifier(scope, model, inputs, zipmap=zipmap)
    return _parse_lightgbm_simple_model(scope, model, inputs, split=split)


def parse_lightgbm(
    model,
    initial_types=None,
    target_opset=None,
    custom_conversion_functions=None,
    custom_shape_calculators=None,
    zipmap=True,
    split=None,
):
    raw_model_container = LightGbmModelContainer(model)
    topology = Topology(
        raw_model_container,
        default_batch_size="None",
        initial_types=initial_types,
        target_opset=target_opset,
        custom_conversion_functions=custom_conversion_functions,
        custom_shape_calculators=custom_shape_calculators,
    )
    scope = topology.declare_scope("__root__")

    inputs = []
    for var_name, initial_type in initial_types:
        inputs.append(scope.declare_local_variable(var_name, initial_type))

    for variable in inputs:
        raw_model_container.add_input(variable)

    outputs = _parse_lightgbm(scope, model, inputs, zipmap=zipmap, split=split)

    for variable in outputs:
        raw_model_container.add_output(variable)

    return topology
