import numpy as np
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from collections import Counter


class DecisionTree:
    def __init__(self, max_depth=None, min_samples_split=2, n_features=None):
        self.max_depth = max_depth
        self.min_samples_split = min_samples_split
        self.n_features = n_features  # 随机森林中使用，限制每次分裂时考虑的特征数量
        self.root = None

    def fit(self, X, y):
        # 处理n_features参数，如果是字符串则转换为整数
        if isinstance(self.n_features, str):
            if self.n_features == "sqrt":
                self.n_features = int(np.sqrt(X.shape[1]))
            elif self.n_features == "log2":
                self.n_features = int(np.log2(X.shape[1]))
            else:
                self.n_features = X.shape[1]
        # 如果没有设置n_features，则使用所有特征
        self.n_features = X.shape[1] if self.n_features is None else min(self.n_features, X.shape[1])
        self.root = self._grow_tree(X, y)

    def _grow_tree(self, X, y, depth=0):
        n_samples, n_features = X.shape
        n_labels = len(np.unique(y))

        # 停止条件
        if (depth >= self.max_depth
                or n_labels == 1
                or n_samples < self.min_samples_split):
            leaf_value = self._most_common_label(y)
            return LeafNode(value=leaf_value)

        # 随机选择特征子集
        feat_idxs = np.random.choice(n_features, self.n_features, replace=False)

        # 找到最佳分裂点
        best_gain, best_feat, best_thresh = self._best_split(X, y, feat_idxs)

        # 如果无法找到有效的分裂点，创建叶节点
        if best_gain == -1:
            leaf_value = self._most_common_label(y)
            return LeafNode(value=leaf_value)

        # 创建决策节点
        left_idxs, right_idxs = self._split(X[:, best_feat], best_thresh)

        # 检查分裂是否有效
        if len(left_idxs) == 0 or len(right_idxs) == 0:
            leaf_value = self._most_common_label(y)
            return LeafNode(value=leaf_value)

        left = self._grow_tree(X[left_idxs, :], y[left_idxs], depth + 1)
        right = self._grow_tree(X[right_idxs, :], y[right_idxs], depth + 1)
        return DecisionNode(feature_idx=best_feat, threshold=best_thresh, left=left, right=right)

    def _best_split(self, X, y, feat_idxs):
        best_gain = -1
        split_idx, split_thresh = None, None

        for feat_idx in feat_idxs:
            X_column = X[:, feat_idx]
            thresholds = np.unique(X_column)

            for threshold in thresholds:
                gain = self._information_gain(y, X_column, threshold)

                if gain > best_gain:
                    best_gain = gain
                    split_idx = feat_idx
                    split_thresh = threshold

        return best_gain, split_idx, split_thresh

    def _information_gain(self, y, X_column, split_thresh):
        # 计算父节点的熵
        parent_entropy = self._entropy(y)

        # 生成划分
        left_idxs, right_idxs = self._split(X_column, split_thresh)

        if len(left_idxs) == 0 or len(right_idxs) == 0:
            return 0

        # 计算加权平均子节点熵
        n = len(y)
        n_l, n_r = len(left_idxs), len(right_idxs)
        e_l, e_r = self._entropy(y[left_idxs]), self._entropy(y[right_idxs])
        child_entropy = (n_l / n) * e_l + (n_r / n) * e_r

        # 计算信息增益
        ig = parent_entropy - child_entropy
        return ig

    def _split(self, X_column, split_thresh):
        left_idxs = np.argwhere(X_column <= split_thresh).flatten()
        right_idxs = np.argwhere(X_column > split_thresh).flatten()
        return left_idxs, right_idxs

    def _entropy(self, y):
        hist = np.bincount(y)
        ps = hist / len(y)
        return -np.sum([p * np.log2(p) for p in ps if p > 0])

    def _most_common_label(self, y):
        if len(y) == 0:
            return 0  # 默认返回0作为标签，或者可以根据需要修改
        counter = Counter(y)
        most_common = counter.most_common(1)[0][0]
        return most_common

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

    def _traverse_tree(self, x, node):
        if isinstance(node, LeafNode):
            return node.value

        if x[node.feature_idx] <= node.threshold:
            return self._traverse_tree(x, node.left)
        return self._traverse_tree(x, node.right)


class DecisionNode:
    def __init__(self, feature_idx, threshold, left, right):
        self.feature_idx = feature_idx
        self.threshold = threshold
        self.left = left
        self.right = right


class LeafNode:
    def __init__(self, value):
        self.value = value


class RandomForest:
    def __init__(self, n_trees=100, max_depth=None, min_samples_split=2,
                 max_features=None):
        self.n_trees = n_trees
        self.max_depth = max_depth
        self.min_samples_split = min_samples_split
        self.max_features = max_features
        self.trees = []

    def fit(self, X, y):
        self.trees = []
        for _ in range(self.n_trees):
            tree = DecisionTree(
                max_depth=self.max_depth,
                min_samples_split=self.min_samples_split,
                n_features=self.max_features
            )

            # 自助采样（bootstrap）
            X_sample, y_sample = self._bootstrap_samples(X, y)

            # 训练决策树
            tree.fit(X_sample, y_sample)
            self.trees.append(tree)

    def _bootstrap_samples(self, X, y):
        n_samples = X.shape[0]
        idxs = np.random.choice(n_samples, n_samples, replace=True)
        return X[idxs], y[idxs]

    def predict(self, X):
        tree_preds = np.array([tree.predict(X) for tree in self.trees])
        tree_preds = np.swapaxes(tree_preds, 0, 1)

        # 多数投票
        y_pred = [np.bincount(pred).argmax() for pred in tree_preds]
        return np.array(y_pred)


# 加载数据
iris = load_iris()
X, y = iris.data, iris.target
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 训练随机森林
rf = RandomForest(n_trees=100, max_depth=10, max_features="sqrt")
rf.fit(X_train, y_train)

# 预测并评估
y_pred = rf.predict(X_test)
accuracy = np.sum(y_pred == y_test) / len(y_test)
print(f"随机森林准确率: {accuracy:.4f}")

# 对比单棵决策树
dt = DecisionTree(max_depth=10)
dt.fit(X_train, y_train)
y_pred_dt = dt.predict(X_test)
accuracy_dt = np.sum(y_pred_dt == y_test) / len(y_test)
print(f"单棵决策树准确率: {accuracy_dt:.4f}")