import numpy as np

class Node:
    def __init__(self, feature_index=None, threshold=None, left=None, right=None, value=None):
        # for decision node
        self.feature_index = feature_index
        self.threshold = threshold
        self.left = left
        self.right = right
        # for leaf node
        self.value = value

class DecisionTreeClassifier:
    def __init__(self, min_samples_split=2, max_depth=2, feature_selection_strategy=None):
        self.root = None
        self.min_samples_split = min_samples_split
        self.max_depth = max_depth
        self.feature_selection_strategy = feature_selection_strategy

    def fit(self, X, y):
        self.root = self._build_tree(X, y)

    def _build_tree(self, X, y, current_depth=0):
        num_samples, num_features = X.shape
        # split until stopping conditions are met
        if num_samples >= self.min_samples_split and current_depth <= self.max_depth:
            best_split = self._get_best_split(X, y, num_samples, num_features)
            if best_split["info_gain"] > 0:
                left_subtree = self._build_tree(X[best_split["dataset_left"]], y[best_split["dataset_left"]], current_depth+1)
                right_subtree = self._build_tree(X[best_split["dataset_right"]], y[best_split["dataset_right"]], current_depth+1)
                return Node(best_split["feature_index"], best_split["threshold"], left_subtree, right_subtree)
        # compute leaf node
        leaf_value = self._calculate_leaf_value(y)
        return Node(value=leaf_value)

    def _get_best_split(self, X, y, num_samples, num_features):
        best_split = {}
        max_info_gain = -float("inf")
        
        # Determine number of features to consider
        if self.feature_selection_strategy == 'sqrt':
            num_features_considered = int(np.sqrt(num_features))
        elif self.feature_selection_strategy == 'log2':
            num_features_considered = int(np.log2(num_features))
        else:  # None or invalid strategy falls back to using all features
            num_features_considered = num_features

         # Randomly select features to consider for the split
        features_considered = np.random.choice(range(num_features), size=num_features_considered, replace=False)

        for feature_index in features_considered:
            feature_values = X[:, feature_index]
            possible_thresholds = np.unique(feature_values)
            for threshold in possible_thresholds:
                dataset_left, dataset_right = self._split_dataset(X, y, feature_index, threshold)
                if len(dataset_left) > 0 and len(dataset_right) > 0:
                    y, left_y, right_y = y, y[dataset_left], y[dataset_right]
                    current_info_gain = self._information_gain(y, left_y, right_y)
                    if current_info_gain > max_info_gain:
                        best_split["feature_index"] = feature_index
                        best_split["threshold"] = threshold
                        best_split["dataset_left"], best_split["dataset_right"] = dataset_left, dataset_right
                        best_split["info_gain"] = current_info_gain
                        max_info_gain = current_info_gain
        return best_split

    def _split_dataset(self, X, y, feature_index, threshold):
        left_indices = np.where(X[:, feature_index] <= threshold)[0]
        right_indices = np.where(X[:, feature_index] > threshold)[0]
        return left_indices, right_indices

    def _information_gain(self, parent, left_child, right_child):
        weight_l = len(left_child) / len(parent)
        weight_r = len(right_child) / len(parent)
        gain = self._entropy(parent) - (weight_l*self._entropy(left_child) + weight_r*self._entropy(right_child))
        return gain

    def _entropy(self, y):
        class_labels = np.unique(y)
        entropy = 0
        for cls in class_labels:
            p_cls = len(y[y == cls]) / len(y)
            entropy += -p_cls * np.log2(p_cls)
        return entropy

    def _calculate_leaf_value(self, y):
        y = list(y)
        return max(y, key=y.count)

    def predict(self, X):
        predictions = [self._make_prediction(x, self.root) for x in X]
        return np.array(predictions)

    def _make_prediction(self, x, tree):
        if tree.value is not None: return tree.value
        feature_val = x[tree.feature_index]
        if feature_val <= tree.threshold:
            return self._make_prediction(x, tree.left)
        else:
            return self._make_prediction(x, tree.right)