# Licensed to Elasticsearch B.V under one or more agreements.
# Elasticsearch B.V licenses this file to you under the Apache 2.0 License.
# See the LICENSE file in the project root for more information

from typing import List, Union

import numpy as np

from eland.ml._optional import import_optional_dependency
from eland.ml._model_serializer import Tree, TreeNode, Ensemble

sklearn = import_optional_dependency("sklearn")
xgboost = import_optional_dependency("xgboost")

from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor
from sklearn.utils.validation import check_is_fitted
from xgboost import Booster, XGBRegressor, XGBClassifier


class ModelTransformer:
    def __init__(
        self,
        model,
        feature_names: List[str],
        classification_labels: List[str] = None,
        classification_weights: List[float] = None,
    ):
        self._feature_names = feature_names
        self._model = model
        self._classification_labels = classification_labels
        self._classification_weights = classification_weights

    def is_supported(self):
        return isinstance(
            self._model,
            (
                DecisionTreeClassifier,
                DecisionTreeRegressor,
                RandomForestRegressor,
                RandomForestClassifier,
                XGBClassifier,
                XGBRegressor,
                Booster,
            ),
        )


class SKLearnTransformer(ModelTransformer):
    """
    Base class for SKLearn transformers.
    warning: Should not use this class directly. Use derived classes instead
    """

    def __init__(
        self,
        model,
        feature_names: List[str],
        classification_labels: List[str] = None,
        classification_weights: List[float] = None,
    ):
        """
        Base class for SKLearn transformations

        :param model: sklearn trained model
        :param feature_names: The feature names for the model
        :param classification_labels: Optional classification labels (if not encoded in the model)
        :param classification_weights: Optional classification weights
        """
        super().__init__(
            model, feature_names, classification_labels, classification_weights
        )
        self._node_decision_type = "lte"

    def build_tree_node(self, node_index: int, node_data: dict, value) -> TreeNode:
        """
        This builds out a TreeNode class given the sklearn tree node definition.

        Node decision types are defaulted to "lte" to match the behavior of SKLearn

        :param node_index: The node index
        :param node_data: Opaque node data contained in the sklearn tree state
        :param value: Opaque node value (i.e. leaf/node values) from tree state
        :return: TreeNode object
        """
        if value.shape[0] != 1:
            raise ValueError(
                f"unexpected multiple values returned from leaf node '{node_index}'"
            )
        if node_data[0] == -1:  # is leaf node
            if (
                value.shape[1] == 1
            ):  # classification requires more than one value, so assume regression
                leaf_value = float(value[0][0])
            else:
                # the classification value, which is the index of the largest value
                leaf_value = int(np.argmax(value))
            return TreeNode(
                node_index,
                decision_type=self._node_decision_type,
                leaf_value=leaf_value,
            )
        else:
            return TreeNode(
                node_index,
                decision_type=self._node_decision_type,
                left_child=int(node_data[0]),
                right_child=int(node_data[1]),
                split_feature=int(node_data[2]),
                threshold=float(node_data[3]),
            )


class SKLearnDecisionTreeTransformer(SKLearnTransformer):
    """
    class for transforming SKLearn decision tree models into Tree model formats supported by Elasticsearch.
    """

    def __init__(
        self,
        model: Union[DecisionTreeRegressor, DecisionTreeClassifier],
        feature_names: List[str],
        classification_labels: List[str] = None,
    ):
        """
        Transforms a Decision Tree model (Regressor|Classifier) into a ES Supported Tree format
        :param model: fitted decision tree model
        :param feature_names: model feature names
        :param classification_labels: Optional classification labels
        """
        super().__init__(model, feature_names, classification_labels)

    def transform(self) -> Tree:
        """
        Transform the provided model into an ES supported Tree object
        :return: Tree object for ES storage and use
        """
        target_type = (
            "regression"
            if isinstance(self._model, DecisionTreeRegressor)
            else "classification"
        )
        check_is_fitted(self._model, ["tree_"])
        tree_classes = None
        if self._classification_labels:
            tree_classes = self._classification_labels
        if isinstance(self._model, DecisionTreeClassifier):
            check_is_fitted(self._model, ["classes_"])
            if tree_classes is None:
                tree_classes = [str(c) for c in self._model.classes_]
        nodes = list()
        tree_state = self._model.tree_.__getstate__()
        for i in range(len(tree_state["nodes"])):
            nodes.append(
                self.build_tree_node(i, tree_state["nodes"][i], tree_state["values"][i])
            )

        return Tree(self._feature_names, target_type, nodes, tree_classes)


