# Experiments for "Linear Regression, Read Correctly": every number in the post comes from here (numpy 2.1, scipy 1.18, seeded).
#   python linear_regression.py [OUT.json]  ->  prints the results; writes the t-quantile table and mtcars for the lab
import sys, json, numpy as np
from scipy import stats
from mtcars import MTCARS

def ols(X, y):
    """Least squares with an intercept column already in X: coefficients, SE, t, p, 95% CI, RSE, R^2, adjusted R^2, F and its p."""
    n, k = X.shape; b, *_ = np.linalg.lstsq(X, y, rcond=None); r = y - X @ b; rss = r @ r; tss = ((y - y.mean()) ** 2).sum()
    s2 = rss / (n - k); se = np.sqrt(np.diag(s2 * np.linalg.inv(X.T @ X))); t = b / se; p = 2 * stats.t.sf(np.abs(t), n - k)
    q = stats.t.ppf(0.975, n - k); F = ((tss - rss) / (k - 1)) / (rss / (n - k)) if k > 1 else np.nan
    return dict(b=b, se=se, t=t, p=p, lo=b - q * se, hi=b + q * se, rse=np.sqrt(s2), r2=1 - rss / tss, adj=1 - (rss / (n - k)) / (tss / (n - 1)),
                F=F, pF=stats.f.sf(F, k - 1, n - k) if k > 1 else np.nan, df=n - k)
def show(name, res, names):
    print(f"   {name}: RSE {res['rse']:.3f} (df {res['df']}), R^2 {res['r2']:.4f}, adjusted {res['adj']:.4f}, F {res['F']:.2f} (p {res['pF']:.2e})")
    for i, nm in enumerate(names):
        print(f"      {nm:10s} {res['b'][i]:9.4f}  SE {res['se'][i]:.4f}  t {res['t'][i]:7.2f}  p {res['p'][i]:.2e}  95% CI [{res['lo'][i]:.4f}, {res['hi'][i]:.4f}]")

d = np.array([r[1:] for r in MTCARS]); mpg, cyl, disp, hp, wt = d.T; one = np.ones(32)
print("1. mtcars")
show("mpg ~ wt", ols(np.c_[one, wt], mpg), ["intercept", "wt"])
print(f"   r(mpg, wt) = {np.corrcoef(mpg, wt)[0, 1]:.4f}; r^2 = {np.corrcoef(mpg, wt)[0, 1] ** 2:.4f}; pearson test p {stats.pearsonr(mpg, wt)[1]:.2e}")
show("mpg ~ hp", ols(np.c_[one, hp], mpg), ["intercept", "hp"])
show("mpg ~ wt + hp", ols(np.c_[one, wt, hp], mpg), ["intercept", "wt", "hp"])
print(f"   r(wt, hp) = {np.corrcoef(wt, hp)[0, 1]:.3f}; VIF = {1 / (1 - np.corrcoef(wt, hp)[0, 1] ** 2):.2f}")
show("mpg ~ wt + disp + hp + cyl", ols(np.c_[one, wt, disp, hp, cyl], mpg), ["intercept", "wt", "disp", "hp", "cyl"])
Xo = np.c_[wt, disp, hp, cyl]; vif = [1 / (1 - ols(np.c_[one, np.delete(Xo, j, 1)], Xo[:, j])["r2"]) for j in range(4)]
print("   VIF wt, disp, hp, cyl:", [round(v, 2) for v in vif])
c6, c8 = (cyl == 6).astype(float), (cyl == 8).astype(float)
show("mpg ~ wt + cyl6 + cyl8 (dummies, 4 cylinders = baseline)", ols(np.c_[one, wt, c6, c8], mpg), ["intercept", "wt", "cyl6", "cyl8"])
show("mpg ~ wt + cyl (cylinders as a number)", ols(np.c_[one, wt, cyl], mpg), ["intercept", "wt", "cyl"])

X1 = np.c_[one, wt]; res = ols(X1, mpg); x0 = np.array([1.0, 3.0]); yh = x0 @ res["b"]; q = stats.t.ppf(0.975, 30)
sef = res["rse"] * np.sqrt(x0 @ np.linalg.inv(X1.T @ X1) @ x0)
print(f"   a 3,000 lb car: predicted mpg {yh:.2f}; 95% confidence interval for the mean [{yh - q * sef:.2f}, {yh + q * sef:.2f}]; "
      f"95% prediction interval for one car [{yh - q * np.hypot(res['rse'], sef):.2f}, {yh + q * np.hypot(res['rse'], sef):.2f}]")
X2 = np.c_[one, wt, wt ** 2]; show("mpg ~ wt + wt^2", ols(X2, mpg), ["intercept", "wt", "wt^2"])

rng = np.random.default_rng(0)
# 2. what "95% confidence" means: coverage over repeated samples, exact t vs +-2 SE vs +-1.96 SE
print("2. coverage of intervals for the slope (true beta1 = 0.5), 20,000 simulated data sets per n")
for n in (5, 10, 30, 100):
    cov = {"t": 0, "2": 0, "1.96": 0}; q = stats.t.ppf(0.975, n - 2)
    for _ in range(20000):
        x = rng.uniform(0, 10, n); y = 2 + 0.5 * x + rng.normal(0, 1, n); xc = x - x.mean(); b1 = (xc @ y) / (xc @ xc)
        r = y - (y.mean() + b1 * xc); se = np.sqrt((r @ r) / (n - 2) / (xc @ xc))
        for k, m in (("t", q), ("2", 2.0), ("1.96", 1.96)): cov[k] += abs(b1 - 0.5) <= m * se
    print(f"   n={n:3d}: t quantile {q:.3f}; coverage exact t {cov['t'] / 200:.1f}%, +-2 SE {cov['2'] / 200:.1f}%, +-1.96 SE {cov['1.96'] / 200:.1f}%")

