import json
import math
from pathlib import Path

import numpy as np


SEED = 2977


def gaussian_kernel(X, Y=None):
    if Y is None:
        Y = X
    d2 = ((X[:, None, :] - Y[None, :, :]) ** 2).sum(axis=2)
    return np.exp(-0.5 * d2)


def hermite_rule_2d(q):
    # Gauss-Hermite integrates exp(-t^2) f(t); t=sqrt(2) w gives
    # standard-normal integration, with weights normalized to sum to one.
    roots, weights = np.polynomial.hermite.hermgauss(q)
    nodes_1d = np.sqrt(2.0) * roots
    weights_1d = weights / np.sqrt(np.pi)
    nodes = np.array([(a, b) for a in nodes_1d for b in nodes_1d])
    weights_2d = np.array([wa * wb for wa in weights_1d for wb in weights_1d])
    weights_2d /= weights_2d.sum()
    return nodes, weights_2d


def quadrature_features(X, q):
    W, a = hermite_rule_2d(q)
    # G(w,x)=exp(i w^T x), Z_mi=sqrt(a_i)G(w_i,x_m).
    Z = np.exp(1j * (X @ W.T)) * np.sqrt(a)[None, :]
    return Z


def random_features(X, n, rng):
    W = rng.normal(size=(n, X.shape[1]))
    Z = np.exp(1j * (X @ W.T)) / np.sqrt(n)
    return Z


def eig_metrics(K, Kh, top=20):
    lam = np.linalg.eigvalsh(K)[::-1]
    lhat = np.linalg.eigvalsh((Kh + Kh.conj().T) / 2.0)[::-1]
    j = min(top, len(lam))
    top_abs = float(np.mean(np.abs(lam[:j] - lhat[:j])))
    top_rel = float(np.mean(np.abs(lam[:j] - lhat[:j]) / np.maximum(lam[:j], 1e-12)))
    op = float(np.linalg.norm(K - Kh, ord=2))
    # Weyl's inequality is checked for every ordered eigenvalue.
    weyl_max = float(np.max(np.abs(lam - lhat)))
    return {"top20_abs_eigen_error": top_abs, "top20_relative_eigen_error": top_rel,
            "operator_error": op, "weyl_bound_holds": bool(weyl_max <= op + 2e-8),
            "largest_exact_eigenvalue": float(lam[0]), "largest_approx_eigenvalue": float(lhat[0])}


def krr_mse(Xtr, ytr, Xte, yte, Ztr, Zte, ridge=1e-3):
    # Complex feature-space ridge regression; same representation for both methods.
    A = Ztr.conj().T @ Ztr + ridge * np.eye(Ztr.shape[1])
    theta = np.linalg.solve(A, Ztr.conj().T @ ytr)
    pred = np.real(Zte @ theta)
    return float(np.mean((pred - yte) ** 2))


def main():
    rng = np.random.default_rng(SEED)
    # Fixed points make all methods directly comparable. A smooth target is
    # deliberately compatible with the RBF kernel used by the Fourier integral.
    X = rng.normal(size=(120, 2))
    K = gaussian_kernel(X)
    rows = []
    for q in (3, 4, 6, 8):
        Zq = quadrature_features(X, q)
        mq = eig_metrics(K, Zq @ Zq.conj().T)
        random_runs = []
        for s in range(20):
            Zr = random_features(X, q * q, np.random.default_rng(SEED + 1000 * q + s))
            random_runs.append(eig_metrics(K, Zr @ Zr.conj().T))
        rr = {key: float(np.mean([r[key] for r in random_runs]))
              for key in ("top20_abs_eigen_error", "top20_relative_eigen_error", "operator_error",
                          "largest_approximate_eigenvalue") if key in random_runs[0]}
        rows.append({"features": q * q, "quadrature": mq, "random_mean": rr,
                     "random_std_top20_relative": float(np.std([r["top20_relative_eigen_error"] for r in random_runs]))})

    # Equal-width learning check at 64 features, with a held-out split.
    Xtr, Xte = X[:80], X[80:]
    ytr = np.sin(Xtr[:, 0]) + 0.25 * np.cos(1.5 * Xtr[:, 1])
    yte = np.sin(Xte[:, 0]) + 0.25 * np.cos(1.5 * Xte[:, 1])
    Zqtr, Zqte = quadrature_features(Xtr, 8), quadrature_features(Xte, 8)
    q_mse = krr_mse(Xtr, ytr, Xte, yte, Zqtr, Zqte)
    random_mses = []
    for s in range(20):
        # Generate the same frequencies for train and test in each replicate.
        W = np.random.default_rng(SEED + 7000 + s).normal(size=(64, 2))
        Ztr = np.exp(1j * (Xtr @ W.T)) / 8.0
        Zte = np.exp(1j * (Xte @ W.T)) / 8.0
        random_mses.append(krr_mse(Xtr, ytr, Xte, yte, Ztr, Zte))

    # Independent cheap algebra sanity check of Z Z* and the probability rule.
    W, a = hermite_rule_2d(4)
    Z = np.exp(1j * (X[:12] @ W.T)) * np.sqrt(a)[None, :]
    # Independent matrix-form evaluation of the same quadrature sum.
    direct = np.exp(1j * (X[:12] @ W.T)) @ np.diag(a) @ np.exp(-1j * (W @ X[:12].T))
    algebra_error = float(np.max(np.abs(Z @ Z.conj().T - direct)))

    result = {"seed": SEED, "kernel": "exp(-||x-y||^2/2)", "rows": rows,
              "learning_64_features": {"quadrature_mse": q_mse,
                                        "random_mean_mse": float(np.mean(random_mses)),
                                        "random_std_mse": float(np.std(random_mses))},
              "sanity": {"weights_sum": float(a.sum()), "weights_nonnegative": bool(np.all(a >= 0)),
                         "ZZh_identity_max_error": algebra_error}}
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()