class SKLearnForestTransformer(SKLearnTransformer):
    """
    Base class for transforming SKLearn forest models into Ensemble model formats supported by Elasticsearch.

    warning: do not use this class directly. Use a derived class instead
    """

    def __init__(
        self,
        model: Union[RandomForestClassifier, RandomForestRegressor],
        feature_names: List[str],
        classification_labels: List[str] = None,
        classification_weights: List[float] = None,
    ):
        super().__init__(
            model, feature_names, classification_labels, classification_weights
        )

    def build_aggregator_output(self) -> dict:
        raise NotImplementedError("build_aggregator_output must be implemented")

    def determine_target_type(self) -> str:
        raise NotImplementedError("determine_target_type must be implemented")

    def transform(self) -> Ensemble:
        check_is_fitted(self._model, ["estimators_"])
        estimators = self._model.estimators_
        ensemble_classes = None
        if self._classification_labels:
            ensemble_classes = self._classification_labels
        if isinstance(self._model, RandomForestClassifier):
            check_is_fitted(self._model, ["classes_"])
            if ensemble_classes is None:
                ensemble_classes = [str(c) for c in self._model.classes_]
        ensemble_models = [
            SKLearnDecisionTreeTransformer(m, self._feature_names).transform()
            for m in estimators
        ]
        return Ensemble(
            self._feature_names,
            ensemble_models,
            self.build_aggregator_output(),
            target_type=self.determine_target_type(),
            classification_labels=ensemble_classes,
            classification_weights=self._classification_weights,
        )


class SKLearnForestRegressorTransformer(SKLearnForestTransformer):
    """
    Class for transforming RandomForestRegressor models into an ensemble model supported by Elasticsearch
    """

    def __init__(self, model: RandomForestRegressor, feature_names: List[str]):
        super().__init__(model, feature_names)

    def build_aggregator_output(self) -> dict:
        return {
            "weighted_sum": {
                "weights": [1.0 / len(self._model.estimators_)]
                * len(self._model.estimators_),
            }
        }

    def determine_target_type(self) -> str:
        return "regression"


class SKLearnForestClassifierTransformer(SKLearnForestTransformer):
    """
    Class for transforming RandomForestClassifier models into an ensemble model supported by Elasticsearch
    """

    def __init__(
        self,
        model: RandomForestClassifier,
        feature_names: List[str],
        classification_labels: List[str] = None,
    ):
        super().__init__(model, feature_names, classification_labels)

    def build_aggregator_output(self) -> dict:
        return {"weighted_mode": {"num_classes": len(self._model.classes_)}}

    def determine_target_type(self) -> str:
        return "classification"


