import heapq
import importlib
import json
import random
import warnings
from abc import ABC, abstractmethod
from collections.abc import Sequence
from importlib.metadata import pass_none
from itertools import count

import numpy as np
import sklearn
from graphviz import Source
from sklearn import tree

from ruletree.base.RuleTreeBase import RuleTreeBase
from ruletree.tree.RuleTreeNode import RuleTreeNode
from ruletree.base.RuleTreeBaseStump import RuleTreeBaseStump
from ruletree.utils.data_utils import json_NumpyEncoder
from ruletree.utils.define import DATA_TYPE_IMAGE, DATA_TYPE_TABULAR, DATA_TYPE_TS


class RuleTree(RuleTreeBase, ABC):
    def __init__(self,
                 max_leaf_nodes,
                 min_samples_split,
                 max_depth,
                 prune_useless_leaves,
                 base_stumps, #isinstance of RuleTreeBaseStump, list of instances of RuleTreeBaseStump,
                                # list of tuple (probability, RuleTreeBaseStump) where sum(probabilities)==1
                 stump_selection, # ['random', 'best']
                 random_state,
                 ):
        self.tiebreaker = count()
        self.root:RuleTreeNode = None
        self.queue = list()
        self.classes_ = None
        self.n_classes_ = None

        self.max_leaf_nodes = float("inf") if max_leaf_nodes is None else max_leaf_nodes
        self.min_samples_split = min_samples_split
        self.max_depth = max_depth
        self.prune_useless_leaves = prune_useless_leaves

        self.base_stumps = base_stumps
        self.stump_selection = stump_selection

        self.random_state = random_state
        random.seed(random_state)

    def _set_stump(self):
        # base stump type checking
        class_to_check = self._get_stumps_base_class()
        _base_stump = []
        _p = []
        if isinstance(self.base_stumps, RuleTreeBaseStump):
            assert isinstance(self.base_stumps, class_to_check)
            _base_stump.append(self.base_stumps)
            _p.append(1.)
        elif type(self.base_stumps) == list:
            assert len(self.base_stumps) > 0
            _equal_p = 1 / len(self.base_stumps)
            for el in self.base_stumps:
                if isinstance(el, Sequence):
                    assert isinstance(el[1], RuleTreeBaseStump)
                    assert isinstance(el[1], class_to_check)
                    _p.append(el[0])
                    _base_stump.append(el[1])
                else:
                    assert isinstance(el, RuleTreeBaseStump)
                    assert isinstance(el, class_to_check)
                    _p.append(_equal_p)
                    _base_stump.append(el)

        assert sum(_p) == 1.
        self.base_stumps = [(p, stump) for p, stump in zip(np.cumsum(_p), _base_stump)]

    def _get_random_stump(self, X):
        val = random.random()

        supported_stumps = self._filter_types(X)

        for p, clf in supported_stumps:
            if val <= p:
                return sklearn.clone(clf)

    def _filter_types(self, X):
        if len(X.shape) == 2:
            data_type = DATA_TYPE_TABULAR
        elif len(X.shape) == 3:
            data_type = DATA_TYPE_TS
        elif len(X.shape) == 4:
            data_type = DATA_TYPE_IMAGE
        else:
            raise TypeError(f"Invalid data type for shape {X.shape}")

        compatible_stumps = [(p, x) for p, x in self.base_stumps if x.supports(data_type)]

        if len(compatible_stumps) == 0:
            raise TypeError(f"No compatible stumps found for "
                            f"shape {X.shape}.\r\n{[x.__name__ for _, x in self.base_stumps]}")

        p_total = sum([x for x, _ in compatible_stumps])
        if p_total < 1:
            compatible_stumps = [(p/p_total, x) for p, x in compatible_stumps]

        return compatible_stumps


    def fit(self, X: np.array, y: np.array = None, **kwargs):
        self.classes_ = np.unique(y)
        self.n_classes_ = len(self.classes_)
        self.n_features = X.shape[1]
        self._set_stump()

        idx = np.arange(X.shape[0])

        self.root = self.prepare_node(y, idx, "R")

        self.queue_push(self.root, idx)

        nbr_curr_nodes = 0
        while len(self.queue) > 0 and nbr_curr_nodes + len(self.queue) < self.max_leaf_nodes:
            idx, current_node = self.queue_pop()

            if len(idx) < self.min_samples_split:
                self.make_leaf(current_node)
                nbr_curr_nodes += 1
                continue

            if nbr_curr_nodes + len(self.queue) + 1 >= self.max_leaf_nodes:
                self.make_leaf(current_node)
                nbr_curr_nodes += 1
                continue

            if self.max_depth is not None and current_node.get_depth() >= self.max_depth:
                self.make_leaf(current_node)
                nbr_curr_nodes += 1
                continue

            if self.check_additional_halting_condition(y=y, curr_idx=idx):
                self.make_leaf(current_node)
                nbr_curr_nodes += 1
                continue

            clf = self.make_split(X, y, idx=idx, **kwargs)
            labels = clf.apply(X[idx])

            if self.is_split_useless(X=X, clf=clf, idx=idx):
                self.make_leaf(current_node)
                nbr_curr_nodes += 1
                continue

            idx_l, idx_r = idx[labels == 1], idx[labels == 2]

            current_node.set_stump(clf)
            current_node.node_l = self.prepare_node(y, idx_l, current_node.node_id + "l", )
            current_node.node_r = self.prepare_node(y, idx_r, current_node.node_id + "r", )
            current_node.node_l.parent, current_node.node_r.parent = current_node, current_node
            current_node.balance_score = (np.min(np.unique(labels, return_counts= True)[1]) / labels.shape[0])

            self.queue_push(current_node.node_l, idx_l)
            self.queue_push(current_node.node_r, idx_r)

        if self.prune_useless_leaves:
            self.root = self.root.simplify()

        self._post_fit_fix()

        return self

    def predict(self, X: np.ndarray):
        labels, leaves, proba = self._predict(X, self.root)

        return labels

    def apply(self, X: np.ndarray):
        labels, leaves, proba = self._predict(X, self.root)

        return leaves

    def predict_proba(self, X: np.ndarray):
        labels, leaves, proba = self._predict(X, self.root)
        proba_matrix = np.zeros((X.shape[0], self.n_classes_))
        for classe in self.classes_:
            proba_matrix[labels == classe, self.classes_ == classe] = proba[labels == classe]

        return proba_matrix

    def _predict(self, X: np.ndarray, current_node: RuleTreeNode):
        if current_node.is_leaf():
            n = len(X)
            return np.array([current_node.prediction] * n), \
                np.array([current_node.node_id] * n), \
                np.array([current_node.prediction_probability] * n)

        else:
            labels, leaves, proba = (
                np.full(len(X), fill_value=-1,
                        dtype=object if type(current_node.prediction) is str else type(current_node.prediction)),
                np.zeros(len(X), dtype=object),
                np.ones(len(X), dtype=float) * -1
            )

            clf = current_node.stump
            labels_clf = clf.apply(X)
            X_l, X_r = X[labels_clf == 1], X[labels_clf == 2]
            if X_l.shape[0] != 0:
                labels[labels_clf == 1], leaves[labels_clf == 1], proba[labels_clf == 1] = self._predict(X_l,
                                                                                                         current_node.node_l)
            if X_r.shape[0] != 0:
                labels[labels_clf == 2], leaves[labels_clf == 2], proba[labels_clf == 2] = self._predict(X_r,
                                                                                                         current_node.node_r)

            return labels, leaves, proba

    def get_rules(self):
        return self.root.get_rule()

    def make_leaf(self, node: RuleTreeNode) -> RuleTreeNode:
        return node

    def queue_pop(self):
        el = heapq.heappop(self.queue)
        return el[-2:]

    def check_additional_halting_condition(self, y, curr_idx: np.ndarray):
        return False

    def _post_fit_fix(self):
        return

    @abstractmethod
    def make_split(self, X: np.ndarray, y, idx: np.ndarray, **kwargs) -> tree:
        pass

    @abstractmethod
    def prepare_node(self, y: np.ndarray, idx: np.ndarray, node_id: str) -> RuleTreeNode:
        pass

    @abstractmethod
    def queue_push(self, node: RuleTreeNode, idx: np.ndarray):
        pass

    @abstractmethod
    def is_split_useless(self, X: np.ndarray, clf: tree, idx: np.ndarray):
        pass

    @abstractmethod
    def _get_stumps_base_class(self):
        return RuleTreeBaseStump

    def _get_tree_paths(self, current_node = None):
        if current_node.is_leaf():
            return [[current_node.node_id]]
        
        if current_node is None:
            current_node = self.root
                

        if current_node.node_l:  
            left_paths = self._get_tree_paths(current_node.node_l)
            right_paths = self._get_tree_paths(current_node.node_r)
            
            for path in left_paths:
                path.append(current_node.node_id) 
                
            for path in right_paths:
                path.append(current_node.node_id)
                
            paths = left_paths + right_paths
            
        return paths
    
    def _node_dict(self, current_node=None, d=None, i=0):
        if d is None:
            d = {}
            
        if current_node is None:
            current_node = self.root
    
        # Assign an index to the current node
        if current_node.node_id not in d:
            d[current_node.node_id] = i
            i += 1
    
        if not current_node.is_leaf():
            i = self._node_dict(current_node.node_l, d, i)[-1]
            i = self._node_dict(current_node.node_r, d, i)[-1]
    
        return d, i

    @abstractmethod
    def _get_prediction_probas(self, current_node, probas=None):
        pass

    def _tree_value(self):
        probas = self._get_prediction_probas(self.root)
        if isinstance(probas[0], (list, np.ndarray)):
            return np.array(probas).reshape(-1, 1, len(probas[0]))
        else:
            return np.array(probas).reshape(-1, 1, 1)
            
    def _tree_feature(self, current_node = None, feats=None):
        if feats is None:
            feats = []
            
        if current_node is None:
            current_node = self.root
        
        if current_node.is_leaf():
            feats.append(-2)
        else:
            feats.append(current_node.stump.feature_original[0])
            self._tree_feature(current_node.node_l, feats = feats)
            self._tree_feature(current_node.node_r, feats = feats)
        
        return feats
        
    
   
   
    def _compute_importances(self, current_node=None, importances=None):
       if importances is None:
           importances = np.zeros(self.n_features, dtype=np.float64)

       if current_node is None:
           current_node = self.root

       if not current_node.is_leaf():
           #following the implementation of https://github.com/scikit-learn/scikit-learn/blob/main/sklearn/tree/_tree.pyx#L1251
           feature = current_node.stump.feature_original[0]
           
           imp_parent, imp_child_l, imp_child_r = current_node.stump.tree_.impurity
           n_parent, n_child_l, n_child_r = current_node.stump.tree_.weighted_n_node_samples
           
           info_gain = (
               n_parent * imp_parent
               - n_child_l * imp_child_l
               - n_child_r * imp_child_r
           )
           
           info_gain 
           
           importances[feature] += info_gain

           # Recur for left and right children
           self._compute_importances(current_node.node_l, importances)
           self._compute_importances(current_node.node_r, importances)

       return importances 
   
    def compute_feature_importances(self, normalize=True):
        
        importances = self._compute_importances()
        root_weighted_samples = self.root.stump.tree_.weighted_n_node_samples[0]
        importances /= root_weighted_samples
        
        if normalize:
            total = np.sum(importances)
            if total > 0.0:
                importances /= total

        return importances
        
    def local_interpretation(self, X, joint_contribution = False):
        
        node_dict = self._node_dict()[0] #-> Turn 'R': 0, 'Rl' : 1, 'Rr' :2 and so on

        leaves = np.array([node_dict[x] for x in self.apply(X)])
        paths = [[node_dict[x] for x in path] for path in self._get_tree_paths(self.root)]
        
        for path in paths:
            path.reverse()
        
        leaf_to_path = {}
        
        for path in paths:
            leaf_to_path[path[-1]] = path
            
        values = self._tree_value().squeeze(axis = 1)
        
        return leaves, paths, leaf_to_path, values
        
    
    def eval_contributions(self, 
                           leaves, 
                           paths, 
                           leaf_to_path, 
                           values, 
                           biases,
                           line_shape,
                           joint_contribution = False):
        
        direct_prediction = values[leaves]
        values_list = list(values)
        feature_index = list(self._tree_feature())
        
        contributions = []
        if joint_contribution:
            for row, leaf in enumerate(leaves):
                path = leaf_to_path[leaf]
                
                
                path_features = set()
                contributions.append({})
                for i in range(len(path) - 1):
                    path_features.add(feature_index[path[i]])
                    contrib = values_list[path[i+1]] - \
                             values_list[path[i]]
                    #path_features.sort()
                    contributions[row][tuple(sorted(path_features))] = \
                        contributions[row].get(tuple(sorted(path_features)), 0) + contrib
            return direct_prediction, biases, contributions
        
        else:
            unique_leaves = np.unique(leaves)
            unique_contributions = {}
            
            for row, leaf in enumerate(unique_leaves):
                for path in paths:
                    if leaf == path[-1]:
                        break
                
                contribs = np.zeros(line_shape)
                for i in range(len(path) - 1):
                    
                    contrib = values_list[path[i+1]] - \
                             values_list[path[i]]
                    contribs[feature_index[path[i]]] += contrib
                unique_contributions[leaf] = contribs
                
            for row, leaf in enumerate(leaves):
                contributions.append(unique_contributions[leaf])

            return direct_prediction, biases, np.array(contributions)

    @classmethod
    def print_rules(cls, rules: dict, columns_names: list = None, ndigits=2, indent: int = 0, ):
        #columns_names now included in stump get_rule(), should we keep it here?
        #ndigits/float precision included in stump get_rule(), should we keep it here?
      
        indentation = "".join(["|   " for _ in range(indent)])
        
        if rules["is_leaf"]:
            pred = rules['prediction']
            print(f"{indentation} output: "
                  f"{pred if type(pred) in [np.str_, np.string_, str] else round(pred, ndigits=ndigits)}")
            
        else:
            print(f"{indentation}|--- {rules['textual_rule']}")
            cls.print_rules(rules=rules['left_node'], columns_names=columns_names, indent=indent + 1)
            print(f"{indentation}|--- {rules['not_textual_rule']}")
            cls.print_rules(rules=rules['right_node'], columns_names=columns_names, indent=indent + 1)
    
    
    @classmethod
    def print_rules_old(cls, rules: dict, columns_names: list = None, ndigits=2, indent: int = 0, ):
        names = lambda x: f"X_{x}"
        names_pivots = lambda x: f"P_{x}"

        if columns_names is not None:
            names = lambda x: columns_names[x]

        indentation = "".join(["|   " for _ in range(indent)])

        if rules["is_leaf"]:
            pred = rules['prediction']

            print(f"{indentation} output: "
                  f"{pred if type(pred) in [np.str_, np.string_, str] else round(pred, ndigits=ndigits)}")
        else:
            comparison = "==" if rules['is_categorical'] else "<="
            not_comparison = "!=" if rules['is_categorical'] else ">"
            feature_idx = rules['feature_idx']
            coefficients = rules['coefficients'] if 'coefficients' in rules else None
            thr = rules['threshold']
            
            print(feature_idx)

            if isinstance(feature_idx, (int, np.integer)):
                print(f"{indentation}|--- {rules['textual_rule']}")
                cls.print_rules(rules=rules['left_node'], columns_names=columns_names, indent=indent + 1)

                print(f"{indentation}|--- {rules['not_textual_rule']}")
                cls.print_rules(rules=rules['right_node'], columns_names=columns_names, indent=indent + 1)
                
            elif (isinstance(feature_idx, str) and feature_idx.endswith('_P')): #TODO: cambia
                print(f"{indentation}|--- {names_pivots(feature_idx[:-2])} {comparison} "
                      f"{thr if isinstance(thr, (np.str_, np.string_, str)) else round(thr, ndigits=ndigits)}"
                      f"\t{rules['samples']}")
                cls.print_rules(rules=rules['left_node'], columns_names=columns_names, indent=indent + 1)
            
                print(f"{indentation}|--- {names_pivots(feature_idx[:-2])} {not_comparison} "
                      f"{thr if isinstance(thr, (np.str_, np.string_, str)) else round(thr, ndigits=ndigits)}")
                cls.print_rules(rules=rules['right_node'], columns_names=columns_names, indent=indent + 1)

            
            elif isinstance(feature_idx, tuple) and all(isinstance(i, (int, np.integer)) for i in feature_idx):
                
                print(f"{indentation}|--- near {names_pivots(feature_idx[0])}"
                      
                      f"\t{rules['samples']}")
                cls.print_rules(rules=rules['left_node'], columns_names=columns_names, indent=indent + 1)

                print(f"{indentation}|--- near {names_pivots(feature_idx[1])}"
                     )
                cls.print_rules(rules=rules['right_node'], columns_names=columns_names, indent=indent + 1)
                
            elif isinstance(feature_idx, list) and all(isinstance(i, (int, np.integer)) for i in feature_idx):
                feature_coefficient_pairs = [f"{round(coef, ndigits=ndigits)} * {names(idx)}" for coef, idx in zip(coefficients, feature_idx)]
                feature_coefficient_str = " + ".join(feature_coefficient_pairs)
                
                print(f"{indentation}|--- {feature_coefficient_str} {comparison} "
                      f"{thr if isinstance(thr, (str, np.str_, np.string_)) else round(thr, ndigits=ndigits)}"
                      f"\t{rules['samples']}")
                cls.print_rules(rules=rules['left_node'], columns_names=columns_names, indent=indent + 1)
                
                print(f"{indentation}|--- {feature_coefficient_str} {not_comparison} "
                      f"{thr if isinstance(thr, (str, np.str_, np.string_)) else round(thr, ndigits=ndigits)}")
                cls.print_rules(rules=rules['right_node'], columns_names=columns_names, indent=indent + 1)
                
            elif isinstance(feature_idx, list) and feature_idx[0].endswith('_P'):
                feature_coefficient_pairs = [f"{round(coef, ndigits=ndigits)} * {names_pivots(idx[:-2])}" for coef, idx in zip(coefficients, feature_idx)]
                feature_coefficient_str = " + ".join(feature_coefficient_pairs)
                
                print(f"{indentation}|--- {feature_coefficient_str} {comparison} "
                      f"{thr if isinstance(thr, (str, np.str_, np.string_)) else round(thr, ndigits=ndigits)}"
                      f"\t{rules['samples']}")
                cls.print_rules(rules=rules['left_node'], columns_names=columns_names, indent=indent + 1)
                
                print(f"{indentation}|--- {feature_coefficient_str} {not_comparison} "
                      f"{thr if isinstance(thr, (str, np.str_, np.string_)) else round(thr, ndigits=ndigits)}")
                cls.print_rules(rules=rules['right_node'], columns_names=columns_names, indent=indent + 1)
                
            else: #if list -> oblique
                raise Exception("Unimplemented")

    @classmethod
    def decode_ruletree(cls, vector):

        #need to check if n_classes_ is actually necessary
        
        #n_classes_ = np.array([n_classes_], dtype=np.intp)
        
        idx_to_node = {index: RuleTreeNode(node_id=None,
                                        prediction=None, prediction_probability=None, parent=-1) 
                       for index in range(len(vector[0]))}
        
        idx_to_node[0].node_id = 'R'
        idx_to_node[0].parent = -1
        
        return idx_to_node
                            
    def encode_ruletree(self):
        nodes = (2 ** (self.max_depth + 1)) - 1
        vector = np.zeros((2, nodes), dtype=object)
        
        index = {'R': 0}  # root index
        parent = {}
        
        if not hasattr(self, 'root'):
            raise ValueError('This RuleTree instance must be fitted before encoding.')
        else:
            root_node = self.root
        root_node.encode_node(index, parent, vector, self)
        
        
        for node in range(vector.shape[1]):
            if vector[0][node] == -1:
                # children update
                if ((2*node + 1) < (vector.shape[1] - 1)) and ((2*node + 2) < vector.shape[1]):
                    vector[0][2*node + 1] = -1
                    vector[1][2*node + 1] = vector[1][node]
                    vector[0][2*node + 2] = -1
                    vector[1][2*node + 2] = vector[1][node]

                    parent[2*node + 1] = node
                    parent[2*node + 2] = node

                    # update the node itself
                    vector[0][node] = vector[0][parent[node]]
                    vector[1][node] = vector[1][parent[node]]
                    
        self.vector = vector
        
        return vector


    def to_dict(self, filename=None):
        node_list = [self.root]

        args = {
            "max_leaf_nodes": self.max_leaf_nodes,
            "min_samples_split": self.min_samples_split,
            "max_depth": self.max_depth,
            "prune_useless_leaves": self.prune_useless_leaves,
            "base_stumps": None,
            "stump_selection": self.stump_selection,
            "random_state": self.random_state,
        }

        warnings.warn("As for now base_stump is not serializable")

        dictionary = {
            "tree_type": self.__class__.__module__,
            "args": args,
            "classes_": self.classes_.tolist(),
            "n_classes_": self.n_classes_,
            "nodes": [],
        }

        while len(node_list) > 0:
            node = node_list.pop()
            if not node.is_leaf():
                node_list += [node.node_l, node.node_r]

            dictionary["nodes"].append(node.node_to_dict())

        with open(filename, 'w') as f:
            json.dump(dictionary, f, cls=json_NumpyEncoder)

        return dictionary

    @classmethod
    def from_dict(cls, filename):
        with open(filename, 'r') as f:
            dictionary = json.load(f)

        class_c = getattr(importlib.import_module(dictionary['tree_type']), dictionary['tree_type'].split('.')[-1])
        tree = class_c(**dictionary['args'])
        tree.classes_ = dictionary['classes_']
        tree.n_classes_ = dictionary['n_classes_']

        nodes = {node["node_id"]: RuleTreeNode.dict_to_node(node) for node in dictionary['nodes']}

        for node_instance, node_info in zip(nodes.values(), dictionary['nodes']):
            if not node_info["is_leaf"]:
                node_instance.node_l = nodes[node_info["left_node"]]
                node_instance.node_r = nodes[node_info["right_node"]]
                nodes[node_info["left_node"]].parent = node_instance
                nodes[node_info["right_node"]].parent = node_instance

        tree.root = nodes['R']

        return tree

    def export_graphviz(self, columns_names=None, scaler=None, float_precision=3, filename:str|None=None):
        dot = self.root.export_graphviz(columns_names=columns_names, scaler=scaler, float_precision=float_precision)
        if filename is None:
            return Source(dot.to_string()).view()
        else:
            Source(dot.to_string()).render(filename=filename)

