import numpy as np

from Model.FastForest import FastTree
from scipy.sparse import isspmatrix
from ._tree import Tree
from sklearn.tree.tree import check_array
from sklearn.tree.tree import _tree

DTYPE = _tree.DTYPE
issparse = isspmatrix


class FastForest:
    def __init__(self):
        self._tree_values_len = []
        self._tree_values = []
        self._tree_nodes_len = []
        self._tree_nodes = []

    def __getstate__(self):
        d = {
            'vl': self._tree_values_len,
            'v': self._tree_values,
            'nl': self._tree_nodes_len,
            'n': self._tree_nodes
        }
        return d

    def __setstate__(self, d):
        self._tree_values_len = d['vl']
        self._tree_values = d['v']
        self._tree_nodes_len = d['nl']
        self._tree_nodes = d['n']

    @property
    def n_estimators(self):
        return len(self._tree_nodes_len)

    def import_from_forest(self, forest):
        estimators_ = [FastTree.FastTree(e) for e in forest.estimators_
                       if (e.classes_[0] == 0 and e.n_classes_ == 2)]
        if max([len(x.tree_.nodes) for x in estimators_]) >= 32766 and \
                max([len(x._values) for x in estimators_]) >= 32766:
            print('Compress random forest error!')
            raise OverflowError

        values_array = [x._values for x in estimators_]
        self._tree_values_len = np.array([len(x) for x in values_array], np.int16)
        self._tree_values = np.array(np.concatenate(values_array), np.float32)

        tree_array = [x.tree_.nodes for x in estimators_]
        self._tree_nodes_len = np.array([len(x) for x in tree_array], np.int16)
        self._tree_nodes = np.concatenate(tree_array)

    def predict_proba(self, X):
        X = np.array(X, dtype=DTYPE)
        result = Tree().apply(X, self._tree_nodes, self._tree_nodes_len, self._tree_values_len)
        all_proba = self._tree_values.take(result)
        all_proba = all_proba.reshape((self.n_estimators, -1))

        all_proba = all_proba.sum(axis=0) / self.n_estimators
        return all_proba


class FastForest2:
    def __init__(self):
        self.estimators_ = []
        pass

    @property
    def n_estimators(self):
        return len(self.estimators_)

    def import_from_forest(self, forest):
        self.estimators_ = [FastTree.FastTree(e) for e in forest.estimators_
                            if (e.classes_[0] == 0 and e.n_classes_ == 2)]
        pass

    def __getstate__(self):
        """Getstate re-implementation, for pickling."""
        d = {}
        d['f'] = self.estimators_[0].n_features_

        values_array = [x._values for x in self.estimators_]
        values_len_array = np.array([len(x) for x in values_array], np.int16)
        d['vl'] = values_len_array
        d['vv'] = np.array(np.concatenate(values_array), np.float32)

        tree_array = [x.tree_.nodes for x in self.estimators_]
        tree_len_array = np.array([len(x) for x in tree_array], np.int16)
        d['tl'] = tree_len_array
        d['tv'] = np.concatenate(tree_array)

        return d

    def __setstate__(self, d):
        """Setstate re-implementation, for unpickling."""
        self.estimators_ = []
        vi = 0
        ti = 0
        for vl, tl in zip(d['vl'], d['tl']):
            values = d['vv'][vi:(vi + vl)]
            nodes = d['tv'][ti:(ti + tl)]
            tree = FastTree.FastTree(None)
            tree.load_from_raw(d['f'], values, nodes)
            self.estimators_.append(tree)
            vi += vl
            ti += tl

    def predict_proba(self, X):
        all_proba = np.zeros((X.shape[0]), dtype=np.float64)
        for e in self.estimators_:
            r = e.predict_proba(X)
            all_proba += r
        all_proba /= len(self.estimators_)

        return all_proba

    def test(self, X):
        result_ori = self.estimators_ori.predict_proba(X)[:, 1]
        result_new = self.estimators_.predict_proba(X)
        if np.sum(result_ori - result_new) > 0:
            print(result_ori)
            print(result_new)

        return result_new