class XGBoostForestTransformer(ModelTransformer):
    """
    Base class for transforming XGBoost models into ensemble models supported by Elasticsearch

    warning: do not use directly. Use a derived classes instead
    """

    def __init__(
        self,
        model: Booster,
        feature_names: List[str],
        base_score: float = 0.5,
        objective: str = "reg:squarederror",
        classification_labels: List[str] = None,
        classification_weights: List[float] = None,
    ):
        super().__init__(
            model, feature_names, classification_labels, classification_weights
        )
        self._node_decision_type = "lt"
        self._base_score = base_score
        self._objective = objective

    def get_feature_id(self, feature_id: str) -> int:
        if feature_id[0] == "f":
            try:
                return int(feature_id[1:])
            except ValueError:
                raise RuntimeError(f"Unable to interpret '{feature_id}'")
        else:
            try:
                return int(feature_id)
            except ValueError:
                raise RuntimeError(f"Unable to interpret '{feature_id}'")

    def extract_node_id(self, node_id: str, curr_tree: int) -> int:
        t_id, n_id = node_id.split("-")
        if t_id is None or n_id is None:
            raise RuntimeError(
                f"cannot determine node index or tree from '{node_id}' for tree {curr_tree}"
            )
        try:
            t_id = int(t_id)
            n_id = int(n_id)
            if t_id != curr_tree:
                raise RuntimeError(
                    f"extracted tree id {t_id} does not match current tree {curr_tree}"
                )
            return n_id
        except ValueError:
            raise RuntimeError(
                f"cannot determine node index or tree from '{node_id}' for tree {curr_tree}"
            )

    def build_tree_node(self, row, curr_tree: int) -> TreeNode:
        node_index = row["Node"]
        if row["Feature"] == "Leaf":
            return TreeNode(node_idx=node_index, leaf_value=float(row["Gain"]))
        else:
            return TreeNode(
                node_idx=node_index,
                decision_type=self._node_decision_type,
                left_child=self.extract_node_id(row["Yes"], curr_tree),
                right_child=self.extract_node_id(row["No"], curr_tree),
                threshold=float(row["Split"]),
                split_feature=self.get_feature_id(row["Feature"]),
            )

    def build_tree(self, nodes: List[TreeNode]) -> Tree:
        return Tree(feature_names=self._feature_names, tree_structure=nodes)

    def build_base_score_stump(self) -> Tree:
        return Tree(
            feature_names=self._feature_names,
            tree_structure=[TreeNode(0, leaf_value=self._base_score)],
        )

    def build_forest(self) -> List[Tree]:
        """
        This builds out the forest of trees as described by XGBoost into a format
        supported by Elasticsearch

        :return: A list of Tree objects
        """
        if self._model.booster not in {"dart", "gbtree"}:
            raise ValueError("booster must exist and be of type dart or gbtree")

        tree_table = self._model.trees_to_dataframe()
        transformed_trees = list()
        curr_tree = None
        tree_nodes = list()
        for _, row in tree_table.iterrows():
            if row["Tree"] != curr_tree:
                if len(tree_nodes) > 0:
                    transformed_trees.append(self.build_tree(tree_nodes))
                curr_tree = row["Tree"]
                tree_nodes = list()
            tree_nodes.append(self.build_tree_node(row, curr_tree))
            # add last tree
        if len(tree_nodes) > 0:
            transformed_trees.append(self.build_tree(tree_nodes))
        # We add this stump as XGBoost adds the base_score to the regression outputs
        if self._objective.startswith("reg"):
            transformed_trees.append(self.build_base_score_stump())
        return transformed_trees

    def build_aggregator_output(self) -> dict:
        raise NotImplementedError("build_aggregator_output must be implemented")

    def determine_target_type(self) -> str:
        raise NotImplementedError("determine_target_type must be implemented")

    def is_objective_supported(self) -> bool:
        return False

    def transform(self) -> Ensemble:
        if self._model.booster not in {"dart", "gbtree"}:
            raise ValueError("booster must exist and be of type dart or gbtree")

        if not self.is_objective_supported():
            raise ValueError(f"Unsupported objective '{self._objective}'")

        forest = self.build_forest()
        return Ensemble(
            feature_names=self._feature_names,
            trained_models=forest,
            output_aggregator=self.build_aggregator_output(),
            classification_labels=self._classification_labels,
            classification_weights=self._classification_weights,
            target_type=self.determine_target_type(),
        )


class XGBoostRegressorTransformer(XGBoostForestTransformer):
    def __init__(self, model: XGBRegressor, feature_names: List[str]):
        super().__init__(
            model.get_booster(), feature_names, model.base_score, model.objective
        )

    def determine_target_type(self) -> str:
        return "regression"

    def is_objective_supported(self) -> bool:
        return self._objective in {
            "reg:squarederror",
            "reg:linear",
            "reg:squaredlogerror",
            "reg:logistic",
        }

    def build_aggregator_output(self) -> dict:
        return {"weighted_sum": {}}


class XGBoostClassifierTransformer(XGBoostForestTransformer):
    def __init__(
        self,
        model: XGBClassifier,
        feature_names: List[str],
        classification_labels: List[str] = None,
    ):
        super().__init__(
            model.get_booster(),
            feature_names,
            model.base_score,
            model.objective,
            classification_labels,
        )

    def determine_target_type(self) -> str:
        return "classification"

    def is_objective_supported(self) -> bool:
        return self._objective in {"binary:logistic", "binary:hinge"}

    def build_aggregator_output(self) -> dict:
        return {"logistic_regression": {}}
