# Experiments for "Decision Trees and Ensembles": every number in the post comes from here (numpy 2.1, seeded).
# CART regression trees, cost-complexity pruning, bagging, random forests and gradient boosting are implemented here directly.
#   python trees.py [OUT.json]
import sys, json, numpy as np

def best_split(X, y, idx, feats, min_leaf):
    """Best (feature, threshold) by RSS reduction among `feats`; threshold = midpoint between neighbouring values."""
    best = (None, None, np.inf); ys = y[idx]
    for j in feats:
        o = idx[np.argsort(X[idx, j], kind="stable")]; xs, yo = X[o, j], y[o]; n = len(o)
        cs, cs2 = np.cumsum(yo), np.cumsum(yo ** 2); tot, tot2 = cs[-1], cs2[-1]
        k = np.arange(min_leaf, n - min_leaf + 1)
        if len(k) == 0: continue
        k = k[xs[k - 1] < xs[np.minimum(k, n - 1)]] if n > 1 else k
        k = k[k < n]
        if len(k) == 0: continue
        sl, sl2 = cs[k - 1], cs2[k - 1]; rss = (sl2 - sl ** 2 / k) + ((tot2 - sl2) - (tot - sl) ** 2 / (n - k))
        i = int(np.argmin(rss))
        if rss[i] < best[2] - 1e-12: best = (j, (xs[k[i] - 1] + xs[k[i]]) / 2, rss[i])
    return best
def grow(X, y, idx=None, depth=0, max_depth=30, min_leaf=5, mtry=None, rng=None):
    if idx is None: idx = np.arange(len(y))
    node = {"n": len(idx), "pred": float(y[idx].mean()), "rss": float(((y[idx] - y[idx].mean()) ** 2).sum())}
    if depth >= max_depth or len(idx) < 2 * min_leaf or node["rss"] <= 1e-12: return node
    feats = np.arange(X.shape[1]) if mtry is None else rng.choice(X.shape[1], mtry, replace=False)
    j, t, rss = best_split(X, y, idx, feats, min_leaf)
    if j is None: return node
    L, R = idx[X[idx, j] <= t], idx[X[idx, j] > t]
    node.update(j=int(j), t=float(t), left=grow(X, y, L, depth + 1, max_depth, min_leaf, mtry, rng), right=grow(X, y, R, depth + 1, max_depth, min_leaf, mtry, rng))
    return node
def predict(node, X):
    out = np.empty(len(X))
    for i, x in enumerate(X):
        nd = node
        while "j" in nd: nd = nd["left"] if x[nd["j"]] <= nd["t"] else nd["right"]
        out[i] = nd["pred"]
    return out
def leaves(nd): return 1 if "j" not in nd else leaves(nd["left"]) + leaves(nd["right"])
def leaf_rss(nd): return nd["rss"] if "j" not in nd else leaf_rss(nd["left"]) + leaf_rss(nd["right"])
def prune_path(tree):
    """Weakest-link (cost-complexity) pruning: the sequence of (alpha, subtree) of Breiman et al. (1984)."""
    import copy; T = copy.deepcopy(tree); path = [(0.0, copy.deepcopy(T))]
    while "j" in T:
        def walk(nd, acc):
            if "j" in nd: acc.append(((nd["rss"] - leaf_rss(nd)) / (leaves(nd) - 1), nd)); walk(nd["left"], acc); walk(nd["right"], acc)
            return acc
        a = min(g for g, _ in walk(T, []))
        for g, nd in walk(T, []):
            if g <= a + 1e-12: nd.pop("j"); nd.pop("t"); nd.pop("left"); nd.pop("right")
        path.append((a, copy.deepcopy(T)))
    return path
def subtree_at(path, alpha): return [T for a, T in path if a <= alpha + 1e-12][-1]

rng = np.random.default_rng(0); out = {}
# 1. a salary-like data set: experience matters first, then (for experienced players) hits
n = 260; years = rng.integers(1, 24, n).astype(float); hits = rng.integers(1, 240, n).astype(float)
logsal = 5.1 + 1.0 * (years >= 4.5) + 0.7 * (years >= 4.5) * (hits >= 117.5) + 0.25 * rng.normal(0, 1, n)
X1 = np.c_[years, hits]; t2 = grow(X1, logsal, max_depth=2)
def show(nd, pre=""):
    if "j" not in nd: return f"{pre}leaf: n {nd['n']}, mean {nd['pred']:.3f}\n"
    nm = ["years", "hits"][nd["j"]]; return f"{pre}{nm} < {nd['t']:.1f}?\n" + show(nd["left"], pre + "  yes: ") + show(nd["right"], pre + "  no:  ")
print("1. depth-2 tree on the salary-like data (true splits: years 4.5, hits 117.5)\n" + show(t2))
out["salary"] = {"years": years.tolist(), "hits": hits.tolist(), "y": np.round(logsal, 4).tolist()}

