# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/01_APIS/01_basetree.ipynb.

# %% auto 0
__all__ = ['DecisionNode', 'BaseDecisionTree']

# %% ../nbs/01_APIS/01_basetree.ipynb 3
from fastcore.utils import *
import numpy as np
import pandas as pd
from .utils import *

# %% ../nbs/01_APIS/01_basetree.ipynb 5
class DecisionNode():
    """Class that represents a decision node or leaf in the decision tree"""
    def __init__(self, 
                feature_i: Union[str,int]=None,  # Feature index which we want to use as the threshold measure.
                threshold: float =None,          # The value that we will compare feature values at feature_i against to determine the prediction.
                value=None,                      # The class prediction if classification tree, or float value if regression tree.
                true_branch=None,                # Next decision node for samples where features value met the threshold.
                false_branch=None,               # Next decision node for samples where features value did not meet the threshold.
                depth: int=None                  # depth of current node
                ):
        self.feature_i = feature_i    
        self.threshold= threshold             
        self.value = value                            #
        self.true_branch = true_branch  #
        self.false_branch = false_branch# 
        self.estimator =  None                        # An estimator that can reun at leaf node 
        self.is_leaf = False                          # indicator weather currnet leaf is final leaf or not
        self.loss = None                              # loss of gradient bossiting
        self.grad = None                              # gradient
        self.w = None                                 # weight
        self.resid = None                             # Residuals
        self.depth = depth                            
        self.addons = None                            # added information needed for diagnosis
        
    def estimate(self,x):
        return self.value if self.estimator is None else self.estimator.predict(x)

