import numpy as np, warnings
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
warnings.simplefilter("ignore")

f = np.sin
SIG = 0.3
N = 40
COLORS = ["#d62728", "#e6a800", "#2ca02c"]

def fit(x, y, d):
    return np.polyfit(x/3, y, d)          # scale x for numerical stability
def pred(c, x):
    return np.polyval(c, x/3)

# ---- same three datasets as datasets.png (same seed, same draw order)
rng = np.random.default_rng(7)
three = []
for _ in range(3):
    x = rng.uniform(-3, 3, N)
    y = f(x) + rng.normal(0, SIG, N)
    three.append((x, y))

xs = np.linspace(-3, 3, 400)
fig, ax = plt.subplots(figsize=(8, 5), dpi=160)
ax.plot(xs, f(xs), color="#222222", lw=2.2, label=r"$f(x)=\sin(x)$", zorder=1)
for i, ((x, y), col) in enumerate(zip(three, COLORS)):
    ax.scatter(x, y, s=34, color=col, edgecolor="#222222", linewidth=0.5, alpha=0.85, label=f"Dataset {i+1}", zorder=2)
ax.set_xlabel("x"); ax.set_ylabel("y"); ax.set_title(r"Three datasets drawn from $y = f(x) + \epsilon$")
ax.grid(alpha=0.25); ax.legend(frameon=False)
for sp in ("top", "right"): ax.spines[sp].set_visible(False)
fig.tight_layout(); fig.savefig("datasets.png"); plt.close(fig)
style = dict(sharex=True, sharey=True)

# ---- figure 2: fits of three model complexities on the 3 datasets
degs = [1, 3, 8]
titles = {1: "Degree 1 (a straight line)", 3: "Degree 3", 8: "Degree 8"}
fig, axes = plt.subplots(1, 3, figsize=(13, 4.3), dpi=160, **style)
for ax, d in zip(axes, degs):
    ax.plot(xs, f(xs), color="#222222", lw=2.2, zorder=1)
    for (x, y), col in zip(three, COLORS):
        ax.scatter(x, y, s=22, color=col, edgecolor="#222222", linewidth=0.4, alpha=0.8, zorder=2)
        ax.plot(xs, pred(fit(x, y, d), xs), color=col, lw=2, zorder=3)
    ax.set_title(titles[d]); ax.set_xlabel("x"); ax.grid(alpha=0.25)
    for s in ("top", "right"): ax.spines[s].set_visible(False)
axes[0].set_ylabel("y"); axes[0].set_ylim(-2, 2)
fig.tight_layout(); fig.savefig("fits.png"); plt.close(fig)

# ---- big simulation
rng = np.random.default_rng(123)
R = 4000
maxd = 8
grid = np.linspace(-3, 3, 61)
test_x = rng.uniform(-3, 3, 4000)

fhat_grid = {d: np.empty((R, len(grid))) for d in range(1, maxd+1)}
train_mse = {d: np.empty(R) for d in range(1, maxd+1)}
x0 = 2.0
fhat_x0 = {d: np.empty(R) for d in range(1, maxd+1)}
y0 = f(x0) + rng.normal(0, SIG, R)            # one fresh noisy test label per dataset
test_mse_pool = {d: np.empty(R) for d in range(1, maxd+1)}
ytest_pool = f(test_x) + rng.normal(0, SIG, test_x.size)
for r in range(R):
    x = rng.uniform(-3, 3, N); y = f(x) + rng.normal(0, SIG, N)
    for d in range(1, maxd+1):
        c = fit(x, y, d)
        fhat_grid[d][r] = pred(c, grid)
        fhat_x0[d][r] = pred(c, np.array([x0]))[0]
        train_mse[d][r] = np.mean((y - pred(c, x))**2)
        test_mse_pool[d][r] = np.mean((f(grid) + rng.normal(0, SIG, grid.size) - fhat_grid[d][r])**2)

# ---- numbers
fx0 = f(x0)
print("At x0=2, f(x0)=%.4f, sigma^2=%.2f" % (f(2.0), SIG**2))
rows = []
for d in range(1, maxd+1):
    m = fhat_x0[d].mean()
    b2 = (m - fx0)**2
    v = fhat_x0[d].var()
    emp = np.mean((y0 - fhat_x0[d])**2)
    rows.append((d, b2, v, SIG**2, b2+v+SIG**2, emp))
    print("deg %d  bias2=%.4f var=%.4f noise=%.3f sum=%.4f empirical=%.4f" % rows[-1])

print("\nAveraged over x in [-3,3]")
bias2_d, var_d = [], []
for d in range(1, maxd+1):
    m = fhat_grid[d].mean(0)
    b2 = np.mean((m - f(grid))**2)
    v = np.mean(fhat_grid[d].var(0))
    bias2_d.append(b2); var_d.append(v)
    print("deg %d bias2=%.4f var=%.4f noise=%.3f sum=%.4f | train=%.4f test(pool)=%.4f" %
          (d, b2, v, SIG**2, b2+v+SIG**2, train_mse[d].mean(), test_mse_pool[d].mean()))

D = np.arange(1, maxd+1)
bias2_d = np.array(bias2_d); var_d = np.array(var_d)
tot = bias2_d + var_d + SIG**2
trn = np.array([train_mse[d].mean() for d in D]); tst = np.array([test_mse_pool[d].mean() for d in D])

fig, axes = plt.subplots(1, 2, figsize=(12, 4.4), dpi=160)
ax = axes[0]
ax.plot(D, bias2_d, color="#d62728", lw=2.2, marker="o", label=r"Bias$^2$")
ax.plot(D, var_d, color="#2c7fb8", lw=2.2, marker="o", label="Variance")
ax.axhline(SIG**2, color="#777777", lw=1.5, ls=":", label=r"Noise $\sigma^2$")
ax.plot(D, tot, color="#222222", lw=2.6, marker="o", label="Bias² + Variance + σ²")
ax.set_yscale("log"); ax.set_xlabel("Model complexity (polynomial degree)"); ax.set_ylabel("Error (log scale)")
ax.set_title("Where the test error comes from"); ax.grid(alpha=0.25, which="both"); ax.legend(frameon=False, fontsize=8)
ax = axes[1]
ax.plot(D, trn, color="#2ca02c", lw=2.2, marker="o", label="Training error")
ax.plot(D, tst, color="#222222", lw=2.2, marker="o", label="Test error")
ax.set_yscale("log"); ax.set_xlabel("Model complexity (polynomial degree)")
ax.set_title("Training error vs test error"); ax.grid(alpha=0.25, which="both"); ax.legend(frameon=False)
for a in axes:
    for s in ("top", "right"): a.spines[s].set_visible(False)
fig.tight_layout(); fig.savefig("tradeoff.png"); plt.close(fig)