import json
import time
from pathlib import Path
import numpy as np

SEED = 3104
rng_global = np.random.default_rng(SEED)


def sample_rff(n, d=2, mode="gaussian", rng=None):
    rng = np.random.default_rng() if rng is None else rng
    if mode == "gaussian":
        # A conventional isotropic Gaussian RFF control with deliberately broad bandwidth.
        w = rng.normal(0.0, 3.0, size=(n, d))
    elif mode == "uniform":
        w = rng.uniform(-6.0, 6.0, size=(n, d))
    elif mode == "gevery_matched":
        # density proportional to exp(-2*kappa*||w||^(1/s)); s=1, kappa=.75.
        # In d dimensions, y=2*kappa*r^(1/s) is Gamma(d*s, 1).
        s, kappa = 1.0, 0.75
        y = rng.gamma(shape=d * s, scale=1.0, size=n)
        r = (y / (2.0 * kappa)) ** s
        direction = rng.normal(size=(n, d))
        direction /= np.linalg.norm(direction, axis=1, keepdims=True)
        w = r[:, None] * direction
    elif mode == "sobolev_matched":
        # density proportional to (1+||w||^2)^(-bar_s), bar_s>d/2.
        bar_s = 3.0
        u = rng.beta(d / 2.0, bar_s - d / 2.0, size=n)
        r = np.sqrt(u / (1.0 - u))
        direction = rng.normal(size=(n, d))
        direction /= np.linalg.norm(direction, axis=1, keepdims=True)
        w = r[:, None] * direction
    else:
        raise ValueError(mode)
    b = rng.uniform(0.0, 2.0 * np.pi, size=n)
    return w, b


def features(x, w, b):
    return np.sqrt(2.0) * np.cos(x @ w.T + b[None, :])


def target(x):
    # Smooth (indeed analytic) low-frequency periodic target, representative of a
    # coordinate field whose spectral energy is concentrated near the origin.
    return (0.80 * np.cos(x[:, 0] + 0.20) +
            0.55 * np.cos(x[:, 1] - 0.70) +
            0.35 * np.cos(2.0 * x[:, 0] + x[:, 1] + 0.40) +
            0.18 * np.cos(3.0 * x[:, 1] - 0.30))


def whiten(z, eps_scale=1e-5):
    g = (z.T @ z) / z.shape[0]
    eps = eps_scale * np.trace(g) / g.shape[0]
    vals, vecs = np.linalg.eigh(g + eps * np.eye(g.shape[0]))
    vals = np.maximum(vals, eps)
    zw = z @ (vecs * (1.0 / np.sqrt(vals))) @ vecs.T
    return zw, np.linalg.eigvalsh(g), float(eps)


def ridge_fit(z, y, alpha=1e-5):
    a = z.T @ z + alpha * np.eye(z.shape[1])
    return np.linalg.solve(a, z.T @ y)


def run():
    rng = np.random.default_rng(SEED)
    # Random train/test locations avoid exploiting a grid-specific aliasing artifact.
    xtr = rng.uniform(-np.pi, np.pi, size=(2048, 2))
    xte = rng.uniform(-np.pi, np.pi, size=(4096, 2))
    ytr, yte = target(xtr), target(xte)
    modes = ["gaussian", "uniform", "gevery_matched"]
    counts = [64, 128, 256, 512]
    rows = []
    t0 = time.perf_counter()
    # Three independent draws provide a small reproducibility check.
    for n in counts:
        for mode in modes:
            for whitened in [False, True]:
                errs, conds = [], []
                for rep in range(3):
                    rr = np.random.default_rng(SEED + 10000 * rep + n + hash(mode) % 997)
                    w, b = sample_rff(n, mode=mode, rng=rr)
                    ztr, zte = features(xtr, w, b), features(xte, w, b)
                    raw_eigs = np.linalg.eigvalsh((ztr.T @ ztr) / len(xtr))
                    cond = float(raw_eigs[-1] / max(raw_eigs[0], 1e-15))
                    if whitened:
                        ztr, _, _ = whiten(ztr)
                        # Whitening is an input-coordinate transform learned on train;
                        # use the same transform for test by recomputing explicitly below.
                        g = (features(xtr, w, b).T @ features(xtr, w, b)) / len(xtr)
                        eps = 1e-5 * np.trace(g) / n
                        vals, vecs = np.linalg.eigh(g + eps * np.eye(n))
                        zte = zte @ (vecs * (1.0 / np.sqrt(np.maximum(vals, eps)))) @ vecs.T
                    beta = ridge_fit(ztr, ytr)
                    pred = zte @ beta
                    errs.append(float(np.sqrt(np.mean((pred - yte) ** 2))))
                    conds.append(cond)
                rows.append({"N": n, "mode": mode, "whitened": whitened,
                             "test_rmse_mean": float(np.mean(errs)),
                             "test_rmse_std": float(np.std(errs)),
                             "raw_gram_condition_mean": float(np.mean(conds))})
    # Cheap mathematical sanity checks: empirical tails should decrease in the
    # expected order, and whitening should make the training Gram spectrum nearly flat.
    tail_rng = np.random.default_rng(SEED + 77)
    checks = {}
    for mode in ["gaussian", "gevery_matched", "sobolev_matched"]:
        w, _ = sample_rff(200000, mode=mode, rng=tail_rng)
        r = np.linalg.norm(w, axis=1)
        checks[mode] = {"median_radius": float(np.median(r)),
                        "q99_radius": float(np.quantile(r, .99)),
                        "mean_radius": float(np.mean(r))}
    w, b = sample_rff(256, mode="gevery_matched", rng=np.random.default_rng(SEED))
    zz = features(xtr, w, b)
    before = np.linalg.eigvalsh((zz.T @ zz) / len(xtr))
    zw, after, eps = whiten(zz)
    after_reg = np.linalg.eigvalsh((zw.T @ zw) / len(xtr))
    checks["whitening"] = {
        "raw_condition": float(before[-1] / max(before[0], 1e-15)),
        "whitened_condition": float(after_reg[-1] / max(after_reg[0], 1e-15)),
        "whitened_eigen_min": float(after_reg[0]),
        "whitened_eigen_max": float(after_reg[-1]),
        "ridge_epsilon": eps,
    }
    out = {"seed": SEED, "runtime_sec": time.perf_counter() - t0,
           "rows": rows, "sanity_checks": checks}
    Path("results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))


if __name__ == "__main__":
    run()