# Super class of RegressionTree and ClassificationTree
class BaseDecisionTree(object):
    """Super class of RegressionTree and ClassificationTree."""
    def __init__(self, 
                min_samples_split:int = 20,   # The minimum number of samples needed to make a split when building a tree.
                min_impurity:float = 1e-7,    # The minimum impurity required to split the tree further.
                min_sample_leaf: int = 10,    # Minimum sample required to have a leaf node
                max_depth:int = float("inf"), # The maximum depth of a tree.
                loss: Union[str] = None,      # String of loss or funciton which defines the loss. This amounts to Loss function that is used for Gradient Boosting models to calculate impurity.
                verbose:Union[bool,int]=False, # Verbosity for tree builidng
                digits:Union[int, None] = 2    # To round the values before doing a split 
                ):
        self.root = None  # Root node in dec. tree
        self.min_samples_split = min_samples_split
        self.min_impurity = min_impurity
        self.max_depth = max_depth
        self.min_sample_leaf = min_sample_leaf

        # function definig hwo tree woudl grow
        self._node_function = None
        self._leaf_function = None
        self._split_function = None

        # If y is one-hot encoded (multi-dim) or not (one-dim)
        self.one_dim = None

        # If Gradient Boost
        self.loss = loss

        self.cat_features = None
        self.con_features = None
        self.features_ = None
        self.target_ = None

        # other things
        self.digits, self.verbose  = digits, verbose


    def fit(self, X, y, loss=None):
        """ Build decision tree """
        self.cat_features = X.select_dtypes(include=['object','category']).columns.tolist()
        self.con_features = [col for col in X.columns if col not in self.cat_features]
        self.features_ = X.columns.tolist()
        self.target_ = y.name
        self.one_dim = len(np.shape(y)) == 1
        # # Check if expansion of y is needed
        # if len(np.shape(y)) == 1:
        #     y = np.expand_dims(y, axis=1)

        # Add y as last column of X
        Xy = pd.concat((X, y), axis=1)
        self.root = self._build_tree(Xy)
        self.loss=None

    def _build_tree(self, Xy, current_depth=0,features=None):
        """ Recursive method which builds out the decision tree and splits X and respective y
        on the feature of X which (based on impurity) best separates the data"""

        largest_impurity = 0
        bst_estimator = None    
        bst_weights = None        
        best_criteria = {"feature_i": None, "threshold": None} # Feature index and threshold
        best_sets = {
            "leftXy": None,   # Xy of left subtree
            "rightXy": None # Xy of right subtree
            }# Subsets of the data
        
        n_samples, n_features = np.shape(Xy)
        if len(self.features_) + 1  == n_features:
            if self.verbose : print(f"\nWe are going good at {current_depth} depth")
        else:
            print('Somethign is wrong',n_samples, n_features, self.features_)

        if n_samples >= self.min_samples_split and current_depth <= self.max_depth:
            # Calculate the impurity for each feature
            for feature_i in self.features_:
                if self.verbose : print(f'Evaluating {feature_i}')
                # All values of feature_i       
                if feature_i in self.cat_features:
                    unique_values = get_sorted_cats(Xy[feature_i].values)
                    thresh = []
                else:
                    if self.digits is not None:
                        unique_values = np.unique(np.round(Xy[feature_i],self.digits))
                    else:
                        unique_values = np.unique(Xy[feature_i])
                    
                # Iterate through all unique values of feature column i and
                # calculate the impurity
                for threshold in unique_values:
                    # Divide X and y depending on if the feature value of X at index feature_i
                    # meets the threshold
                    if feature_i in self.cat_features:
                        thresh.append(threshold)
                        Xy1, Xy2 = divide_on_feature(Xy, feature_i, thresh)  
                    else:
                        Xy1, Xy2 = divide_on_feature(Xy, feature_i, threshold)

                    if len(Xy1) > 0 and len(Xy2) > 0:
                        
                        # Select the y-values of the two sets
                        # y1 = Xy1[:, n_features:]
                        # y2 = Xy2[:, n_features:]
                
                        # Calculate impurity
                        # impurity = self._impurity_calculation(y, y1, y2)
                        
                        # evaluate node function
                        impurity, estimator, weights, addons = self._node_function(Xy, Xy1,Xy2,target,features=features)

                        # If this threshold resulted in a higher information gain than previously
                        # recorded save the threshold value and the feature
                        # index
                        if impurity > largest_impurity:
                            largest_impurity, bst_estimator, bst_weights, bst_addons = impurity, estimator, weights, addons
                            best_criteria["feature_i"]= feature_i
                            best_criteria["threshold"] =  thresh if feature_i in self.cat_features else threshold
                            best_sets["leftXy"]=  Xy1 # X of left subtree
                            best_sets["rightXy"]=  Xy2 # X of left subtree
                            
                if self.verbose: print(f"Impurity Update at {best_criteria} with {largest_impurity}")
                            
                                
        if largest_impurity > self.min_impurity:
            if self.verbose: print(f'Best Impurtiy is with {best_criteria} with {largest_impurity} Split sizes: Left: {best_sets["leftXy"].shape} : Right {best_sets["rightXy"].shape}')
            # Build subtrees for the right and left branches
            true_branch = self._build_tree(best_sets["leftXy"], current_depth + 1)
            false_branch = self._build_tree(best_sets["rightXy"], current_depth + 1)
            node = DecisionNode(feature_i=best_criteria["feature_i"], threshold=best_criteria[
                                "threshold"], true_branch=true_branch, false_branch=false_branch,depth=current_depth)
            node.estimator = bst_estimator
            node.w = bst_weights
            node.addons = bst_addons
            return node
        else:
            if self.verbose: print(f"No Impurtiy update. Hence a leaf with {best_criteria} with {largest_impurity}")

        # We're at leaf => determine value
        leaf_value,estimator,weights,addons = self._leaf_function(Xy, target, features=features)
        leaf_node = DecisionNode(value=leaf_value)
        leaf_node.estimator = estimator
        leaf_node.w = weights
        leaf_node.addons = addons
        return leaf_node


    def predict_value(self, x, tree=None,verbose=False):
        """ Do a recursive search down the tree and make a prediction of the data sample by the
            value of the leaf that we end up at """

        if tree is None:
            tree = self.root

        # If we have a value (i.e we're at a leaf) => return value as the prediction
        if tree.value is not None:
            return tree.estimate(x)

        # Choose the feature that we will test
        if isinstance(x, pd.Series):
            feature_value = x.loc[tree.feature_i]
        else:
            feature_value = x[tree.feature_i].values[0]
        if verbose:
            print(f'Feature value for {tree.feature_i} :{feature_value}')

        # Determine if we will follow left or right branch
        branch = tree.false_branch
        if isinstance(feature_value, int) or isinstance(feature_value, float):
            if feature_value >= tree.threshold:  branch = tree.true_branch
        elif isinstance(tree.threshold,list):
            if feature_value in tree.threshold: branch = tree.true_branch
        elif feature_value == tree.threshold: 
            branch = tree.true_branch

        # Test subtree
        return self.predict_value(x, branch,verbose)

    def predict(self, X):
        """ Classify samples one by one and return the set of labels """
        y_pred = [self.predict_value(sample) for sample in X]
        return y_pred

    def print_tree(self, tree=None, indent=" "):
        """ Recursively print the decision tree """
        if not tree:
            tree = self.root

        # If we're at leaf => print the label
        if tree.value is not None:
            print (tree.value)
        # Go deeper down the tree
        else:
            # Print test
            print ("%s:%s? " % (tree.feature_i, tree.threshold))
            # Print the true scenario
            print ("%sT->" % (indent), end="")
            self.print_tree(tree.true_branch, indent + indent)
            # Print the false scenario
            print ("%sF->" % (indent), end="")
            self.print_tree(tree.false_branch, indent + indent)
            
