"""
rf.py -- Random Forests from scratch (Breiman 2001).

Backs the notebooks in  Z-ML Trees-Random Forests.  Self-contained (a lean deep-tree grower with per-node
feature subsampling, so it does not depend on the CART notebook's module).

A single deep tree is a low-bias, HIGH-variance learner (see the CART notebook). A random forest averages
many deep trees to cancel that variance -- but averaging only helps if the trees are DE-CORRELATED, so
Breiman adds two independent sources of randomness:

    1. BAGGING  -- each tree is grown on a bootstrap resample of the rows.
    2. RANDOM SUBSPACE -- at every split only a random subset of  max_features  columns is considered.

For an average of B trees each with variance sigma^2 and pairwise correlation rho, the ensemble variance is
    rho sigma^2 + (1-rho)/B sigma^2,
so raising B kills the second term while LOWERING rho (via max_features) attacks the first -- the whole
point of the random subspace trick. Because each bootstrap leaves out ~37% of rows, every tree has an
"out-of-bag" (OOB) test set for free: averaging each row's predictions over only the trees that did NOT
see it gives an unbiased generalisation estimate with no separate validation split.
"""

import numpy as np


class _Node:
    __slots__ = ("feat", "thr", "left", "right", "value", "n", "imp")
    def __init__(self, value, n, imp):
        self.feat = None; self.thr = None; self.left = None; self.right = None
        self.value = value; self.n = n; self.imp = imp


class _Tree:
    """deep CART tree that considers a random subset of max_features features at each split."""
    def __init__(self, mf, min_samples_leaf, max_depth, rng, regression, n_classes):
        self.mf = mf; self.msl = min_samples_leaf; self.max_depth = max_depth
        self.rng = rng; self.regression = regression; self.K = n_classes
        self.importances = None

    def _imp(self, y):
        if self.regression: return np.mean((y - y.mean()) ** 2) if len(y) else 0.0
        c = np.bincount(y, minlength=self.K).astype(float); p = c / c.sum(); return 1.0 - (p * p).sum()

    def _leaf(self, y):
        if self.regression: return float(y.mean())
        return np.bincount(y, minlength=self.K).astype(float) / len(y)

    def _best_split(self, X, y):
        n, p = X.shape; parent = self._imp(y); best = (0.0, None, None)
        feats = self.rng.choice(p, min(self.mf, p), replace=False)
        for j in feats:
            xs = X[:, j]; order = np.argsort(xs, kind="quicksort"); xj = xs[order]; yj = y[order]
            cut = np.where(np.diff(xj) > 0)[0]
            if len(cut) == 0: continue
            if self.regression:
                cs = np.cumsum(yj); cs2 = np.cumsum(yj * yj); tot = cs[-1]; tot2 = cs2[-1]
                nL = cut + 1.0; nR = n - nL
                child = ((cs2[cut] - cs[cut] ** 2 / nL) + ((tot2 - cs2[cut]) - (tot - cs[cut]) ** 2 / nR)) / n
            else:
                oh = np.zeros((n, self.K)); oh[np.arange(n), yj] = 1.0; cc = np.cumsum(oh, axis=0)
                nL = (cut + 1.0)[:, None]; nR = n - nL; cL = cc[cut]; cR = cc[-1] - cL
                giniL = 1.0 - ((cL / nL) ** 2).sum(1); giniR = 1.0 - ((cR / nR) ** 2).sum(1)
                child = (nL[:, 0] * giniL + nR[:, 0] * giniR) / n
            gain = parent - child
            ok = (np.asarray(nL).reshape(-1) >= self.msl) & (np.asarray(nR).reshape(-1) >= self.msl)
            if not ok.any(): continue
            gain = np.where(ok, gain, -np.inf); k = int(np.argmax(gain))
            if gain[k] > best[0]:
                best = (float(gain[k]), int(j), float((xj[cut[k]] + xj[cut[k] + 1]) / 2.0))
        return best

    def _build(self, X, y, depth):
        node = _Node(self._leaf(y), len(y), self._imp(y))
        if len(y) < 2 * self.msl or node.imp == 0.0 or (self.max_depth is not None and depth >= self.max_depth):
            return node
        gain, feat, thr = self._best_split(X, y)
        if feat is None or gain <= 0: return node
        m = X[:, feat] <= thr; node.feat = feat; node.thr = thr
        self.importances[feat] += node.imp * node.n - self._imp(y[m]) * m.sum() - self._imp(y[~m]) * (~m).sum()
        node.left = self._build(X[m], y[m], depth + 1); node.right = self._build(X[~m], y[~m], depth + 1)
        return node

    def fit(self, X, y):
        self.importances = np.zeros(X.shape[1]); self.root = self._build(X, y, 0); return self

    def predict_proba(self, X):
        out = []
        for x in X:
            nd = self.root
            while nd.feat is not None: nd = nd.left if x[nd.feat] <= nd.thr else nd.right
            out.append(nd.value)
        return np.array(out)


