import numpy as np
import pandas as pd


class DecisionTree:
    """decision tree with CART algorithm"""

    def __init__(self):
        self.tree = {}

    def __repr__(self):
        return 'Decision Tree( %r )' % self.tree

    def train(self):
        """create tree and apply post-pruning"""
        '''t_data for training data, v_data for validation data, t_label, v_label in the same way'''
        t_data, t_label, v_data, v_label = self.__load_data()
        self.tree = self.__create_tree(t_data, t_label, v_data, v_label)
        self.__post_prune(v_data, v_label)

    def test(self):
        df = pd.read_csv('../data/sonar.csv', header=None)
        d, l = np.array(df.iloc[:, :-1]), (np.array(df.iloc[:, -1]) == 'R').astype(int)
        '''test the code on the testing set'''
        dataset, labels = np.append(d[79:97], d[189:], axis=0), np.append(l[79:97], l[189:], axis=0)
        '''test the code on the training set'''
        # dataset, labels = np.append(d[0:55], d[97:161], axis=0), np.append(l[0:55], l[97:161], axis=0)
        '''test the code on the validation set'''
        # dataset, labels = np.append(d[55:79], d[161:189], axis=0), np.append(l[55:79], l[161:189], axis=0)
        return self.__get_correct_rate(self.tree, dataset, labels)

    def __post_prune(self, v_data, v_labels):
        """recursion to post-prune"""
        self.tree = self.__prune(self.tree, v_data, v_labels)

    @staticmethod
    def __get_correct_rate(t, dataset, labels):
        """calculate the correct rate for the dataset"""
        predictions = []
        for data in dataset:
            predictions.append(DecisionTree.__classify(t, data) > 0.5)
        return sum(np.array(predictions) == labels) / len(labels)

    @staticmethod
    def get_depth(t):
        """get the depth of the tree"""
        left_depth, right_depth = 0, 0
        if isinstance(t, dict):
            left_depth = DecisionTree.get_depth(t['left'])
            right_depth = DecisionTree.get_depth(t['right'])
        return max(left_depth, right_depth) + 1

    @staticmethod
    def get_leaves_count(t):
        """count the number of leaves"""
        if isinstance(t, float):
            return 1
        if isinstance(t, dict):
            left_count = DecisionTree.get_leaves_count(t['left'])
            right_count = DecisionTree.get_leaves_count(t['right'])
            return left_count + right_count

    @staticmethod
    def __prune(t, dataset, labels):
        """on the node, prune the branch from the bottom to the top"""
        '''if is float, return itself'''
        if isinstance(t, float):
            return t
        pos = np.array([t['feature'], t['value']])
        l_data, l_labels, r_data, r_labels = DecisionTree.__divide_two_node(dataset, labels, pos)
        '''apply pruning to left and right'''
        if isinstance(t['left'], dict):
            t['left'] = DecisionTree.__prune(t['left'], l_data, l_labels)
        if isinstance(t['right'], dict):
            t['right'] = DecisionTree.__prune(t['right'], r_data, r_labels)
        '''if neither left nor right is subtree, try pruning on the current node'''
        if isinstance(t['left'], float) and isinstance(t['right'], float):
            error_no_merge = sum(np.square(l_labels - t['left'])) + sum(np.square(r_labels - t['right']))
            mean_label = (t['left'] + t['right']) / 2
            error_merge = sum(np.square(l_labels - mean_label)) + sum(np.square(r_labels - mean_label))
            if error_merge < error_no_merge:
                return mean_label
            else:
                return t
        return t

    @staticmethod
    def __classify(t, data):
        """recurse to classify the data"""
        if isinstance(t, float):
            return t
        if data[t['feature']] <= t['value']:
            return DecisionTree.__classify(t['left'], data)
        else:
            return DecisionTree.__classify(t['right'], data)

    @staticmethod
    def __load_data():
        """load data and divide into training set and validation set"""
        df = pd.read_csv('../data/sonar.csv', header=None)
        d, l = np.array(df.iloc[:, :-1]), (np.array(df.iloc[:, -1]) == 'R').astype(int)
        t_data, t_label = np.append(d[0:55], d[97:161], axis=0), np.append(l[0:55], l[97:161], axis=0)
        v_data, v_label = np.append(d[55:79], d[161:189], axis=0), np.append(l[55:79], l[161:189], axis=0)
        return t_data, t_label, v_data, v_label

    @staticmethod
    def __get_err(data, labels, feature=None, value=None):
        """given the feature and the value to divide, return the error after division.
        if not given the feature and value, return total error"""
        if feature is None or value is None:
            return np.var(labels) * len(labels)
        left_labels, right_labels = labels[data[:, feature] <= value], labels[data[:, feature] > value]
        if len(left_labels) == 0 or len(right_labels) == 0:
            return np.var(labels) * len(labels)
        return np.var(left_labels) * len(left_labels) + np.var(right_labels) * len(right_labels)

    @staticmethod
    def __get_min_err(data, labels, feature):
        """given the feature, calculate error with every value feature can get,
        then find the min error then return the corresponding value and the min error"""
        errs = []
        for i in np.arange(len(labels)):
            errs.append(DecisionTree.__get_err(data, labels, feature, data[i, feature]))
        errs = np.array(errs)
        return [data[np.argmin(errs), feature], np.min(errs)]

    @staticmethod
    def __get_best_division(t_data, t_label, v_data, v_label):
        """find the best feature and value to divide the dataset, return the feature and value"""
        if len(np.unique(t_label)) == 1:
            return None
        errs_mat = []
        for i in np.arange(t_data.shape[1]):
            errs_mat.append(DecisionTree.__get_min_err(t_data, t_label, i))
        errs_mat = np.array(errs_mat)
        min_index = errs_mat[:, 1].argmin()
        '''run the statement below to create decision tree without pre-pruning'''
        # return np.array([min_index, errs_mat[min_index][0]])
        '''run the statements below to create decision tree with pre-pruning, 
        if after division the performance is far lower than the former, do not divide'''
        val = errs_mat[min_index][0]
        if DecisionTree.__get_err(v_data, v_label, min_index, val) < DecisionTree.__get_err(v_data, v_label):
            return np.array([min_index, errs_mat[min_index][0]])
        else:
            return None

    @staticmethod
    def __divide_two_node(data, labels, pos):
        """divide the data at the position (feature, value)"""
        feature = int(pos[0])
        value = pos[1]
        right_data, right_label = data[data[:, feature] > value], labels[data[:, feature] > value]
        left_data, left_label = data[data[:, feature] <= value], labels[data[:, feature] <= value]
        return left_data, left_label, right_data, right_label

    @staticmethod
    def __create_tree(t_data, t_label, v_data, v_label):
        """recurse to create the tree"""
        division = DecisionTree.__get_best_division(t_data, t_label, v_data, v_label)
        if division is None:
            return t_label.mean()
        node = {'feature': int(division[0]), 'value': division[1]}
        l_data, l_label, r_data, r_label = DecisionTree.__divide_two_node(t_data, t_label, division)
        node['left'] = DecisionTree.__create_tree(l_data, l_label, v_data, v_label)
        node['right'] = DecisionTree.__create_tree(r_data, r_label, v_data, v_label)
        return node


tree = DecisionTree()
tree.train()
print(tree.test())
print(tree.get_leaves_count(tree.tree))
