from sasctl.utils.pyml2ds.basic import TreeParser

from .core import EnsembleParser


class LightgbmTreeParser(TreeParser):
    """Class for parsing lightgbm tree."""

    def _not_leaf(self):
        return "split_feature" in self._node

    def _get_var(self):
        return self._features[self._node["split_feature"]]

    def _go_left(self):
        return self._node["default_left"]

    def _go_right(self):
        return not self._node["default_left"]

    def _left_node(self):
        return self._node["left_child"]

    def _right_node(self):
        return self._node["right_child"]

    def _missing_node(self):
        return None

    def _split_value(self):
        return self._node["threshold"]

    def _decision_type(self):
        return self._node["decision_type"]

    def _leaf_value(self):
        return self._node["leaf_value"]


class LightgbmParser(EnsembleParser):
    """Class for parsing lightgbm model.

    Parameters
    ----------
    booster : lightgbm.basic.Booster
        Booster of lightgbm model.

    """

    def __init__(self, booster):
        super(LightgbmParser, self).__init__()

        self._booster = booster
        self._dump = booster.dump_model()

        objective = self._dump.get("objective")

        if objective != "binary sigmoid:1":
            raise ValueError(
                "Only binary sigmoid objective function is "
                "currently supported. Received '%s'." % objective
            )

        self._features = self._dump["feature_names"]
        self.out_transform = "1 / (1 + exp(-{0}))"

        self._tree_parser = LightgbmTreeParser()
        self._tree_parser._features = self._features

    def _iter_trees(self):
        for tree in self._dump["tree_info"]:
            yield tree["tree_index"], tree["tree_structure"]