class RandomForest:
    def __init__(self, n_estimators=100, max_features="sqrt", min_samples_leaf=1,
                 max_depth=None, bootstrap=True, regression=False, random_state=0):
        self.n_estimators = n_estimators; self.max_features = max_features
        self.min_samples_leaf = min_samples_leaf; self.max_depth = max_depth
        self.bootstrap = bootstrap; self.regression = regression; self.random_state = random_state

    def _mf(self, p):
        if self.max_features == "sqrt": return max(1, int(np.sqrt(p)))
        if self.max_features == "log2": return max(1, int(np.log2(p)))
        if isinstance(self.max_features, float): return max(1, int(self.max_features * p))
        if self.max_features is None: return p
        return int(self.max_features)

    def fit(self, X, y):
        X = np.asarray(X, float); y = np.asarray(y); n, p = X.shape
        if not self.regression:
            self.classes_, y = np.unique(y, return_inverse=True); self.K = len(self.classes_)
        else:
            self.K = 1
        mf = self._mf(p); rng = np.random.default_rng(self.random_state)
        self.trees = []; imp = np.zeros(p)
        oob_sum = np.zeros((n, self.K)) if not self.regression else np.zeros(n)
        oob_cnt = np.zeros(n)
        for b in range(self.n_estimators):
            idx = rng.integers(0, n, n) if self.bootstrap else np.arange(n)
            tr = _Tree(mf, self.min_samples_leaf, self.max_depth,
                       np.random.default_rng(rng.integers(1 << 31)), self.regression, self.K).fit(X[idx], y[idx])
            self.trees.append(tr); imp += tr.importances
            if self.bootstrap:
                oob = np.ones(n, bool); oob[np.unique(idx)] = False
                if oob.any():
                    pr = tr.predict_proba(X[oob])
                    if self.regression: oob_sum[oob] += pr.reshape(-1)
                    else: oob_sum[oob] += pr
                    oob_cnt[oob] += 1
        s = imp.sum(); self.feature_importances_ = imp / s if s > 0 else imp
        self.n_features_ = p
        if self.bootstrap:
            seen = oob_cnt > 0
            if self.regression:
                self.oob_prediction_ = np.where(seen, oob_sum / np.maximum(oob_cnt, 1), np.nan)
            else:
                self.oob_decision_function_ = oob_sum / np.maximum(oob_cnt, 1)[:, None]
                self.oob_seen_ = seen
        return self

    def predict_proba(self, X):
        X = np.asarray(X, float)
        return np.mean([t.predict_proba(X) for t in self.trees], axis=0)

    def predict(self, X):
        if self.regression:
            X = np.asarray(X, float)
            return np.mean([t.predict_proba(X).reshape(-1) for t in self.trees], axis=0)
        return self.classes_[np.argmax(self.predict_proba(X), axis=1)]