#!/usr/bin/env python3

"""
Two-body problem in Newtonian gravity (Keplerian orbit), as well as in General Relativity (GR).
Orbit in GR is calculated using the parameterized post-Newtonian (PPN) equation of motion,
taking into account that the PPN parameters in GR are: beta = 1 and gamma = 1. Both Keplerian
and GR orbits are calculated by numerical integration of the corresponding equation of motion,
using the Keplerian positions and velocities at apocenter as initial conditions.
"""

import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import odeint

# Functions

def apocenter(g,m1,m2,a,e):
    """Calculates Keplerian orbital period, position and velocity at apocenter"""
    gm=g*(m1+m2)
    p=2.0*np.pi*np.sqrt(a**3/gm)
    r=a*(1.0+e)
    v=np.sqrt(gm*(2.0/r-1.0/a))
    return (p,r,v)

def ppncorr(par,t,c,g,beta,gamma,m1,m2):
    """Calculates PPN correction for acceleration (PPN parameters for GR are: beta=1.0 i gamma=1.0)"""
    r=np.sqrt(par[0]**2+par[1]**2)
    vdotv=par[2]**2+par[3]**2
    rdotv=par[0]*par[2]+par[1]*par[3]
    dax=g*m1*((2.0*(beta+gamma)*g*m1/r-gamma*vdotv)*par[0]+2.0*(1.0+gamma)*rdotv*par[2])/(c**2*r**3)
    day=g*m1*((2.0*(beta+gamma)*g*m1/r-gamma*vdotv)*par[1]+2.0*(1.0+gamma)*rdotv*par[3])/(c**2*r**3)
    return (dax,day)

def nwtderiv(par,t,c,g,m1,m2):
    """Calculates acceleration in Newtonian gravity as the negative gradient of the gravitational potential"""
    r=np.sqrt(par[0]**2+par[1]**2)
    ax=-g*(m1+m2)*par[0]/r**3
    ay=-g*(m1+m2)*par[1]/r**3
    return np.array([par[2],par[3],ax,ay])

def grderiv(par,t,c,g,m1,m2):
    """Calculates acceleration in GR by adding the PPN correction to the Newtonian acceleration"""
    beta=1.0
    gamma=1.0
    r=np.sqrt(par[0]**2+par[1]**2)
    dax,day=ppncorr(par,t,c,g,beta,gamma,m1,m2)
    ax=-g*(m1+m2)*par[0]/r**3+dax
    ay=-g*(m1+m2)*par[1]/r**3+day
    return np.array([par[2],par[3],ax,ay])

# Constants

# Speed of light (in AU/yr)
c=63115.2
# Gravitational constant (in AU^3/(Msun*yr^2))
g=4.0*np.pi**2

# Input data

# Masses of two bodies (in Msun)
m1=4.3e6
m2=1.0
# Semi-major axis (in AU)
a=841.0
# Eccentricity
e=0.985
# Integration period (in orbital periods)
ip=150
# Integration step (in yr)
dt=0.01

# Initial conditions at apocenter

# Keplerian orbital period, position and velocity
p,r,v=apocenter(g,m1,m2,a,e)
start=np.array([r,0.0,0.0,v])
print("Keplerian orbital period: {0:.3f} yr\n".format(p))

# Numerical integration of the Keplerian and GR orbits

# Integration times
ti=np.arange(0.0,ip*p+dt,dt)

# Integration of Keplerian orbit at the specified times and with given tolerance
kplorbit=odeint(nwtderiv,start,ti,args=(c,g,m1,m2),rtol=1.0e-12)

# Integration of GR orbit at the specified times and with given tolerance
grorbit=odeint(grderiv,start,ti,args=(c,g,m1,m2),rtol=1.0e-12)

# Figure settings
plt.rcParams.update({'font.family':'serif','font.size':14,'text.usetex':True,'legend.framealpha':None,'legend.edgecolor':'black'})

# Output figure
plt.plot(grorbit[:,0],grorbit[:,1],color="blue",label="GR")
plt.plot(kplorbit[:,0],kplorbit[:,1],color="red",label="Kepler")
plt.axhline(0,color="black",linestyle="dotted")
plt.axvline(0,color="black",linestyle="dotted")
plt.xlabel("$x$ (AU)")
plt.ylabel("$y$ (AU)")
plt.legend(loc="best")
plt.gca().set_aspect("equal")
plt.tight_layout()
plt.savefig("s4714b.jpg",dpi=300,bbox_inches="tight")
plt.show()
