import json, math, time
from pathlib import Path
import numpy as np
from scipy.interpolate import CubicSpline
from sklearn.metrics import mean_squared_error

SEED = 469
rng = np.random.default_rng(SEED)


def fill_distance_1d(points, grid=None):
    """Numerical fill distance on [0,1] for a finite point set."""
    points = np.sort(np.asarray(points))
    if grid is None:
        grid = np.linspace(0.0, 1.0, 20001)
    return np.max(np.min(np.abs(grid[:, None] - points[None, :]), axis=1))


def verify_math():
    # Uniform discretizations have h approximately 1/(2m), while random
    # points have noticeably poorer coverage at the same cardinality.
    ms = np.array([4, 8, 16, 32, 64])
    uniform_h = np.array([fill_distance_1d(np.linspace(0, 1, int(m))) for m in ms])
    random_h = np.array([
        fill_distance_1d(np.sort(rng.uniform(0, 1, int(m)))) for m in ms
    ])
    slope = np.polyfit(np.log(ms), np.log(uniform_h), 1)[0]
    expected = 1.0 / (2.0 * (ms - 1))
    return {
        "m": ms.tolist(),
        "uniform_fill_distance": uniform_h.tolist(),
        "random_fill_distance": random_h.tolist(),
        "uniform_theory": expected.tolist(),
        "loglog_slope": float(slope),
        "max_uniform_theory_error": float(np.max(np.abs(uniform_h - expected))),
    }


def features(z):
    # Compact input class: z in [-1,1]^d, with bounded Fourier features.
    z = np.asarray(z)
    return np.concatenate([z, np.sin(np.pi*z), np.cos(np.pi*z)], axis=1)


def operator(z, x):
    """Smooth nonlinear coefficient-to-function operator on [0,1]."""
    z = np.asarray(z)
    x = np.asarray(x)
    # Coefficients decay slowly enough that refinement reveals new content.
    out = np.zeros((len(z), len(x)))
    for k in range(1, 13):
        phase = (k * x)[None, :]
        coeff = (np.sin((k + 0.35) * z[:, 0])
                 + 0.65 * np.cos((k + 1.2) * z[:, 1])
                 + 0.35 * z[:, 2] * z[:, 3]) / (k ** 1.15)
        out += coeff[:, None] * np.sin(2 * np.pi * phase)
    # nonlinear low-frequency component
    out += 0.30 * (z[:, 0] * z[:, 1])[:, None] * np.cos(2 * np.pi * x)[None, :]
    return out


def rbf_predict(xtr, ytr, xte, length=1.35, ridge=2e-3):
    # Small exact kernel regression is the offline operator A_off.
    def kernel(a, b):
        d2 = np.sum((a[:, None, :] - b[None, :, :]) ** 2, axis=2)
        return np.exp(-d2 / (2 * length * length))
    K = kernel(xtr, xtr)
    alpha = np.linalg.solve(K + ridge * np.eye(len(K)), ytr)
    return kernel(xte, xtr) @ alpha


def relative_l2(pred, truth):
    return float(np.linalg.norm(pred - truth) / np.linalg.norm(truth))


def run_experiment():
    # Verify the claimed fill-distance scaling before any learning experiment.
    math_check = verify_math()
    local = np.random.default_rng(SEED + 1)
    d = 4
    ntrain_max = 128
    ntest = 256
    ztrain = local.uniform(-1, 1, size=(ntrain_max, d))
    ztest = local.uniform(-1, 1, size=(ntest, d))
    ftrain = features(ztrain)
    ftest = features(ztest)
    xfine = np.linspace(0, 1, 513)
    truth_test = operator(ztest, xfine)
    # A fixed fine grid is the reference function for all reconstruction errors.
    resolutions = [16, 32, 64]
    kappas = [0.5, 1.0, 1.5]
    C = 1.0
    records = []
    for m in resolutions:
        x = np.linspace(0, 1, m)
        # Exact output observations and online cubic reconstruction form oracle.
        exact = operator(ztest, x)
        oracle = np.empty_like(truth_test)
        for j in range(ntest):
            oracle[j] = CubicSpline(x, exact[j], bc_type='natural')(xfine)
        oracle_err = relative_l2(oracle, truth_test)
        # Same computational protocol for every kappa.
        for kappa in kappas:
            N = min(ntrain_max, int(math.ceil(C * m ** kappa)))
            t0 = time.perf_counter()
            pred_obs = rbf_predict(ftrain[:N], operator(ztrain[:N], x), ftest)
            pred = np.empty_like(truth_test)
            for j in range(ntest):
                pred[j] = CubicSpline(x, pred_obs[j], bc_type='natural')(xfine)
            wall = time.perf_counter() - t0
            learned_err = relative_l2(pred, truth_test)
            feature_h = np.max(np.min(
                np.sum((ftest[:, None, :] - ftrain[None, :N, :]) ** 2, axis=2), axis=1
            ) ** 0.5)
            records.append({
                "m": m, "kappa": kappa, "N": N,
                "oracle_error": oracle_err,
                "learned_error": learned_err,
                "learned_oracle_gap": learned_err - oracle_err,
                "feature_fill_proxy": float(feature_h),
                "kernel_solve_flops_proxy": int(N ** 3),
                "wall_seconds": wall,
            })
    # Fixed-data baseline uses N equal to the largest budget at each m only for
    # the comparison requested by the idea: N=16 does not grow with resolution.
    fixed = []
    for m in resolutions:
        N = 16
        x = np.linspace(0, 1, m)
        t0 = time.perf_counter()
        obs = rbf_predict(ftrain[:N], operator(ztrain[:N], x), ftest)
        pred = np.empty_like(truth_test)
        for j in range(ntest):
            pred[j] = CubicSpline(x, obs[j], bc_type='natural')(xfine)
        fixed.append({"m": m, "N": N, "error": relative_l2(pred, truth_test),
                      "wall_seconds": time.perf_counter() - t0})
    return {"math_check": math_check, "records": records, "fixed_N_baseline": fixed}


if __name__ == "__main__":
    result = run_experiment()
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))