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

# Toy RG velocity: a translation-invariant field whose real-space influence
# decays as exp(-Lambda * distance). This is the minimal kernel implied by
# quasi-local RG flow and lets us measure truncation error exactly.
SEED = 2708
L = 256
LAMBDAS = np.array([0.10, 0.16, 0.25, 0.40, 0.64, 1.00])
EPSILONS = np.array([0.20, 0.10, 0.05, 0.02, 0.01])


def periodic_kernel(lam, n=L):
    d = np.minimum(np.arange(n), n - np.arange(n)).astype(float)
    k = np.exp(-lam * d)
    # Normalize so the zero mode has unit gain; normalization cancels in the
    # relative truncation error but makes the operator easy to interpret.
    return k / k.sum()


def truncate(k, radius):
    n = len(k)
    d = np.minimum(np.arange(n), n - np.arange(n))
    out = np.where(d <= radius, k, 0.0)
    # Preserve the zero-mode response, as a local approximation would normally
    # absorb the omitted mass into its learned bias/gain.
    return out / out.sum()


def relative_operator_error(k, radius):
    """RMS error over all Fourier modes, normalized by exact operator power."""
    exact = np.fft.rfft(k)
    approx = np.fft.rfft(truncate(k, radius))
    return float(np.linalg.norm(approx - exact) / np.linalg.norm(exact))


def required_radius(lam, eps):
    k = periodic_kernel(lam)
    for r in range(L // 2):
        if relative_operator_error(k, r) <= eps:
            return r
    return L // 2


def verify_field_error(lam, radius, trials=64):
    rng = np.random.default_rng(SEED + int(1000 * lam) + radius)
    k = periodic_kernel(lam)
    kt = truncate(k, radius)
    errs = []
    for _ in range(trials):
        x = rng.normal(size=L)
        y = np.fft.irfft(np.fft.rfft(k) * np.fft.rfft(x), n=L)
        yt = np.fft.irfft(np.fft.rfft(kt) * np.fft.rfft(x), n=L)
        errs.append(np.linalg.norm(y - yt) / np.linalg.norm(y))
    return float(np.mean(errs))


def linear_fit(x, y):
    a, b = np.polyfit(x, y, 1)
    pred = a * x + b
    r2 = 1.0 - np.sum((y - pred) ** 2) / max(np.sum((y - y.mean()) ** 2), 1e-12)
    return float(a), float(b), float(r2)


def pyramid_radius_check():
    # Lambda_s = pi/a_s. The claimed cell-radius scaling is
    # R_s/a_s = c [log L + log(1/eps)] / pi, independent of s.
    eps = 0.02
    c = 1.0
    rows = []
    for s in range(6):
        a = 2 ** s
        lam = math.pi / a
        R = c * (math.log(L) + math.log(1 / eps)) / lam
        r = math.ceil(R / a)
        rows.append({'level': s, 'spacing': a, 'Lambda': lam,
                     'physical_R': R, 'cell_radius': r})
    return rows


def main():
    table = []
    for lam in LAMBDAS:
        for eps in EPSILONS:
            r = required_radius(lam, eps)
            table.append({'Lambda': float(lam), 'epsilon': float(eps),
                          'required_R': r,
                          'measured_error': relative_operator_error(periodic_kernel(lam), r),
                          'field_error': verify_field_error(lam, r)})

    # Prediction 1: at fixed epsilon, R*Lambda is approximately constant.
    fixed = [q for q in table if q['epsilon'] == 0.02]
    inv_lam = np.array([1 / q['Lambda'] for q in fixed])
    radii = np.array([q['required_R'] for q in fixed])
    slope, intercept, r2 = linear_fit(inv_lam, radii)
    products = [q['Lambda'] * q['required_R'] for q in fixed]

    # Prediction 2: at fixed Lambda, R grows linearly with log(1/epsilon).
    lam0 = 0.25
    tolrows = [q for q in table if q['Lambda'] == lam0]
    xlog = np.array([math.log(1 / q['epsilon']) for q in tolrows])
    yR = np.array([q['required_R'] for q in tolrows])
    slope_log, intercept_log, r2_log = linear_fit(xlog, yR)

    # Prediction 3 / baseline: a fixed radius has increasingly bad error as the
    # RG length grows (Lambda decreases), while radius selected by the bound
    # stays below tolerance.
    fixed_radius = 8
    baseline = []
    for lam in LAMBDAS:
        err = relative_operator_error(periodic_kernel(lam), fixed_radius)
        rg_r = required_radius(lam, 0.02)
        rg_err = relative_operator_error(periodic_kernel(lam), rg_r)
        baseline.append({'Lambda': float(lam), 'fixed_radius': fixed_radius,
                         'fixed_radius_error': err, 'rg_radius': rg_r,
                         'rg_error': rg_err})

    result = {
        'setup': {'L': L, 'kernel': 'normalized exp(-Lambda * periodic_distance)',
                  'seed': SEED},
        'radius_sweep': table,
        'predictions': {
            'inverse_cutoff': {
                'statement': 'R is affine in 1/Lambda at fixed epsilon',
                'fit_R_vs_1_over_Lambda': {'slope': slope, 'intercept': intercept, 'R2': r2},
                'Lambda_times_R': products,
                'mean': float(np.mean(products)), 'cv': float(np.std(products) / np.mean(products))},
            'log_tolerance': {
                'statement': 'R is affine in log(1/epsilon) at fixed Lambda',
                'Lambda': lam0,
                'fit_R_vs_log_inverse_epsilon': {'slope': slope_log, 'intercept': intercept_log, 'R2': r2_log}},
            'rescaled_grid': {
                'statement': 'R/a is approximately constant when Lambda=pi/a',
                'levels': pyramid_radius_check()}
        },
        'baseline_comparison': baseline
    }
    Path('results.json').write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))


if __name__ == '__main__':
    main()