# 3. p-values: uniform under the null; tiny effects become 'significant' with enough data
ps = []
for _ in range(20000):
    x = rng.normal(0, 1, 30); y = rng.normal(0, 1, 30); ps.append(stats.linregress(x, y).pvalue)
ps = np.array(ps); print(f"3. no effect, n = 30, 20,000 data sets: share with p < 0.05 {np.mean(ps < 0.05):.4f}, p < 0.01 {np.mean(ps < 0.01):.4f}; deciles {np.round(np.quantile(ps, [0.1, 0.5, 0.9]), 3)}")
x = rng.normal(0, 1, 100_000); y = 0.02 * x + rng.normal(0, 1, 100_000); lr = stats.linregress(x, y)
print(f"   tiny effect, n = 100,000: slope {lr.slope:.4f}, p {lr.pvalue:.1e}, R^2 {lr.rvalue ** 2:.5f}")
print(f"   t thresholds for p = 0.05 / 0.01: df 10 {stats.t.ppf(0.975, 10):.3f} / {stats.t.ppf(0.995, 10):.3f}, df 30 {stats.t.ppf(0.975, 30):.3f} / {stats.t.ppf(0.995, 30):.3f}, "
      f"df inf {stats.norm.ppf(0.975):.3f} / {stats.norm.ppf(0.995):.3f}")

# 4. forward selection on pure noise: p-values after selection are not p-values
n, P, hits, nsel, first = 100, 20, 0, [], []
for _ in range(2000):
    X = rng.normal(0, 1, (n, P)); y = rng.normal(0, 1, n); chosen = []
    while True:
        best = None
        for j in set(range(P)) - set(chosen):
            res = ols(np.c_[np.ones(n), X[:, chosen + [j]]], y); pj = res["p"][-1]
            if best is None or pj < best[1]: best = (j, pj)
        if best is None or best[1] >= 0.05: break
        if not chosen: first.append(best[1])
        chosen.append(best[0])
    hits += bool(chosen); nsel.append(len(chosen))
print(f"4. forward selection (enter if p < 0.05), n = 100, 20 noise predictors, 2,000 data sets: at least one selected {hits / 20:.1f}% "
      f"(1 - 0.95^20 = {100 * (1 - 0.95 ** 20):.1f}%), mean number selected {np.mean(nsel):.2f}, median p of the first selected {np.median(first):.4f}")

# 5. a predictor with zero correlation that matters: suppression
n = 200; x1 = rng.normal(0, 1, n); x2 = 0.7 * x1 + np.sqrt(1 - 0.49) * rng.normal(0, 1, n); y = x1 - 0.7 * x2 + 0.5 * rng.normal(0, 1, n)
s1 = ols(np.c_[np.ones(n), x2], y); s2 = ols(np.c_[np.ones(n), x1, x2], y)
print(f"5. suppression: r(y, x2) = {np.corrcoef(y, x2)[0, 1]:.3f}; simple slope of x2 {s1['b'][1]:.3f} (p {s1['p'][1]:.2f}); in y ~ x1 + x2: "
      f"x2 {s2['b'][2]:.3f} (t {s2['t'][2]:.1f}, p {s2['p'][2]:.1e}), x1 {s2['b'][1]:.3f}")

# 6. R^2 never falls when predictors are added; adjusted and test R^2 do
n, ks, tr, ad, te = 50, [0, 1, 2, 5, 10, 20, 30, 40], {}, {}, {}
for k in ks: tr[k], ad[k], te[k] = [], [], []
for _ in range(2000):
    x = rng.normal(0, 1, n); Z = rng.normal(0, 1, (n, 40)); y = 1 + 0.6 * x + rng.normal(0, 1, n)
    xn = rng.normal(0, 1, 2000); Zn = rng.normal(0, 1, (2000, 40)); yn = 1 + 0.6 * xn + rng.normal(0, 1, 2000)
    for k in ks:
        X = np.c_[np.ones(n), x, Z[:, :k]]; res = ols(X, y); pr = np.c_[np.ones(2000), xn, Zn[:, :k]] @ res["b"]
        tr[k].append(res["r2"]); ad[k].append(res["adj"]); te[k].append(1 - ((yn - pr) ** 2).sum() / ((yn - yn.mean()) ** 2).sum())
print("6. one real predictor + k noise predictors, n = 50 (2,000 data sets): mean R^2 training / adjusted / test")
for k in ks: print(f"   k={k:2d}: {np.mean(tr[k]):.3f} / {np.mean(ad[k]):.3f} / {np.mean(te[k]):.3f}")

# 7. collinearity inflates standard errors
for rho in (0.0, 0.9, 0.99):
    se = []
    for _ in range(2000):
        x1 = rng.normal(0, 1, 100); x2 = rho * x1 + np.sqrt(1 - rho ** 2) * rng.normal(0, 1, 100); y = x1 + x2 + rng.normal(0, 1, 100)
        se.append(ols(np.c_[np.ones(100), x1, x2], y)["se"][1])
    print(f"7. corr(x1, x2) = {rho}: mean SE of beta1 {np.mean(se):.3f}; sqrt(VIF) = {1 / np.sqrt(1 - rho ** 2):.2f}")
if len(sys.argv) > 1:
    json.dump({"tq": [round(float(stats.t.ppf(0.975, k)), 6) for k in range(1, 201)], "mtcars": [[r[0], r[1], r[4], r[5]] for r in MTCARS]}, open(sys.argv[1], "w"), separators=(",", ":"))