#!/usr/bin/env python3
"""Double descent with random Fourier features on 12 noisy 1-D points.

Model: f(x) = sum_j theta_j phi_j(x), phi_j(x) = sqrt(2/p) cos(w_j x + b_j),
w_j ~ N(0, s^2), b_j ~ U[0, 2 pi).  For p < n we fit least squares; for p >= n
there are infinitely many interpolants and we take the minimum-norm one
(pseudo-inverse) -- the solution gradient descent from zero converges to.

As p grows the test error first falls (more capacity), then explodes near
p = n (the only interpolant is forced to wiggle wildly to hit the noisy
points: the feature matrix is square and badly conditioned), then falls again
for p >> n (many interpolants; the minimum-norm one is smooth).  A small
ridge penalty removes the peak.  See Belkin, Hsu, Ma & Mandal (2019).

  python double_descent.py [--quick]
Writes out/double_descent.png and prints the curve.
"""
from __future__ import annotations

import argparse
import os

import numpy as np

OUT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "out")

N_TRAIN = 12
NOISE = 0.1
SCALE = 15.0  # frequency scale s of the random features
P_GRID = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 18, 20, 24, 30, 40,
          60, 100, 150, 250, 400, 700, 1000, 2000]


def target(x):
    return np.sin(2 * np.pi * x)


def make_data(n=N_TRAIN, noise=NOISE, seed=0):
    """Jittered grid: one random point in the middle 70% of each of n cells,
    so inputs are random but never clumped (clumps make every p ill-posed)."""
    rng = np.random.default_rng(seed)
    x = (np.arange(n) + rng.uniform(0.15, 0.85, n)) / n
    y = target(x) + noise * rng.standard_normal(n)
    return x, y


def features(x, w, b):
    return np.sqrt(2.0 / len(w)) * np.cos(np.outer(x, w) + b)


def fit(F, y, ridge=0.0):
    """Min-norm least squares (ridge=0) or ridge regression, via the dual form
    when p > n (cheaper and identical)."""
    n, p = F.shape
    if ridge == 0.0:
        return np.linalg.pinv(F, rcond=1e-13) @ y
    if p > n:
        return F.T @ np.linalg.solve(F @ F.T + ridge * np.eye(n), y)
    return np.linalg.solve(F.T @ F + ridge * np.eye(p), F.T @ y)


def test_error_curve(x, y, p_grid=P_GRID, ridge=0.0, draws=100, scale=SCALE, seed=1):
    """Median (and quartiles) over feature draws of test MSE vs the clean target."""
    rng = np.random.default_rng(seed)
    xt = np.linspace(0, 1, 500)
    yt = target(xt)
    med, q1, q3 = [], [], []
    for p in p_grid:
        errs = []
        for _ in range(draws):
            w = rng.normal(0, scale, p)
            b = rng.uniform(0, 2 * np.pi, p)
            theta = fit(features(x, w, b), y, ridge)
            errs.append(np.mean((features(xt, w, b) @ theta - yt) ** 2))
        med.append(np.median(errs))
        q1.append(np.percentile(errs, 25))
        q3.append(np.percentile(errs, 75))
    return np.array(med), np.array(q1), np.array(q3)


def main(quick=False):
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    os.makedirs(OUT, exist_ok=True)
    draws = 30 if quick else 100
    x, y = make_data()
    curves = {}
    for ridge in (0.0, 1e-3):
        curves[ridge] = test_error_curve(x, y, ridge=ridge, draws=draws)
    med0 = curves[0.0][0]
    p = np.array(P_GRID)
    window = (p >= N_TRAIN - 2) & (p <= 2 * N_TRAIN)  # around the threshold p = n
    i_peak = int(np.argmax(np.where(window, med0, -np.inf)))
    under = p < N_TRAIN
    i_best_under = int(np.argmin(np.where(under, med0, np.inf)))
    print(f"n = {N_TRAIN} points, noise sd {NOISE}, feature scale {SCALE}, {draws} feature draws per p")
    print(f"{'p':>6} {'min-norm':>10} {'ridge 1e-3':>11}")
    for j, pp in enumerate(P_GRID):
        print(f"{pp:>6} {med0[j]:>10.4f} {curves[1e-3][0][j]:>11.4f}")
    print(f"best under-parameterised: p={p[i_best_under]}  test MSE {med0[i_best_under]:.4f}")
    print(f"interpolation peak:       p={p[i_peak]}  test MSE {med0[i_peak]:.4f}"
          f"  ({med0[i_peak] / med0[i_best_under]:.0f}x worse)")
    print(f"far over-parameterised:   p={p[-1]}  test MSE {med0[-1]:.4f}")
    print(f"ridge 1e-3 near p = n:    max test MSE {curves[1e-3][0][window].max():.4f} (no peak)")

    fig, axes = plt.subplots(1, 2, figsize=(13, 4.6))
    ax = axes[0]
    for ridge, lab, c in ((0.0, "min-norm interpolation (ridge 0)", "C3"),
                          (1e-3, "ridge λ=1e-3", "C0")):
        m, a, b = curves[ridge]
        ax.plot(p, m, "o-", ms=3, color=c, label=lab)
        ax.fill_between(p, a, b, color=c, alpha=0.15, lw=0)
    ax.axvline(N_TRAIN, color="k", ls="--", lw=1)
    ax.text(N_TRAIN * 1.08, med0.min() * 0.75, "p = n", fontsize=9)
    ax.set_xscale("log")
    ax.set_yscale("log")
    ax.set_xlabel("number of random features p")
    ax.set_ylabel("test MSE (median, IQR over feature draws)")
    ax.set_title("Double descent: random Fourier features, n = 12")
    ax.legend()
    ax.grid(alpha=0.3, which="both")

    ax = axes[1]
    rng = np.random.default_rng(5)
    xt = np.linspace(0, 1, 500)
    ax.plot(xt, target(xt), "k-", lw=1, label="truth sin(2πx)")
    for pp, c in ((6, "C2"), (N_TRAIN, "C3"), (2000, "C0")):
        w = rng.normal(0, SCALE, pp)
        b = rng.uniform(0, 2 * np.pi, pp)
        th = fit(features(x, w, b), y)
        ax.plot(xt, features(xt, w, b) @ th, color=c, lw=1.4, label=f"p = {pp}")
    ax.plot(x, y, "ko", ms=5, label="training points")
    ax.set_ylim(-3, 3)
    ax.set_title("One fit at each regime (min-norm)")
    ax.legend(fontsize=8, loc="lower left")
    ax.grid(alpha=0.3)
    fig.tight_layout()
    path = os.path.join(OUT, "double_descent.png")
    fig.savefig(path, dpi=130)
    plt.close(fig)
    print(f"saved {path}")
    return {"p": p, "minnorm": med0, "ridge": curves[1e-3][0]}


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--quick", action="store_true")
    main(ap.parse_args().quick)