# Experiments for "Cross-Validation and the Bootstrap": every number in the post comes from here (numpy 2.1, seeded).
#   python resampling.py [OUT.json]
import sys, json, numpy as np
from numpy.polynomial import legendre as L

rng = np.random.default_rng(0)
f = lambda x: np.sin(2 * np.pi * x) + 0.5 * x; SIG = 0.3
fit = lambda x, y, d: L.legfit(2 * x - 1, y, d); pred = lambda c, x: L.legval(2 * x - 1, c)
def kfold_mse(x, y, d, k, r=None):
    n = len(x); idx = (r or rng).permutation(n); folds = np.array_split(idx, k); e = []
    for te in folds:
        tr = np.setdiff1d(idx, te); c = fit(x[tr], y[tr], d); e.append(np.mean((y[te] - pred(c, x[te])) ** 2))
    return float(np.mean(e))                                   # CV_(k) = (1/k) * sum of the k fold MSEs (equal-sized folds)
def loocv_shortcut(x, y, d):
    V = L.legvander(2 * x - 1, d); H = V @ np.linalg.pinv(V); r = y - H @ y; return float(np.mean((r / (1 - np.diag(H))) ** 2))
out = {}
# 1. one data set: true test MSE vs validation-set, LOOCV, 5- and 10-fold estimates over polynomial degree
n = 100; x = rng.uniform(0, 1, n); y = f(x) + rng.normal(0, SIG, n); xt = rng.uniform(0, 1, 100000); yt = f(xt) + rng.normal(0, SIG, 100000)
D = list(range(1, 11)); true = [np.mean((yt - pred(fit(x, y, d), xt)) ** 2) for d in D]
val = [[np.mean((y[te] - pred(fit(x[tr], y[tr], d), x[te])) ** 2) for d in D] for tr, te in ((p[:50], p[50:]) for p in (rng.permutation(n) for _ in range(10)))]
loo = [loocv_shortcut(x, y, d) for d in D]
brute = [np.mean([(y[i] - pred(fit(np.delete(x, i), np.delete(y, i), d), x[i])) ** 2 for i in range(n)]) for d in (3, 8)]
k5 = [kfold_mse(x, y, d, 5) for d in D]; k10 = [kfold_mse(x, y, d, 10) for d in D]
print("1. one data set (n = 100): degree, true test MSE of the fit on all 100 points, estimates")
for i, d in enumerate(D): print(f"   degree {d:2d}: true {true[i]:.4f}  validation (10 splits) {min(v[i] for v in val):.3f}-{max(v[i] for v in val):.3f}  LOOCV {loo[i]:.4f}  5-fold {k5[i]:.4f}  10-fold {k10[i]:.4f}")
print(f"   LOOCV shortcut vs refitting 100 times: degree 3 {loo[2]:.10f} vs {brute[0]:.10f}; degree 8 {loo[7]:.10f} vs {brute[1]:.10f}")
print(f"   chosen degree: true {D[int(np.argmin(true))]}, LOOCV {D[int(np.argmin(loo))]}, 5-fold {D[int(np.argmin(k5))]}, 10-fold {D[int(np.argmin(k10))]}, validation splits {[D[int(np.argmin(v))] for v in val]}")
out["curves"] = {"deg": D, "true": true, "val": val, "loo": loo, "k5": k5, "k10": k10}

# 2. what does CV estimate? 1,000 data sets of n = 50, degree 3
R, n, rows = 1000, 50, []
for _ in range(R):
    x = rng.uniform(0, 1, n); y = f(x) + rng.normal(0, SIG, n); c = fit(x, y, 3)
    err = np.mean((yt[:20000] - pred(c, xt[:20000])) ** 2)                      # this model's own test error
    p = rng.permutation(n); vs = np.mean((y[p[25:]] - pred(fit(x[p[:25]], y[p[:25]], 3), x[p[25:]])) ** 2)
    rows.append((err, vs, kfold_mse(x, y, 3, 5), kfold_mse(x, y, 3, 10), loocv_shortcut(x, y, 3)))
A = np.array(rows); avg = A[:, 0].mean()
print(f"2. degree 3, n = 50, 1,000 data sets: average true test error {avg:.4f} (sd across data sets {A[:, 0].std():.4f})")
for j, nm in ((1, "validation set (25/25)"), (2, "5-fold"), (3, "10-fold"), (4, "LOOCV")):
    print(f"   {nm:22s}: mean {A[:, j].mean():.4f} (bias {A[:, j].mean() - avg:+.4f}), sd {A[:, j].std():.4f}, corr with this model's own error {np.corrcoef(A[:, j], A[:, 0])[0, 1]:+.3f}")

# 3. the wrong and the right way: screening predictors before cross-validation (pure noise)
def nn1_err(Xtr, ytr, Xte, yte):
    d = ((Xte[:, None, :] - Xtr[None]) ** 2).sum(-1); return np.mean(ytr[d.argmin(1)] != yte)
def screen(X, y, m):
    Xc = X - X.mean(0); yc = y - y.mean(); c = np.abs(Xc.T @ yc) / (np.linalg.norm(Xc, axis=0) * np.linalg.norm(yc)); return np.argsort(-c)[:m]
