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

SEED = 2820
A, B, S1, S2 = 1.0, 0.8, 0.5, 2.0
# Under K(u)=integral exp(i omega u) Khat(omega)d omega,
# a exp(-s^2 u^2/2) has nonnegative spectral mass a and omega~N(0,s^2).
C = A + B
NEG_MASS = B / C


def exact_kernel(u, bb=B):
    u = np.asarray(u)
    return A * np.exp(-0.5 * (S1 * u) ** 2) - bb * np.exp(-0.5 * (S2 * u) ** 2)


def draw_freqs(m, rng, bb=B):
    c = A + bb
    neg = rng.random(m) < (bb / c) if c else np.zeros(m, dtype=bool)
    omega = np.where(neg, rng.normal(0.0, S2, m), rng.normal(0.0, S1, m))
    signs = np.where(neg, -1.0, 1.0)
    return omega, signs, c


def features(x, omega):
    c, s = np.cos(np.outer(x, omega)), np.sin(np.outer(x, omega))
    return np.stack((c, s), axis=2).reshape(len(x), -1)


def rff_kernel(x, y, m, signed, seed):
    omega, signs, c = draw_freqs(m, np.random.default_rng(seed))
    zx, zy = features(x, omega), features(y, omega)
    d = np.repeat(signs if signed else np.ones(m), 2)
    return (c / m) * ((zx * d) @ zy.T)


def factor_aggregate(x, v, m, signed, seed):
    omega, signs, c = draw_freqs(m, np.random.default_rng(seed))
    z = features(x, omega)
    d = np.repeat(signs if signed else np.ones(m), 2)
    return (c / m) * ((z * d) @ (z.T @ v))


def slope(ms, errors):
    return float(np.polyfit(np.log(np.asarray(ms)), np.log(np.asarray(errors)), 1)[0])


def main():
    rng = np.random.default_rng(SEED)
    n_pairs, trials = 180, 100
    u = rng.uniform(-3.0, 3.0, n_pairs)
    truth = exact_kernel(u)
    ms = [16, 32, 64, 128, 256, 512]
    signed_rmse, positive_rmse, signed_bias = [], [], []
    for m in ms:
        se, pe, estimates = [], [], []
        for t in range(trials):
            omega, signs, c = draw_freqs(m, np.random.default_rng(SEED + 10000*m + t))
            vals = np.cos(np.outer(u, omega))
            est = (c / m) * (vals * signs).sum(axis=1)
            pos = (c / m) * vals.sum(axis=1)
            se.append(np.mean((est - truth) ** 2)); pe.append(np.mean((pos - truth) ** 2))
            estimates.append(est)
        signed_rmse.append(float(np.sqrt(np.mean(se))))
        positive_rmse.append(float(np.sqrt(np.mean(pe))))
        signed_bias.append(float(np.mean(np.concatenate(estimates) - np.tile(truth, trials))))
    observed_slope = slope(ms, signed_rmse)
    positive_limit = float(np.sqrt(np.mean((A*np.exp(-0.5*(S1*u)**2) + B*np.exp(-0.5*(S2*u)**2) - truth)**2)))

    mass_rows = []
    for bb in [0.0, 0.2, 0.5, 0.8, 1.2]:
        signed_truth = exact_kernel(u, bb)
        positive_kernel = A*np.exp(-0.5*(S1*u)**2) + bb*np.exp(-0.5*(S2*u)**2)
        discrepancy = float(np.sqrt(np.mean((positive_kernel - signed_truth)**2)))
        predicted = float(2.0 * bb * np.sqrt(np.mean(np.exp(-(S2*u)**2))))
        mass_rows.append({'negative_amplitude': bb, 'negative_mass_fraction': bb/(A+bb),
                          'observed_positive_limit_rmse': discrepancy, 'predicted_linear_rmse': predicted})

    n, r, m = 220, 8, 128
    x, v = rng.uniform(-3, 3, n), rng.normal(size=(n, r))
    exact = exact_kernel(x[:, None] - x[None, :]) @ v
    errs, times = [], {'exact': [], 'signed': []}
    for t in range(20):
        st = time.perf_counter(); out = factor_aggregate(x, v, m, True, SEED+50000+t); times['signed'].append(time.perf_counter()-st)
        errs.append(float(np.linalg.norm(out-exact)/np.linalg.norm(exact)))
        st = time.perf_counter(); _ = exact_kernel(x[:, None]-x[None, :]) @ v; times['exact'].append(time.perf_counter()-st)
    result = {'kernel': 'exp(-0.5*(0.5u)^2) - 0.8 exp(-0.5*(2u)^2)', 'C': C, 'negative_mass_fraction': NEG_MASS,
      'prediction_1_M_minus_half': {'M': ms, 'signed_rmse': signed_rmse, 'slope': observed_slope, 'predicted_slope': -0.5, 'signed_bias_at_M512': signed_bias[-1]},
      'prediction_2_positive_features': {'positive_rmse': positive_rmse, 'predicted_asymptotic_rmse': positive_limit},
      'prediction_3_negative_mass_sweep': mass_rows,
      'aggregation': {'n': n, 'r': r, 'M': m, 'relative_errors': errs, 'median_exact_seconds': float(np.median(times['exact'])), 'median_signed_factor_seconds': float(np.median(times['signed']))}}
    Path('results.json').write_text(json.dumps(result, indent=2)); print(json.dumps(result, indent=2))

if __name__ == '__main__': main()