"""
    Implementation of the CART algorithm to train decision tree classifiers.
"""

import numpy as np

from algorithms.default_algorithm import DefaultClassifier
from tree import tree


class CART(DefaultClassifier):
    def __init__(self, max_depth=None, min_samples_stop=0):
        super().__init__(max_depth, min_samples_stop)

    def _best_split(self, X, y, feature_index_occurrences=None, modified_factor=1, father_feature=None):
        """Find the best split for a node.

        "Best" means that the average impurity of the two children, weighted by their
        population, is the smallest possible. Additionally it must be less than the
        impurity of the current node.

        To find the best split, we loop through all the features, and consider all the
        midpoints between adjacent training samples as possible thresholds. We compute
        the Gini impurity of the split generated by that particular feature/threshold
        pair, and return the pair with smallest impurity.

        Returns:
            best_idx: Index of the feature for best split, or None if no split is found.
            best_thr: Threshold to use for the split, or None if no split is found.
        """
        # Need at least two elements to split a node.
        m = y.size
        if m <= 1:
            return None, None

        # Count of each class in the current node.
        num_parent = [np.sum(y == c) for c in range(self.n_classes_)]

        # Gini of current node.
        impurity_parent = 1.0 - sum((n / m) ** 2 for n in num_parent)  # Cart original
        best_impurity = impurity_parent
        best_idx, best_thr = None, None

        # Loop through all features.
        for idx in range(self.n_features_):
            # Sort data along selected feature.
            thresholds, classes = zip(*sorted(zip(X[:, idx], y)))

            # We could actually split the node according to each feature/threshold pair
            # and count the resulting population for each class in the children, but
            # instead we compute them in an iterative fashion, making this for loop
            # linear rather than quadratic.
            num_left = [0] * self.n_classes_
            num_right = num_parent.copy()
            for i in range(1, m):  # possible split positions
                c = classes[i - 1]
                num_left[c] += 1
                num_right[c] -= 1
                impurity_left = 1.0 - sum(
                    (num_left[x] / i) ** 2 for x in range(self.n_classes_)
                )
                impurity_right = 1.0 - sum(
                    (num_right[x] / (m - i)) ** 2 for x in range(self.n_classes_)
                )

                # impurity of a split is the weighted average of the impurity of the children.
                impurity = (i * impurity_left + (m - i) * impurity_right) / m
                # modified_impurity = impurity * modified_factor if feature_index_occurrences[idx] else impurity

                # The following condition is to make sure we don't try to split two
                # points with identical values for that feature, as it is impossible
                # (both have to end up on the same side of a split).
                if thresholds[i] == thresholds[i - 1]:
                    continue

                if impurity < best_impurity:
                    best_impurity = impurity
                    best_idx = idx
                    best_thr = (thresholds[i] + thresholds[i - 1]) / 2  # midpoint

        return best_idx, best_thr

    def _grow_tree(self, X, y, depth=0, feature_index_occurrences=None, modified_factor=1, calculate_gini=True,
                   father_feature=None, gamma_factor=None):
        """Build a decision tree by recursively finding the best split."""
        # Population for each class in current node. The predicted class is the one with
        # largest population.
        num_samples_per_class = [np.sum(y == i) for i in range(self.n_classes_)]
        predicted_class = np.argmax(num_samples_per_class)
        node = tree.Node(
            num_samples=y.size,
            num_samples_per_class=num_samples_per_class,
            predicted_class=predicted_class,
            feature_index_occurrences=feature_index_occurrences.copy()
        )
        if calculate_gini:
            node.gini = self._gini(y)

        # Split recursively until maximum depth is reached.
        if depth < self.max_depth and node.num_samples >= self.min_samples_stop:
            idx, thr = self._best_split(X, y, feature_index_occurrences=feature_index_occurrences,
                                        modified_factor=modified_factor, father_feature=father_feature)
            if idx is not None:
                indices_left = X[:, idx] < thr
                X_left, y_left = X[indices_left], y[indices_left]
                X_right, y_right = X[~indices_left], y[~indices_left]
                node.feature_index = idx
                node.threshold = thr
                node.feature_index_occurrences[idx] += 1
                node.left = self._grow_tree(X_left, y_left, depth + 1,
                                            feature_index_occurrences=node.feature_index_occurrences.copy(),
                                            modified_factor=modified_factor, calculate_gini=calculate_gini,
                                            father_feature=node.feature_index)
                node.right = self._grow_tree(X_right, y_right, depth + 1,
                                             feature_index_occurrences=node.feature_index_occurrences.copy(),
                                             modified_factor=modified_factor, calculate_gini=calculate_gini,
                                             father_feature=node.feature_index)
        return node