wrong, right = [], []
for _ in range(50):
    X = rng.normal(0, 1, (50, 5000)); y = np.repeat([0, 1], 25)[rng.permutation(50)]
    top = screen(X, y, 100); folds = np.array_split(rng.permutation(50), 5); ew, er = [], []
    for te in folds:
        tr = np.setdiff1d(np.arange(50), te)
        ew.append(nn1_err(X[tr][:, top], y[tr], X[te][:, top], y[te]))           # wrong: screened on all 50, including the test fold
        t2 = screen(X[tr], y[tr], 100); er.append(nn1_err(X[tr][:, t2], y[tr], X[te][:, t2], y[te]))   # right: screened inside the fold
    wrong.append(np.mean(ew)); right.append(np.mean(er))
print(f"3. 50 samples, 5,000 noise predictors, labels random, keep the 100 most correlated, 1-nearest neighbour, 5-fold CV (50 data sets):")
print(f"   screening before CV: CV error {np.mean(wrong):.3f}; screening inside each fold: {np.mean(right):.3f}; true error 0.5")

# 4. the bootstrap: the portfolio fraction alpha (true 0.6)
alpha = lambda X, Y: (np.var(Y, ddof=1) - np.cov(X, Y)[0, 1]) / (np.var(X, ddof=1) + np.var(Y, ddof=1) - 2 * np.cov(X, Y)[0, 1])
C = np.array([[1, 0.5], [0.5, 1.25]]); pop = [alpha(*rng.multivariate_normal([0, 0], C, 100).T) for _ in range(1000)]
S = rng.multivariate_normal([0, 0], C, 100); a0 = alpha(*S.T); boot = []
for _ in range(1000): i = rng.integers(0, 100, 100); boot.append(alpha(*S[i].T))
boot = np.array(boot)
print(f"4. alpha: 1,000 new samples from the population: mean {np.mean(pop):.4f}, sd {np.std(pop, ddof=1):.4f}")
print(f"   one sample: alpha-hat {a0:.4f}; 1,000 bootstrap resamples: sd {boot.std(ddof=1):.4f}, 95% percentile interval [{np.percentile(boot, 2.5):.3f}, {np.percentile(boot, 97.5):.3f}]")
bse = []
for _ in range(200):
    S2 = rng.multivariate_normal([0, 0], C, 100); bse.append(np.std([alpha(*S2[rng.integers(0, 100, 100)].T) for _ in range(500)], ddof=1))
print(f"   bootstrap SE from 200 different samples (500 resamples each): mean {np.mean(bse):.4f}, middle 90% [{np.percentile(bse, 5):.4f}, {np.percentile(bse, 95):.4f}]")
out["boot"] = {"sample": np.round(S, 4).tolist(), "pop": np.round(pop, 4).tolist(), "a0": a0}
# 5. bootstrap SE of the mtcars slope vs the formula
sys.path.insert(0, ".")
from mtcars import MTCARS
d = np.array([r[1:] for r in MTCARS]); mpg, wt = d[:, 0], d[:, 4]; X1 = np.c_[np.ones(32), wt]
b = np.linalg.lstsq(X1, mpg, rcond=None)[0]; r = mpg - X1 @ b; se = np.sqrt(r @ r / 30 / ((wt - wt.mean()) ** 2).sum())
bs = []
for _ in range(5000): i = rng.integers(0, 32, 32); bs.append(np.linalg.lstsq(X1[i], mpg[i], rcond=None)[0][1])
print(f"5. mtcars slope of wt: formula SE {se:.4f}; bootstrap (5,000 resamples of cars) SE {np.std(bs, ddof=1):.4f}; percentile 95% [{np.percentile(bs, 2.5):.3f}, {np.percentile(bs, 97.5):.3f}]")
# 6. how much of the data is in a bootstrap sample
for n in (10, 100, 1000):
    sim = np.mean([len(np.unique(rng.integers(0, n, n))) / n for _ in range(2000)])
    print(f"6. n = {n:4d}: P(an observation is in a bootstrap sample) = {1 - (1 - 1 / n) ** n:.4f}, simulated share {sim:.4f}")
print(f"   limit 1 - 1/e = {1 - np.exp(-1):.4f}")
# 7. where the bootstrap fails: the maximum of a uniform sample
cov, same = 0, []
for _ in range(1000):
    u = rng.uniform(0, 1, 50); m = u.max(); bm = np.array([u[rng.integers(0, 50, 50)].max() for _ in range(1000)])
    same.append(np.mean(bm == m)); lo, hi = 2 * m - np.percentile(bm, 97.5), 2 * m - np.percentile(bm, 2.5); cov += lo <= 1 <= hi
print(f"7. maximum of 50 uniforms: share of bootstrap maxima equal to the sample maximum {np.mean(same):.3f} (theory {1 - (1 - 1 / 50) ** 50:.3f}); "
      f"95% basic bootstrap interval covers the true maximum 1 in {cov / 10:.1f}% of 1,000 samples")
if len(sys.argv) > 1: json.dump(out, open(sys.argv[1], "w"), separators=(",", ":"), default=float)