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

SEED = 2748
np.random.seed(SEED)

# Fold normal form from the paper: u_{t+1}=u+g*eps+b*u^2.
def fold_step(u, eps, g=1.0, b=1.0):
    return u + g * eps + b * u * u

def fixed_point(eps, g=1.0, b=1.0):
    return -math.sqrt(-g * eps / b)

def stable_gap(eps, g=1.0, b=1.0):
    u = fixed_point(eps, g, b)
    return 1.0 - abs(1.0 + 2.0 * b * u)

def passage_time(eps, g=1.0, b=1.0, max_steps=10000000):
    u, n = 0.0, 0
    while u < 1.0 and n < max_steps:
        u = fold_step(u, eps, g, b)
        n += 1
    return n

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

def jvp_scalar(v, u, eps, g=1.0, b=1.0):
    return (1.0 + 2.0 * b * u) * v

def predictor(u, delta=0.01, g=1.0, b=1.0):
    # Exact scalar JVP, equivalent to one power iteration.
    lam = abs(jvp_scalar(1.0, u, 0.0, g, b))
    return math.pi / max(delta, 1.0 - lam), lam

def residual(u, old):
    return abs(u - old) / (abs(u) + 1e-6)

def stable_run(eps, mode, tol=1e-5, rtol=0.03, tau_threshold=25.0,
               max_steps=50000):
    """Inference on the stable branch, stopping when |u-u*| <= tol.

    Fixed is a conservative cap, residual is the standard local early exit,
    and spectral additionally requires a small local relaxation forecast.
    """
    target = fixed_point(eps)
    u, old = 0.0, 0.0
    for t in range(1, max_steps + 1):
        old, u = u, fold_step(u, eps)
        r = residual(u, old)
        accurate = abs(u - target) <= tol
        if mode == 'fixed':
            if accurate: return t, True
        elif mode == 'residual':
            if accurate and r < rtol: return t, True
        else:
            tau, lam = predictor(u)
            # Stable branch has lam<1; reject critical/unstable estimates.
            if accurate and r < rtol and lam < 1.0 and tau < tau_threshold:
                return t, True
    return max_steps, False

def oscillatory_fallback():
    # A stable rotation has oscillating coordinates; a fold scalar detector
    # should disable its predictor on this non-fold trajectory.
    z = np.array([1.0, 0.0]); theta, rho = 0.55, 0.99
    R = rho * np.array([[math.cos(theta), -math.sin(theta)],
                        [math.sin(theta), math.cos(theta)]])
    xs = []
    for _ in range(20):
        z = R @ z; xs.append(float(z[0]))
    sign_changes = sum(xs[i] * xs[i-1] < 0 for i in range(1, len(xs)))
    return sign_changes >= 1, sign_changes

def main():
    # Prediction 1: gap ~ eps^(1/2); prediction 2: passage time ~ eps^(-1/2).
    eps = np.logspace(-5, -2, 12)
    gaps = np.array([stable_gap(-e) for e in eps])
    passages = np.array([passage_time(e) for e in eps], dtype=float)
    # Prediction 3: changing b gives opposite +/-1/2 prefactor exponents.
    e0 = 2e-5
    bs = np.array([0.25, 0.5, 1., 2., 4.])
    b_gaps = np.array([stable_gap(-e0, b=b) for b in bs])
    b_passages = np.array([passage_time(e0, b=b) for b in bs], dtype=float)
    products = b_gaps * b_passages

    # Direct stable-branch test of pi/(1-lambda): actual threshold time has
    # the same inverse-gap scaling, while pi is an asymptotic calibration.
    rows = []
    for e in [1e-4, 3e-5, 1e-5]:
        target = fixed_point(-e); u = 0.0
        for t in range(1, 50000):
            u = fold_step(u, -e)
            if abs(u-target) <= 1e-5:
                tau, lam = predictor(u)
                rows.append({'eps':e, 'actual_steps':t,
                             'predicted_tau':tau, 'lambda_hat':lam,
                             'gap':1-lam})
                break

    workload = np.logspace(-5, -3, 20)
    controller = {}
    for mode in ['fixed', 'residual', 'spectral']:
        vals = [stable_run(float(-e), mode)[0] for e in workload]
        controller[mode] = {'mean_steps':float(np.mean(vals)),
                            'median_steps':float(np.median(vals)), 'steps':vals}
    result = {
      'seed':SEED,
      'predictions':{
        'gap_epsilon_exponent':{'predicted':0.5,'observed':log_slope(eps,gaps)},
        'passage_epsilon_exponent':{'predicted':-0.5,'observed':log_slope(eps,passages)},
        'gap_b_exponent':{'predicted':0.5,'observed':log_slope(bs,b_gaps)},
        'passage_b_exponent':{'predicted':-0.5,'observed':log_slope(bs,b_passages)},
        'Pi':{'predicted':math.pi,'observed':products.tolist(),
              'mean':float(np.mean(products)),
              'relative_error':float(abs(np.mean(products)-math.pi)/math.pi)}},
      'stable_predictor_samples':rows, 'controller':controller,
      'oscillatory_fallback':dict(detected=oscillatory_fallback()[0],
                                  sign_changes=oscillatory_fallback()[1]),
      'settings':{'delta':0.01,'residual_tol':0.03,'accuracy_tol':1e-5,
                  'tau_threshold':25.0}}
    Path('results.json').write_text(json.dumps(result,indent=2))
    print(json.dumps(result,indent=2))

if __name__ == '__main__': main()