import json
import math
import numpy as np


def skew_balance_check(seed=0, n=10000):
    rng = np.random.default_rng(seed)
    errors = []
    for _ in range(n):
        U = rng.normal(0, 5)
        du = rng.normal(0, 4)
        p = abs(rng.normal()) + 1e-12
        qf = p * math.exp(-du / 2.0)
        qr = p * math.exp(du / 2.0)
        errors.append(abs((math.log(qf) - U) -
                          (math.log(qr) - (U + du))))
    return float(max(errors)), float(np.mean(errors))


def simulate(rho, events=160000, burn=5000, seed=1, nring=128):
    rng = np.random.default_rng(seed)
    p = rng.normal()
    x = 0
    signs, waits, positions = [], [], []
    for k in range(events + burn):
        rate = abs(p)
        dt = rng.exponential(1.0 / rate) if rate > 1e-14 else 0.0
        direction = 1 if p >= 0 else -1
        x = (x + direction) % nring
        p = rho * p + math.sqrt(max(0.0, 1.0 - rho * rho)) * rng.normal()
        if k >= burn:
            signs.append(direction)
            waits.append(dt)
            positions.append(x)
    signs, waits, positions = map(np.asarray, (signs, waits, positions))
    corr = float(np.mean(signs[:-1] * signs[1:]))
    run_length = float(1.0 / (1.0 - corr))
    centered = signs - signs.mean()
    var = np.mean(centered * centered)
    ac = [float(np.mean(centered[:-lag] * centered[lag:]) / var)
          for lag in range(1, 100)]
    sign_tau = float(1.0 + 2.0 * sum(a for a in ac if a > 0))
    angle = 2 * math.pi * positions / nring
    obs = np.cos(angle)
    c = obs - obs.mean()
    v = np.mean(c*c)
    acx = [float(np.mean(c[:-lag] * c[lag:]) / v)
           for lag in range(1, 100)]
    x_tau = float(1.0 + 2.0 * sum(a for a in acx if a > 0))
    return {
        "rho": rho, "sign_corr": corr, "run_length": run_length,
        "sign_tau": sign_tau, "position_tau": x_tau,
        "mean_wait": float(waits.mean()),
        "event_rate_per_time": float(len(waits) / waits.sum())
    }


def main():
    max_err, mean_err = skew_balance_check()
    # For Gaussian OU-refresh momenta, corr(sign_t, sign_{t+1})
    # = 2 asin(rho)/pi. Hence mean consecutive same-sign run length is
    # 1/(1-corr). These are direct mechanism predictions.
    rows = []
    for rho in [0.0, 0.25, 0.5, 0.75, 0.9]:
        out = simulate(rho, seed=100 + int(100 * rho))
        out["pred_corr"] = 2.0 / math.pi * math.asin(rho)
        out["pred_run_length"] = 1.0 / (1.0 - out["pred_corr"])
        out["corr_abs_error"] = abs(out["sign_corr"] - out["pred_corr"])
        rows.append(out)
    result = {
        "skew_balance_max_abs_log_error": max_err,
        "skew_balance_mean_abs_log_error": mean_err,
        "prediction_rows": rows,
        "baseline_rho0_position_tau": rows[0]["position_tau"],
        "idea_rho075_position_tau": rows[3]["position_tau"],
        "note": "Event-rate is reported empirically; per-event Gaussian refresh has heavy-tailed waiting times, so E|p| is not the event rate."
    }
    with open("results.json", "w") as f:
        json.dump(result, f, indent=2)
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()