import numpy as np
from dataclasses import dataclass
from typing import Optional


@dataclass
class TreeNode:
    feature: Optional[int] = None
    threshold: Optional[float] = None
    left: Optional["TreeNode"] = None
    right: Optional["TreeNode"] = None
    prediction: Optional[float] = None


def gini(y):
    if len(y) == 0:
        return 0.0
    _, counts = np.unique(y, return_counts=True)
    p = counts / len(y)
    return 1.0 - np.sum(p ** 2)


def best_split(X, y):
    n, d = X.shape
    best_gain, best_feat, best_thresh = -1.0, None, None
    base_impurity = gini(y)
    for feat in range(d):
        thresholds = np.unique(X[:, feat])
        for t in thresholds:
            left_mask = X[:, feat] <= t
            right_mask = ~left_mask
            if left_mask.sum() == 0 or right_mask.sum() == 0:
                continue
            left_imp = gini(y[left_mask])
            right_imp = gini(y[right_mask])
            weighted = (
                left_mask.sum() * left_imp
                + right_mask.sum() * right_imp
            ) / n
            gain = base_impurity - weighted
            if gain > best_gain:
                best_gain, best_feat, best_thresh = gain, feat, t
    return best_feat, best_thresh, best_gain


def build_tree(X, y, depth=0, max_depth=5, min_samples=2):
    if (
        depth >= max_depth
        or len(y) < min_samples
        or gini(y) == 0.0
    ):
        values, counts = np.unique(y, return_counts=True)
        return TreeNode(prediction=values[np.argmax(counts)])

    feat, thresh, gain = best_split(X, y)
    if feat is None or gain <= 0:
        values, counts = np.unique(y, return_counts=True)
        return TreeNode(prediction=values[np.argmax(counts)])

    left_mask = X[:, feat] <= thresh
    left = build_tree(
        X[left_mask], y[left_mask], depth + 1, max_depth, min_samples
    )
    right = build_tree(
        X[~left_mask], y[~left_mask], depth + 1, max_depth, min_samples
    )
    return TreeNode(feature=feat, threshold=thresh, left=left, right=right)


def predict_one(node, x):
    while node.prediction is None:
        node = node.left if x[node.feature] <= node.threshold else node.right
    return node.prediction


def predict_tree(node, X):
    return np.array([predict_one(node, x) for x in X])
