import json, math
import numpy as np
from sphere_exp_residual import sphere_exp
np.set_printoptions(precision=8, suppress=True)

def exp_np(x,v):
    r=np.linalg.norm(v)
    return math.cos(r)*x + (math.sin(r)/r if r>1e-14 else 1.)*v

def bfield(x):
    c=np.array([.31,-.47,.22])
    return c-x*np.dot(x,c)

def intrinsic(x,h,n):
    for _ in range(n): x=exp_np(x,h*bfield(x))
    return x

def chart(x): return x[:2]/(1-x[2])
def invchart(q):
    r=np.dot(q,q)
    return np.array([2*q[0]/(1+r),2*q[1]/(1+r),(r-1)/(1+r)])
def jacobian_inv(q):
    J=np.empty((3,2)); eps=1e-6
    for j in range(2):
        e=np.zeros(2); e[j]=eps
        J[:,j]=(invchart(q+e)-invchart(q-e))/(2*eps)
    return J
def transformed(q): return np.array([q[0]+.55*q[0]**2,q[1]+.30*q[0]*q[1]])
def inv_transformed(p):
    q=p.copy()
    for _ in range(12):
        f=transformed(q)-p
        J=np.array([[1+1.1*q[0],0],[.30*q[1],1+.30*q[0]]])
        q-=np.linalg.solve(J,f)
    return q
def jac_transformed(q): return np.array([[1+1.1*q[0],0],[.30*q[1],1+.30*q[0]]])
def coordinate_euler(x0,h,n,which):
    q=chart(x0) if which=='q' else transformed(chart(x0))
    for _ in range(n):
        qq=q if which=='q' else inv_transformed(q)
        x=invchart(qq); dq=np.linalg.lstsq(jacobian_inv(qq),bfield(x),rcond=None)[0]
        q=q+h*dq if which=='q' else q+h*(jac_transformed(qq)@dq)
    return invchart(q if which=='q' else inv_transformed(q))
def slope(hs,errs): return float(np.polyfit(np.log(hs),np.log(np.maximum(errs,1e-18)),1)[0])

def parallel_transport_from_start(x0,v0,theta):
    r=np.linalg.norm(v0); u=v0/r
    return r*(-math.sin(theta)*x0+math.cos(theta)*u)

def main():
    x0=np.array([.35,-.25,.9027735]); x0/=np.linalg.norm(x0); T=.8
    hs=np.array([.1,.05,.025,.0125])
    ref=intrinsic(x0,1e-5,int(T/1e-5)); exp_err=[]; euler_err=[]; norm_exp=[]; norm_euler=[]; chart_gap=[]
    for h in hs:
        n=round(T/h); xe=intrinsic(x0,h,n); xu=x0.copy()
        for _ in range(n): xu=xu+h*bfield(xu)
        q1=coordinate_euler(x0,h,n,'q'); q2=coordinate_euler(x0,h,n,'p')
        exp_err.append(np.linalg.norm(xe-ref)); euler_err.append(np.linalg.norm(xu-ref))
        norm_exp.append(abs(np.linalg.norm(xe)-1)); norm_euler.append(abs(np.linalg.norm(xu)-1)); chart_gap.append(np.linalg.norm(q1-q2))
    # Correct geodesic composition test: velocity is parallel transported at every point.
    v=np.array([.2,.3,-.1]); v-=x0*np.dot(x0,v); r=np.linalg.norm(v); exact=exp_np(x0,v); geo_err=[]
    for h in hs:
        z=x0.copy()
        for k in range(round(1/h)):
            vk=parallel_transport_from_start(x0,v,k*h*r)
            z=exp_np(z,h*vk)
        geo_err.append(np.linalg.norm(z-exact))
    # Tangent isotropic noise: empirical covariance in an orthonormal tangent frame.
    rng=np.random.default_rng(7); e1=np.array([1.,0,0]); e1-=x0*np.dot(x0,e1); e1/=np.linalg.norm(e1); e2=np.cross(x0,e1)
    Z=rng.normal(size=(200000,2)); cov=np.cov(Z,rowvar=False)
    tangent_cov=np.cov((Z[:,0,None]*e1+Z[:,1,None]*e2).T)
    out={'step_sizes':hs.tolist(),'predictions':{
      'P1_norm_preservation':{'predicted':'Exp error roundoff; additive error nonzero','exp_abs_errors':norm_exp,'additive_abs_errors':norm_euler},
      'P2_first_order_global_error':{'predicted':'global endpoint error O(h), slope near 1','exp_errors':exp_err,'additive_errors':euler_err,'exp_slope':slope(hs,exp_err),'additive_slope':slope(hs,euler_err)},
      'P3_chart_effect':{'predicted':'coordinate Euler discrepancy O(h), slope near 1','chart_gaps':chart_gap,'slope':slope(hs,chart_gap)},
      'P4_geodesic_exactness':{'predicted':'Exp with parallel-transported velocity is step independent','errors':geo_err},
      'P5_tangent_noise':{'predicted':'orthonormal tangent coordinates have covariance I','coordinate_covariance':cov.tolist(),'ambient_tangent_covariance_error':float(np.linalg.norm(tangent_cov-(np.eye(3)-np.outer(x0,x0))/1.0))}
    },'summary':{'max_exp_norm_error':max(norm_exp),'coarsest_additive_norm_error':norm_euler[0],'coarsest_exp_error':exp_err[0],'coarsest_additive_error':euler_err[0],'max_geodesic_error':max(geo_err)}}
    print(json.dumps(out,indent=2))
if __name__=='__main__': main()