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

@dataclass
class SimResult:
    rms: float
    tau: float
    values: np.ndarray

def grad(x, m):
    return np.sign(x) * abs(x) ** (m - 1)

def simulate(m, alpha, sigma=1.0, n=120000, burn=30000, seed=0):
    rng = np.random.default_rng(seed)
    x = 0.0
    out = np.empty(n - burn)
    for t in range(n):
        x -= alpha * (grad(x, m) + sigma * rng.normal())
        if not np.isfinite(x) or abs(x) > 1e8:
            return SimResult(float('nan'), float('nan'), np.empty(0))
        if t >= burn:
            out[t - burn] = x
    z = out - out.mean()
    var = np.mean(z*z)
    tau = 1.0
    if var > 0:
        for lag in range(1, min(4000, len(z)//10)):
            ac = np.mean(z[:-lag] * z[lag:]) / var
            if ac <= 0:
                break
            tau += 2*ac
    return SimResult(float(np.sqrt(np.mean(out*out))), float(tau), out)

def slope(x, y):
    return float(np.polyfit(np.log(x), np.log(y), 1)[0])

def calibrate(target_r, m, sigma=1.0, C=1.0, amin=1e-5, amax=0.25):
    # stated law in 1D: r = C * alpha^(1/m) * sigma
    return float(np.clip((target_r/(C*sigma))**m, amin, amax))

def main():
    # Small alphas keep Euler updates stable and make asymptotic laws visible.
    ms = [2, 3, 4, 5]
    alphas = np.array([0.002, 0.004, 0.008, 0.016])
    rows = []
    for m in ms:
        radii, taus = [], []
        for j, a in enumerate(alphas):
            r = simulate(m, a, seed=1000 + 10*m + j)
            radii.append(r.rms); taus.append(r.tau)
            rows.append({'m':m, 'alpha':float(a), 'radius':r.rms, 'tau':r.tau})
        rows.append({'m':m, 'radius_slope':slope(alphas, radii),
                     'tau_slope':slope(alphas, taus),
                     'pred_radius_slope':1/m, 'pred_tau_slope':-(m-1)})

    # Empirically calibrate the unknown C at a reference rate, then target a radius.
    calibrated = []
    for m in ms:
        ref_alpha = 0.004
        ref = simulate(m, ref_alpha, seed=4000+m, n=180000, burn=50000)
        C_hat = ref.rms / (ref_alpha**(1/m))
        a = calibrate(0.12, m, C=C_hat)
        got = simulate(m, a, seed=4500+m, n=180000, burn=50000)
        calibrated.append({'m':m, 'C_hat':C_hat, 'alpha':a, 'target_radius':0.12,
                           'observed_radius':got.rms,
                           'relative_error':abs(got.rms-0.12)/0.12})

    # Noise prediction: radius is proportional to sigma^(2/m), since v=Sigma=sigma^2.
    noise_sweep = []
    for m in [2, 3, 4, 5]:
        a = 0.004
        sigmas = np.array([0.5, 1.0, 2.0])
        radii = []
        for j, sig in enumerate(sigmas):
            radii.append(simulate(m, a, sigma=float(sig), seed=6000+10*m+j,
                                  n=160000, burn=45000).rms)
        noise_sweep.append({'m':m, 'observed_slope':slope(sigmas, radii),
                            'predicted_slope':2/m, 'sigmas':sigmas.tolist(),
                            'radii':radii})

    # Adaptation test: infer m from curvature q(rho) for H=|x|^m/m,
    # where q(rho) proportional to rho^(m-2), then invert the radius law.
    adapt = []
    target = 0.12
    for m in ms:
        rho = 0.03
        q1 = rho**(m-2)
        q2 = (2*rho)**(m-2)
        mhat = 2 + math.log((q2+1e-12)/(q1+1e-12), 2) if m != 2 else 2.0
        a = calibrate(target, mhat)
        rr = simulate(m, a, seed=5000+m, n=140000, burn=40000)
        adapt.append({'true_m':m, 'mhat':mhat, 'alpha':a, 'target_radius':target,
                      'observed_radius':rr.rms, 'relative_radius_error':abs(rr.rms-target)/target})

    # Same target-radius request: calibrated rate versus quadratic miscalibration.
    compare = []
    for m in [3, 4, 5]:
        ai = calibrate(target, m)
        aq = target**2
        ri = simulate(m, ai, seed=8000+m, n=140000, burn=40000).rms
        rq = simulate(m, aq, seed=9000+m, n=140000, burn=40000).rms
        compare.append({'m':m, 'idea_alpha':ai, 'idea_radius':ri,
                        'quadratic_alpha':aq, 'quadratic_radius':rq,
                        'idea_error':abs(ri-target), 'quadratic_error':abs(rq-target)})

    result = {'radius_and_mixing_sweep':rows, 'adaptation':adapt,
              'empirical_C_calibration':calibrated, 'noise_sweep':noise_sweep,
              'target_radius_comparison':compare,
              'notes':'RMS stationary radius and integrated autocorrelation time from 1D constant-step SGD.'}
    Path('results.json').write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))

if __name__ == '__main__':
    main()