import json
import numpy as np
from scipy.optimize import root
from experiment import E, I, Iinv, loglog_slope

z0 = np.r_[np.array([0.2, -0.1, 0.3]), np.array([0.7, 1.1, 1.5])]

def grad_h(z):
    q, p = z[:3], z[3:]
    return np.r_[0.12*np.sin(q), Iinv*p]

def field(z):
    g = grad_h(z)
    coad = np.einsum('a,aij,j->i', z[3:], E, g[3:])
    return np.r_[g[3:], -g[:3] + coad]

def step(z, dt):
    def r(x):
        m = (z+x)/2
        return x-z-dt*field(m)
    sol = root(r, z+dt*field(z), method='hybr')
    assert sol.success and np.max(np.abs(r(sol.x))) < 1e-9
    return sol.x

def rollout(dt, T=5.0):
    z=z0.copy()
    for _ in range(round(T/dt)): z=step(z,dt)
    return z

# Fine reference and successively refined midpoint trajectories.
ref=rollout(0.0005)
dts=np.array([0.08,0.04,0.02,0.01])
errs=np.array([np.linalg.norm(rollout(float(dt))-ref) for dt in dts])
slope=loglog_slope(dts, errs)
out={'dts':dts.tolist(),'errors_vs_fine_reference':errs.tolist(),
     'predicted_global_error_exponent':2.0,'observed_global_error_exponent':float(slope)}
open('nonlinear_results.json','w').write(json.dumps(out,indent=2))
print(json.dumps(out,indent=2))