import json, math
from pathlib import Path
import numpy as np
from scipy.optimize import minimize

# Two stable limit cycles are represented by stable radii r=1 and r=2,
# with a common angular velocity. The separatrix is r=1.5.
R1, RS, R2 = 1.0, 1.5, 2.0
OMEGA = 1.0
DT = 0.01

def fr(r):
    return -(r - R1) * (r - RS) * (r - R2)

def dfr(r):
    return -(3*r*r - 2*(R1+RS+R2)*r + (R1*RS + R1*R2 + RS*R2))

def action_integral(a, b, n=100000):
    x = np.linspace(a, b, n)
    return float(np.trapz(-fr(x), x))

def discrete_action(x, T):
    # B=1, so this is exactly the stated 1/4 integral discretization.
    h = T / (len(x)-1)
    v = np.diff(x) / h
    mid = 0.5*(x[:-1] + x[1:])
    return float(h * np.sum((v - fr(mid))**2) / 4.0)

def optimize_action(a, b, T=20.0, n=81):
    # Fixed endpoints and a smooth monotone initial path. The optimizer can
    # discover the minimum-action uphill path for the finite endpoint problem.
    x0 = np.linspace(a, b, n)
    h = T/(n-1)
    def fun(y):
        x = np.r_[a, y, b]
        return discrete_action(x, T)
    ans = minimize(fun, x0[1:-1], method='L-BFGS-B', options={'maxiter': 1200, 'ftol': 1e-12})
    return float(ans.fun), bool(ans.success)

def simulate_rates(D, paths=160, steps=125000, seed=0):
    # Vectorized first-passage counting. Hysteretic labels avoid recrossing noise.
    rng = np.random.default_rng(seed)
    r = np.full(paths, R1)
    state = np.ones(paths, dtype=np.int8)
    counts = np.zeros(2, dtype=np.int64)
    for _ in range(steps):
        r += fr(r)*DT + math.sqrt(2*D*DT)*rng.standard_normal(paths)
        # radial coordinate is reflected at zero (irrelevant for rare events)
        r = np.abs(r)
        hit2 = (state == 1) & (r >= RS)
        hit1 = (state == 2) & (r <= RS)
        counts[0] += np.count_nonzero(hit2)
        counts[1] += np.count_nonzero(hit1)
        state[hit2] = 2
        state[hit1] = 1
    total_time = paths * steps * DT
    return counts / total_time, counts

def main():
    # Prediction 1: action barrier from cycle 1 to separatrix.
    barrier = action_integral(R1, RS)
    reverse_barrier = action_integral(R2, RS)
    # Prediction 2: transverse Floquet multiplier exp(f'(cycle)*period).
    period = 2*math.pi/OMEGA
    floquet1 = math.exp(dfr(R1)*period)
    floquet2 = math.exp(dfr(R2)*period)
    opt_action, opt_ok = optimize_action(1.02, 1.48)
    finite_barrier = action_integral(1.02, 1.48)

    # Prediction 3: log k is affine in 1/D in the rare-event regime, while
    # sufficiently large D leaves that regime. Five-plus values are required.
    Ds = np.array([0.004, 0.006, 0.010, 0.015, 0.030, 0.080, 0.200])
    rates = []
    counts = []
    for i, D in enumerate(Ds):
        rr, cc = simulate_rates(float(D), paths=48, steps=50000, seed=100+i)
        rates.append(rr)
        counts.append(cc.tolist())
    rates = np.asarray(rates)
    counts = np.asarray(counts)
    # Fit only forward rates with enough events; report rare subset and all-data fits.
    valid = (counts[:,0] >= 8) & (rates[:,0] > 0)
    rare = valid & (Ds <= 0.015)
    high = valid & (Ds >= 0.080)
    def fit(mask):
        if np.count_nonzero(mask) < 2: return [float('nan')]*3
        p = np.polyfit(1/Ds[mask], np.log(rates[mask,0]), 1)
        pred = np.polyval(p, 1/Ds[mask])
        r2 = 1 - np.sum((np.log(rates[mask,0])-pred)**2)/np.sum((np.log(rates[mask,0])-np.mean(np.log(rates[mask,0])))**2)
        return [float(p[0]), float(p[1]), float(r2)]
    fit_rare = fit(rare)
    fit_all = fit(valid)
    fit_high = fit(high)
    crossover_D = (RS-R1)**2/(2*period)
    result = {
      'model': 'dr=-(r-1)(r-1.5)(r-2)dt + sqrt(2D)dW; dtheta=dt',
      'predictions': {
        'action_barrier_1_to_separatrix': {'predicted': barrier, 'estimated_discrete_path_1.02_to_1.48': opt_action, 'finite_endpoint_integral': finite_barrier, 'relative_error_vs_finite_integral': abs(opt_action-finite_barrier)/finite_barrier},
        'floquet_transverse': {'period': period, 'cycle_1_derivative': dfr(R1), 'cycle_2_derivative': dfr(R2), 'cycle_1_multiplier': floquet1, 'cycle_2_multiplier': floquet2, 'neutral_phase_multiplier': 1.0},
        'arrhenius_slope': {'predicted_slope': -barrier, 'rare_fit_slope': fit_rare[0], 'rare_fit_intercept': fit_rare[1], 'rare_fit_R2': fit_rare[2], 'all_fit_slope': fit_all[0], 'all_fit_R2': fit_all[2], 'high_noise_fit_slope': fit_high[0], 'high_noise_fit_R2': fit_high[2], 'predicted_crossover_D_from_cycle_width': crossover_D}
      },
      'sweep': [{'D':float(D), 'forward_rate':float(rates[i,0]), 'reverse_rate':float(rates[i,1]), 'forward_count':int(counts[i,0]), 'reverse_count':int(counts[i,1]), 'x=1/D':float(1/D)} for i,D in enumerate(Ds)],
      'settings': {'dt':DT, 'paths':48, 'steps':50000, 'total_time_each_D':48*50000*DT, 'seed_base':100},
      'mechanism_checks': {'action_agreement_pass': abs(opt_action-finite_barrier)/finite_barrier < .20, 'floquet_stable_pass': floquet1 < 1 and floquet2 < 1, 'rare_arrhenius_linear_pass': fit_rare[2] > .90, 'large_noise_deviation_from_rare_slope': abs(fit_high[0]-fit_rare[0]) > .20*abs(fit_rare[0])}
    }
    Path('results.json').write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))

if __name__ == '__main__':
    main()