from sasctl.utils.pyml2ds.basic import TreeParser
from .core import EnsembleParser


class PmmlTreeParser(TreeParser):
    """Class for parsing pmml gradient boosting tree."""

    def _not_leaf(self):
        return self._node.get('defaultChild')

    def _get_var(self):
        return self._node.find('Node').find('SimplePredicate').get('field')

    def _go_left(self):
        return self._node.find('Node').get('id') == self._node.get(
            'defaultChild')

    def _go_right(self):
        return not self._node.find('Node').get('id') == self._node.get(
            'defaultChild')

    def _left_node(self):
        return self._node.findall('Node')[0]

    def _right_node(self):
        return self._node.findall('Node')[1]

    def _missing_node(self):
        return None

    def _split_value(self):
        return self._node.find('Node').find('SimplePredicate').get('value')

    def _decision_type(self):
        ops = {
            'lessThan': '<',
            'lessOrEqual': '<=',
            'greaterThan': '>',
            'greaterOrEqual': '>='
        }

        return ops[self._node.find('Node').find('SimplePredicate').get('operator')]

    def _leaf_value(self):
        return self._node.get('score')


class PmmlParser(EnsembleParser):
    """Class for parsing pmml gradient boosting models.

    Parameters
    ----------
    tree_root : etree.Element
        Root node of pmml gradient boosting forest.

    """
    def __init__(self, tree_root):
        super(PmmlParser, self).__init__()

        self._tree_root = tree_root
        for elem in tree_root.getiterator():
            if hasattr(elem.tag, 'find'):
                i = elem.tag.find('}')
                if i >= 0:
                    elem.tag = elem.tag[i + 1:]

        self._forest = tree_root.find('MiningModel/Segmentation')[0].find('MiningModel')

        rescaleConstant = self._forest.find('Targets/Target').get('rescaleConstant')
        self.out_transform = "1 / (1 + exp(-{}))".format("({0} + " + "{})".format(rescaleConstant))

        self._tree_parser = PmmlTreeParser()

    def _iter_trees(self):
        for booster_id, tree_elem in enumerate(self._forest.find('Segmentation')):
            yield booster_id, tree_elem.find('TreeModel/Node')
