# Moral's Mark - TOE E.4 "Mercury's Perihelion".
# The Sun and eight planets as summed dents. Prints each planet's share of Mercury's perihelion turning,
# the total, beta for the remaining 43 arcsec/century, and the extra turning that one beta gives
# Mercury, Venus, the Earth and Mars.  Needs numpy and scipy. Takes several minutes.
import numpy as np, time
from scipy.integrate import solve_ivp
GMs=1.32712440018e20
# GM of planets (Earth includes Moon)
GM=np.array([2.2032e13,3.24859e14,4.0350323e14,4.282837e13,1.26686534e17,3.7931187e16,5.793939e15,6.836529e15])
names=['Mercury','Venus','Earth+Moon','Mars','Jupiter','Saturn','Uranus','Neptune']
AU=1.495978707e11; DAY=86400.0; YR=365.25*DAY
# J2000 mean elements (Standish): a[AU], e, I, L, long.peri, long.node  (deg)
EL=np.array([
[0.38709927,0.20563593,7.00497902,252.25032350,77.45779628,48.33076593],
[0.72333566,0.00677672,3.39467605,181.97909950,131.60246718,76.67984255],
[1.00000261,0.01671123,-0.00001531,100.46457166,102.93768193,0.0],
[1.52371034,0.09339410,1.84969142,-4.55343205,-23.94362959,49.55953891],
[5.20288700,0.04838624,1.30439695,34.39644051,14.72847983,100.47390909],
[9.53667594,0.05386179,2.48599187,49.95424423,92.59887831,113.66242448],
[19.18916464,0.04725744,0.77263783,313.23810451,170.95427630,74.01692503],
[30.06992276,0.00859048,1.77004347,-55.12002969,44.96476227,131.78422574]])
def kepler_state(a,e,I,L,w_,O,mu):
    I,L,w_,O=np.radians([I,L,w_,O]); w=w_-O; M=np.mod(L-w_,2*np.pi)
    E=M
    for _ in range(50): E=E-(E-e*np.sin(E)-M)/(1-e*np.cos(E))
    xp=a*(np.cos(E)-e); yp=a*np.sqrt(1-e*e)*np.sin(E)
    n=np.sqrt(mu/a**3); r=a*(1-e*np.cos(E))
    vxp=-a*n*np.sin(E)/(1-e*np.cos(E)); vyp=a*n*np.sqrt(1-e*e)*np.cos(E)/(1-e*np.cos(E))
    cO,sO,cw,sw,cI,sI=np.cos(O),np.sin(O),np.cos(w),np.sin(w),np.cos(I),np.sin(I)
    R=np.array([[cO*cw-sO*sw*cI,-cO*sw-sO*cw*cI],[sO*cw+cO*sw*cI,-sO*sw+cO*cw*cI],[sw*sI,cw*sI]])
    return R@np.array([xp,yp]), R@np.array([vxp,vyp])
def initial():
    pos=[np.zeros(3)]; vel=[np.zeros(3)]
    for k in range(8):
        r,v=kepler_state(EL[k,0]*AU,*EL[k,1:],GMs+GM[k]); pos.append(r); vel.append(v)
    pos=np.array(pos); vel=np.array(vel); m=np.concatenate([[GMs],GM])
    pos-= (m[:,None]*pos).sum(0)/m.sum(); vel-=(m[:,None]*vel).sum(0)/m.sum()
    return np.concatenate([pos.ravel(),vel.ravel()]), m
def make_rhs(m,beta=0.0,active=None):
    n=len(m)
    act=np.ones(n,bool) if active is None else active
    def rhs(t,y):
        p=y[:3*n].reshape(n,3); v=y[3*n:]
        d=p[None,:,:]-p[:,None,:]; r2=(d*d).sum(2); np.fill_diagonal(r2,1.0); r3=r2**1.5
        w=(m*act)[None,:]/r3; np.fill_diagonal(w,0.0)
        a=(w[:,:,None]*d).sum(1)
        if beta:
            rs=p[1:]-p[0]; rr=np.linalg.norm(rs,axis=1)[:,None]
            extra=-2*beta*rs/rr**4           # delta-depth = -beta/r^2 around the Sun
            a[1:]+=extra
        return np.concatenate([v,a.ravel()])
    return rhs
