import json
import math
import numpy as np
from scipy.stats import skew, kurtosis

SEED = 1499
N = 200000
EPS = np.array([0.01, 0.02, 0.04, 0.08, 0.16, 0.32])
ALPHAS = [0.5, 1.0, 1.5]

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

def sigma2(e, a, c=1.): return c * e**(2-a)/(2-a)
def rho3(e, a, c=1.): return c * e**(3-a)/(3-a)
def jump_rate(e, a, c=1.): return c * e**(-a)/a

def main():
    out = {"N": N, "eps": EPS.tolist(), "results": {}}
    for a in ALPHAS:
        rng = np.random.default_rng(SEED + round(100*a))
        rows = []
        for e in EPS:
            s2, r3 = sigma2(e,a), rho3(e,a)
            # Directly sample the proposed Gaussian replacement. This is a
            # cheap implementation of the small-jump line in a neural SDE.
            g = rng.normal(0., math.sqrt(s2), N)
            z = g / math.sqrt(s2)
            rows.append({
              "epsilon": float(e),
              "sigma2_theory": s2,
              "gaussian_variance_mc": float(np.var(g)),
              "variance_relative_error": float(abs(np.var(g)-s2)/s2),
              # These are the paper's quantitative rate proxies:
              # naive residual scale sigma; compensated Berry-Esseen/W1 scale
              # rho/sigma^2; standardized residual normality scale rho/sigma^3.
              "naive_scale_sigma": math.sqrt(s2),
              "compensated_w1_proxy_rho_over_sigma2": r3/s2,
              "normality_proxy_rho_over_sigma3": r3/(s2**1.5),
              "normal_sample_abs_skew": float(abs(skew(z))),
              "normal_sample_excess_kurtosis": float(kurtosis(z)),
              "large_jump_rate": jump_rate(e,a),
            })
        e = EPS
        def vals(k): return np.array([r[k] for r in rows])
        obs = {
          "variance_exponent": slope(e, vals("sigma2_theory")),
          "naive_error_exponent": slope(e, vals("naive_scale_sigma")),
          "compensated_W1_proxy_exponent": slope(e, vals("compensated_w1_proxy_rho_over_sigma2")),
          "residual_normality_exponent": slope(e, vals("normality_proxy_rho_over_sigma3")),
          "large_jump_rate_exponent": slope(e, vals("large_jump_rate")),
          "mean_gaussian_relative_variance_error": float(np.mean(vals("variance_relative_error"))),
        }
        pred = {
          "variance_exponent": 2-a,
          "naive_error_exponent": 1-a/2,
          "compensated_W1_proxy_exponent": 1.0,
          "residual_normality_exponent": a/2,
          "large_jump_rate_exponent": -a,
        }
        out["results"][str(a)] = {"observed": obs, "predicted": pred, "rows": rows}
    with open("results.json", "w") as f: json.dump(out, f, indent=2)
    for a,d in out["results"].items():
        print(f"alpha={a}")
        for k in d["predicted"]:
            print(f"  {k}: observed={d['observed'][k]:.5f} predicted={d['predicted'][k]:.5f}")
        print(f"  MC relative variance error={d['observed']['mean_gaussian_relative_variance_error']:.5f}")
        for r in d["rows"]:
            print(f"    eps={r['epsilon']:.3g} sigma2={r['sigma2_theory']:.5g} "
                  f"naive={r['naive_scale_sigma']:.5g} comp_proxy={r['compensated_w1_proxy_rho_over_sigma2']:.5g} "
                  f"normality={r['normality_proxy_rho_over_sigma3']:.5g} rate={r['large_jump_rate']:.5g}")

if __name__ == "__main__": main()