# 2. growing, pruning and cross-validating on a smooth 2-D function
f2 = lambda X: np.sin(3 * X[:, 0]) + np.cos(2 * X[:, 1]) + X[:, 0] * X[:, 1]
X = rng.uniform(-1.5, 1.5, (200, 2)); y = f2(X) + rng.normal(0, 0.4, 200); Xt = rng.uniform(-1.5, 1.5, (20000, 2)); yt = f2(Xt) + rng.normal(0, 0.4, 20000)
big = grow(X, y, min_leaf=1); path = prune_path(big)
print(f"2. full tree: {leaves(big)} leaves, training MSE {np.mean((y - predict(big, X)) ** 2):.4f}, test MSE {np.mean((yt - predict(big, Xt)) ** 2):.4f}")
folds = np.array_split(rng.permutation(200), 10); alphas = sorted(set(a for a, _ in path)); cv = np.zeros(len(alphas))
for te in folds:
    tr = np.setdiff1d(np.arange(200), te); p = prune_path(grow(X[tr], y[tr], min_leaf=1))
    for i, a in enumerate(alphas): cv[i] += np.sum((y[te] - predict(subtree_at(p, a), X[te])) ** 2) / 200
rows = []
for i, a in enumerate(alphas):
    T = subtree_at(path, a); rows.append((a, leaves(T), np.mean((y - predict(T, X)) ** 2), cv[i], np.mean((yt - predict(T, Xt)) ** 2)))
R = np.array(rows); ib = int(np.argmin(R[:, 3])); it = int(np.argmin(R[:, 4]))
print(f"   pruning path: {len(R)} subtrees. CV picks alpha {R[ib, 0]:.3f}: {int(R[ib, 1])} leaves, CV MSE {R[ib, 3]:.4f}, test MSE {R[ib, 4]:.4f}; best possible on test: {int(R[it, 1])} leaves, test MSE {R[it, 4]:.4f}")
for k in (1, 2, 4, 8, 16, 32):
    j = int(np.argmin(np.abs(R[:, 1] - k))); print(f"   ~{k:2d} leaves ({int(R[j, 1])}): train {R[j, 2]:.4f}  CV {R[j, 3]:.4f}  test {R[j, 4]:.4f}")
out["prune"] = {"X": np.round(X, 4).tolist(), "y": np.round(y, 4).tolist(), "leaves": R[:, 1].tolist(), "train": R[:, 2].tolist(), "cv": R[:, 3].tolist(), "test": R[:, 4].tolist(), "alpha": R[:, 0].tolist()}

# 3. instability: the same data, resampled
roots = []
for _ in range(200):
    i = rng.integers(0, 200, 200); t = grow(X[i], y[i], max_depth=1); roots.append((t["j"], t["t"]))
r = np.array(roots); print(f"3. root split over 200 bootstrap resamples: on x1 in {np.mean(r[:, 0] == 0) * 100:.0f}%, on x2 in {np.mean(r[:, 0] == 1) * 100:.0f}%; "
      f"threshold range x1 [{r[r[:, 0] == 0, 1].min():.2f}, {r[r[:, 0] == 0, 1].max():.2f}]" + (f", x2 [{r[r[:, 0] == 1, 1].min():.2f}, {r[r[:, 0] == 1, 1].max():.2f}]" if (r[:, 0] == 1).any() else ""))

# 4. trees vs linear models
def lin_fit(X, y): Z = np.c_[np.ones(len(X)), X]; b = np.linalg.lstsq(Z, y, rcond=None)[0]; return lambda W: np.c_[np.ones(len(W)), W] @ b
res = {}
for nm, g in (("linear truth", lambda Z: 1 + 2 * Z[:, 0] - 1.5 * Z[:, 1]), ("box truth", lambda Z: 1.0 * ((Z[:, 0] > 0.3) & (Z[:, 1] < -0.2)) * 2.5)):
    e_t, e_l = [], []
    for _ in range(50):
        Xa = rng.uniform(-1.5, 1.5, (200, 2)); ya = g(Xa) + rng.normal(0, 0.4, 200); Xb = rng.uniform(-1.5, 1.5, (5000, 2)); yb = g(Xb) + rng.normal(0, 0.4, 5000)
        pa = prune_path(grow(Xa, ya, min_leaf=1)); fo = np.array_split(rng.permutation(200), 5); al = sorted(set(a for a, _ in pa)); c = np.zeros(len(al))
        for te in fo:
            tr = np.setdiff1d(np.arange(200), te); pp = prune_path(grow(Xa[tr], ya[tr], min_leaf=1))
            for k, a in enumerate(al): c[k] += np.sum((ya[te] - predict(subtree_at(pp, a), Xa[te])) ** 2)
        T = subtree_at(pa, al[int(np.argmin(c))]); e_t.append(np.mean((yb - predict(T, Xb)) ** 2)); e_l.append(np.mean((yb - lin_fit(Xa, ya)(Xb)) ** 2))
    res[nm] = (np.mean(e_t), np.mean(e_l)); print(f"4. {nm}: pruned tree test MSE {np.mean(e_t):.3f}, linear regression {np.mean(e_l):.3f} (noise 0.160)")

