import numpy as np
from scipy.sparse import isspmatrix
from sklearn.tree.tree import _tree
from sklearn.tree.tree import check_array

from ._tree import Tree

DTYPE = _tree.DTYPE
issparse = isspmatrix


class FastTree:
    def __init__(self, tree_ori):
        if tree_ori is None:
            self.n_features_ = 0
            self._values = np.zeros(0, np.float64)
            self.tree_ = Tree()
            return

        assert tree_ori.classes_[0] == 0.
        assert tree_ori.n_outputs_ == 1
        self.n_features_ = tree_ori.n_features_
        # self.tree_ori_ = tree_ori.tree_

        self._values = np.zeros(0, np.float64)
        self._left_child = np.zeros(0, np.int16)
        self._right_child = np.zeros(0, np.int16)
        self._feature = np.zeros(0, np.uint16)

        tree_state = tree_ori.tree_.__getstate__()
        # Optimize tree_state
        tree_nodes = tree_state['nodes']
        tree_values = self.optimize_values(tree_state['values'])
        self.optimize_tree_nodes(tree_nodes, tree_values)

        self.tree_ = Tree()
        self.tree_.import_tree(self._left_child, self._right_child, self._feature)

        del self._left_child
        del self._right_child
        del self._feature

    def load_from_raw(self, n_features, values, nodes):
        self.n_features_ = n_features
        self._values = values
        self.tree_.__setstate__(nodes)
        pass

    def _validate_X_predict(self, X, check_input):
        """Validate X whenever one tries to predict, apply, predict_proba"""
        if check_input:
            X = check_array(X, dtype=DTYPE, accept_sparse="csr")
            if issparse(X) and (X.indices.dtype != np.intc or
                                X.indptr.dtype != np.intc):
                raise ValueError("No support for np.int64 index based "
                                 "sparse matrices")

        n_features = X.shape[1]
        if self.n_features_ != n_features:
            raise ValueError("Number of features of the model must "
                             "match the input. Model n_features is %s and "
                             "input n_features is %s "
                             % (self.n_features_, n_features))

        return X

    def optimize_values(self, values):
        values_new = values[:, 0, 1] / np.sum(values, axis=2)[:, 0]
        return values_new

    def optimize_tree_nodes(self, nodes, values):
        values_id_table = np.zeros(nodes.shape, np.int32) - 1
        shift_table = np.zeros(nodes.shape, np.int32)
        j = 0
        for i, n in enumerate(nodes):
            if n[0] < 0:
                # Is node
                values_id_table[i] = j
                shift_table[i:] += 1
                j += 1

        # Optimize values
        values_new = np.zeros(j, np.float64)
        for i, j in enumerate(values_id_table):
            values_new[j] = values[i]

        # Optimize nodes
        left_child = []
        right_child = []
        feature = []
        for i, n in enumerate(nodes):
            c = n['left_child']
            if c < 0:
                # Is node
                continue

            if nodes[c][0] < 0:
                # Is node
                c_new = -values_id_table[c]
            else:
                c_new = c - shift_table[c]
            left_child.append(c_new)

            c = n['right_child']
            if nodes[c][0] < 0:
                # Is node
                c_new = -values_id_table[c]
            else:
                c_new = c - shift_table[c]
            right_child.append(c_new)

            feature.append(n['feature'])

        assert len(values_new) < 32767  # We use int16
        self._values = np.array(values_new, np.float64)
        self._left_child = np.array(left_child, np.int16)
        self._right_child = np.array(right_child, np.int16)
        self._feature = np.array(feature, np.int16)
        pass

    def predict_proba(self, X, check_input=True):
        X = self._validate_X_predict(X, check_input)

        result_new = self.tree_.predict(X)
        proba_new = self._values.take(result_new, axis=0, mode='clip')

        """
        proba = self.tree_ori_.predict(X)
        normalizer = proba.sum(axis=1)[:, np.newaxis]
        normalizer[normalizer == 0.0] = 1.0
        proba /= normalizer
        proba_ori = proba[:, 1]

        if np.sum(proba_new - proba_ori) > 0.:
            print(proba_new)
            print(proba_ori)
        """

        return proba_new