def run(years,beta=0.0,active=None,step_days=2.0,rtol=1e-11):
    y0,m=initial(); rhs=make_rhs(m,beta,active)
    T=years*YR; te=np.arange(0,T,step_days*DAY)
    t0=time.time(); s=solve_ivp(rhs,(0,T),y0,t_eval=te,method='DOP853',rtol=rtol,atol=1.0)
    return s.t/YR, s.y, m, time.time()-t0
def perihelion_rate(t,y,m,k):
    n=len(m); p=y[:3*n].reshape(n,3,-1); v=y[3*n:].reshape(n,3,-1)
    r=(p[k]-p[0]).T; vv=(v[k]-v[0]).T; mu=m[0]+m[k]
    h=np.cross(r,vv); rn=np.linalg.norm(r,axis=1)[:,None]
    ev=np.cross(vv,h)/mu-r/rn
    w=np.unwrap(np.arctan2(ev[:,1],ev[:,0]))
    e=np.linalg.norm(ev,axis=1)
    return np.degrees(np.polyfit(t,w,1)[0])*3600*100, e   # arcsec per century
def elements_k(t,y,m,k):
    n=len(m); p=y[:3*n].reshape(n,3,-1); v=y[3*n:].reshape(n,3,-1)
    r=(p[k]-p[0]).T; vv=(v[k]-v[0]).T; mu=m[0]+m[k]
    h=np.cross(r,vv); rn=np.linalg.norm(r,axis=1)[:,None]
    ev=np.cross(vv,h)/mu-r/rn; e=np.linalg.norm(ev,axis=1)
    hn=h/np.linalg.norm(h,axis=1)[:,None]
    nodev=np.cross(np.array([0,0,1.0]),hn)                     # ascending node direction
    Om=np.arctan2(nodev[:,1],nodev[:,0])
    nn=nodev/np.linalg.norm(nodev,axis=1)[:,None]
    om=np.arctan2(np.einsum('ij,ij->i',np.cross(nn,ev),hn),np.einsum('ij,ij->i',nn,ev))
    varpi=np.unwrap(Om+om)
    return varpi,e,np.unwrap(Om),np.unwrap(om)
def rate(t,ang): return np.degrees(np.polyfit(t,ang,1)[0])*3600*100

if __name__=='__main__':
    # Memo TOE E.4 - Mercury's Perihelion. Sun and eight planets as summed dents.
    Y=1000
    base=np.zeros(9,bool); base[:2]=True
    t,y,m,_=run(Y,active=base,step_days=5,rtol=1e-13); vp,_,_,_=elements_k(t,y,m,1); floor=rate(t,vp)
    print(f"Sun only (floor): {floor:+.3f} arcsec/century")
    for k in range(2,9):
        act=base.copy(); act[k]=True
        t,y,m,_=run(Y,active=act,step_days=5,rtol=1e-13); vp,_,_,_=elements_k(t,y,m,1)
        print(f"  + {names[k-1]:11s} {rate(t,vp)-floor:8.2f}")
    t,y,m,_=run(Y,step_days=5,rtol=1e-13); vp,e,_,_=elements_k(t,y,m,1)
    print(f"ALL PLANETS: {rate(t,vp):.2f}; eccentricity change per century {np.polyfit(t,e,1)[0]*100:+.3e}")
    a=EL[0,0]*AU; e0=EL[0,1]; P=2*np.pi*np.sqrt(a**3/GMs)/YR
    beta=np.radians(43.0/3600)/(100/P)*GMs*a*(1-e0**2)/(2*np.pi); print(f"beta = {beta:.4e}")
    t2,y2,m2,_=run(Y,beta=beta,step_days=5,rtol=1e-13)
    for k,nm in [(1,'Mercury'),(2,'Venus'),(3,'Earth+Moon'),(4,'Mars')]:
        v1,_,_,_=elements_k(t,y,m,k); v2,_,_,_=elements_k(t2,y2,m2,k)
        print(f"  {nm:11s} extra from steeper dent {rate(t2,v2)-rate(t,v1):.3f}")