# 5. impurity measures for one split (Hastie et al. 2009, sec. 9.2.3): parent 400/400, split A (300,100)|(100,300), split B (200,400)|(200,0)
def imp(c):
    p = np.array(c) / sum(c); return 1 - p.max(), float(np.sum(p * (1 - p))), float(-np.sum([q * np.log2(q) for q in p if q > 0]))
for nm, (l, r) in (("A", ((300, 100), (100, 300))), ("B", ((200, 400), (200, 0)))):
    nl, nr = sum(l), sum(r); w = [(nl * a + nr * b) / 800 for a, b in zip(imp(l), imp(r))]
    print(f"5. split {nm}: weighted misclassification {w[0]:.4f}, Gini {w[1]:.4f}, entropy {w[2]:.4f} (parent: 0.5, 0.5, 1.0)")

# 6. ensembles on Friedman's first benchmark: 10 inputs, 5 of them used
fr = lambda Z: 10 * np.sin(np.pi * Z[:, 0] * Z[:, 1]) + 20 * (Z[:, 2] - 0.5) ** 2 + 10 * Z[:, 3] + 5 * Z[:, 4]
Xf = rng.uniform(0, 1, (300, 10)); yf = fr(Xf) + rng.normal(0, 1, 300); Xg = rng.uniform(0, 1, (5000, 10)); yg = fr(Xg) + rng.normal(0, 1, 5000)
mse = lambda p: float(np.mean((yg - p) ** 2))
print(f"6. Friedman #1, n = 300, 10 inputs (noise variance 1): linear regression {mse(lin_fit(Xf, yf)(Xg)):.3f}")
pf = prune_path(grow(Xf, yf, min_leaf=1)); fo = np.array_split(rng.permutation(300), 5); al = sorted(set(a for a, _ in pf)); c = np.zeros(len(al))
for te in fo:
    tr = np.setdiff1d(np.arange(300), te); pp = prune_path(grow(Xf[tr], yf[tr], min_leaf=1))
    for k, a in enumerate(al): c[k] += np.sum((yf[te] - predict(subtree_at(pp, a), Xf[te])) ** 2)
Tf = subtree_at(pf, al[int(np.argmin(c))]); print(f"   single tree, pruned by 5-fold CV: {leaves(Tf)} leaves, test MSE {mse(predict(Tf, Xg)):.3f}")
B = 300; bag_p = np.zeros(5000); rf_p = np.zeros(5000); oob_sum = np.zeros(300); oob_n = np.zeros(300); curves = {"bag": [], "rf": [], "oob": [], "boost": {}}
for b in range(1, B + 1):
    i = rng.integers(0, 300, 300); oob = np.setdiff1d(np.arange(300), i)
    tb = grow(Xf[i], yf[i], min_leaf=1); bag_p += predict(tb, Xg); oob_sum[oob] += predict(tb, Xf[oob]); oob_n[oob] += 1
    tr_ = grow(Xf[i], yf[i], min_leaf=1, mtry=3, rng=rng); rf_p += predict(tr_, Xg)
    if b in (1, 2, 5, 10, 20, 50, 100, 200, 300):
        m = oob_n > 0; curves["bag"].append((b, mse(bag_p / b))); curves["rf"].append((b, mse(rf_p / b))); curves["oob"].append((b, float(np.mean((yf[m] - oob_sum[m] / oob_n[m]) ** 2))))
for (b, eb), (_, er), (_, eo) in zip(curves["bag"], curves["rf"], curves["oob"]): print(f"   B = {b:3d}: bagging {eb:.3f} (out-of-bag estimate {eo:.3f}), random forest (m = 3) {er:.3f}")
for lam, depth in ((0.1, 2), (0.01, 2), (1.0, 2), (0.1, 1)):
    F = np.full(300, yf.mean()); G = np.full(5000, yf.mean()); cur = []
    for b in range(1, 2001):
        t = grow(Xf, yf - F, max_depth=depth, min_leaf=1); F += lam * predict(t, Xf); G += lam * predict(t, Xg)
        if b in (1, 10, 50, 100, 200, 500, 1000, 2000): cur.append((b, mse(G)))
    curves["boost"][f"{lam}_{depth}"] = cur; print(f"   boosting, learning rate {lam}, depth {depth}: " + ", ".join(f"B={b} {e:.3f}" for b, e in cur))
out["ens"] = {"tree": mse(predict(Tf, Xg)), "linear": mse(lin_fit(Xf, yf)(Xg)), **curves}
if len(sys.argv) > 1: json.dump(out, open(sys.argv[1], "w"), separators=(",", ":"), default=float)