import rebound
sim = rebound.Simulation()
sim.add(m=1., x=1., vz = 2.)

"""
Explanation: Orbital Elements
We can add particles to a simulation by specifying cartesian components:
End of explanation
"""


sim.add(m=1., a=1.)
sim.status()

"""
Explanation: Any components not passed automatically default to 0.  REBOUND can also accept orbital elements.  
Reference bodies
As a reminder, there is a one-to-one mapping between (x,y,z,vx,vy,vz) and orbital elements, and one should always specify what the orbital elements are referenced against (e.g., the central star, the system's barycenter, etc.).  The differences betwen orbital elements referenced to these centers differ by $\sim$ the mass ratio of the largest body to the central mass.  By default, REBOUND always uses Jacobi elements, which for each particle are always referenced to the center of mass of all particles with lower index in the simulation.  This is a useful set for theoretical calculations, and gives a logical behavior as the mass ratio increase, e.g., in the case of a circumbinary planet.  Let's set up a binary,
End of explanation
"""


sim.add(m=1.e-3, a=100.)

"""
Explanation: We always have to pass a semimajor axis (to set a length scale), but any other elements are by default set to 0.  Notice that our second star has the same vz as the first one due to the default Jacobi elements.  Now we could add a distant planet on a circular orbit,
End of explanation
"""


sim.add(primary=sim.particles[1], a=0.01)

"""
Explanation: This planet is set up relative to the binary center of mass (again due to the Jacobi coordinates), which is probably what we want.  But imagine we now want to place a test mass in a tight orbit around the second star.  If we passed things as above, the orbital elements would be referenced to the binary/outer-planet center of mass.  We can override the default by explicitly passing a primary (any instance of the Particle class):
End of explanation
"""


orbits = sim.calculate_orbits()
for orbit in orbits:
    print(orbit)

"""
Explanation: All simulations are performed in Cartesian elements, so to avoid the overhead, REBOUND does not update particles' orbital elements as the simulation progresses.  However, we can always calculate them when required with sim.calculate_orbits().  Note that REBOUND will always output angles in the range $[-\pi,\pi]$, except the inclination which is always in $[0,\pi]$.
End of explanation
"""


print(sim.particles[3].calculate_orbit(sim, primary=sim.particles[1]))

"""
Explanation: Notice that there is always one less orbit than there are particles, since orbits are only defined between pairs of particles.  We see that we got the first two orbits right, but the last one is way off.  The reason is that again the REBOUND default is that we always get Jacobi elements.  But we initialized the last particle relative to the second star, rather than the center of mass of all the previous particles.
To get orbital elements relative to a specific body, you can manually use the calculate_orbit method of the Particle class:
End of explanation
"""


sim = rebound.Simulation()
sim.add(m=1.)
sim.add(a=1., e=0., inc=0.1, Omega=0.3, omega=0.1)
orbits = sim.calculate_orbits()
print(orbits[0])

"""
Explanation: though we could have simply avoided this problem by adding bodies from the inside out (second star, test mass, first star, circumbinary planet).
Edge cases and orbital element sets
Different orbital elements lose meaning in various limits, e.g., a planar orbit and a circular orbit.  REBOUND therefore allows initialization with several different types of variables that are appropriate in different cases.  It's important to keep in mind that the procedure to initialize particles from orbital elements is not exactly invertible, so one can expect discrepant results for elements that become ill defined.  For example,
End of explanation
"""


print(orbits[0].theta)

"""
Explanation: The problem here is that $\omega$ (the angle from the ascending node to pericenter) is ill-defined for a circular orbit, so it's not clear what we mean when we pass it, and we get spurious results (i.e., $\omega = 0$ rather than 0.1, and $f=0.1$ rather than the default 0).  Similarly, $f$, the angle from pericenter to the particle's position, is undefined.  However, the true longitude $\theta$, the broken angle from the $x$ axis to the ascending node = $\Omega + \omega + f$, and then to the particle's position, is always well defined:
End of explanation
"""


sim = rebound.Simulation()
sim.add(m=1.)
sim.add(a=1., e=0., inc=0.1, Omega=0.3, theta = 0.4)
orbits = sim.calculate_orbits()
print(orbits[0].theta)

import rebound
sim = rebound.Simulation()
sim.add(m=1.)
sim.add(a=1., e=0.2, Omega=0.1)
orbits = sim.calculate_orbits()
print(orbits[0])

"""
Explanation: To be clearer and ensure we get the results we expect, we could instead pass theta to specify the longitude we want, e.g.
End of explanation
"""


print(orbits[0].pomega)

"""
Explanation: Here we have a planar orbit, in which case the line of nodes becomes ill defined, so $\Omega$ is not a good variable, but we pass it anyway!  In this case, $\omega$ is also undefined since it is referenced to the ascending node.  Here we get that now these two ill-defined variables get flipped.  The appropriate variable is pomega ($\varpi = \Omega + \omega$), which is the angle from the $x$ axis to pericenter:
End of explanation
"""


import rebound
sim = rebound.Simulation()
sim.add(m=1.)
sim.add(a=1., e=0.2, pomega=0.1)
orbits = sim.calculate_orbits()
print(orbits[0])

"""
Explanation: We can specify the pericenter of the orbit with either $\omega$ or $\varpi$:
End of explanation
"""


sim = rebound.Simulation()
sim.add(m=1.)
sim.add(a=1., e=0.1, Omega=0.3, M = 0.1)
sim.add(a=1., e=0.1, Omega=0.3, l = 0.4)
orbits = sim.calculate_orbits()
print(orbits[0].l)
print(orbits[1].l)

import rebound
sim = rebound.Simulation()
sim.add(m=1.)
sim.add(a=1., e=0.1, omega=1.)
orbits = sim.calculate_orbits()
print(orbits[0])

"""
Explanation: Note that if the inclination is exactly zero, REBOUND sets $\Omega$ (which is undefined) to 0, so $\omega = \varpi$.  
Finally, we can initialize particles using mean, rather than true, longitudes or anomalies (for example, this might be useful for  resonances).  We can either use the mean anomaly $M$, which is referenced to pericenter (again ill-defined for circular orbits), or its better-defined counterpart the mean longitude l $= \lambda = \Omega + \omega + M$, which is analogous to $\theta$ above,
End of explanation
"""


import random
import numpy as np

def simulation(par):
    e,f = par
    e = 10**e
    f = 10**f
    sim = rebound.Simulation()
    sim.add(m=1.)
    a = 1.
    inc = random.random()*np.pi
    Omega = random.random()*2*np.pi
    sim.add(m=0.,a=a,e=e,inc=inc,Omega=Omega, f=f)
    o=sim.calculate_orbits()[0]
    if o.f < 0: # avoid wrapping issues
        o.f += 2*np.pi
    err = max(np.fabs(o.e-e)/e, np.fabs(o.f-f)/f)
    return err

random.seed(1)
N = 100
es = np.linspace(-16.,-1.,N)
fs = np.linspace(-16.,-1.,N)
params = [(e,f) for e in es for f in fs]

pool=rebound.InterruptiblePool()
res = pool.map(simulation, params)
res = np.array(res).reshape(N,N)
res = np.nan_to_num(res)

%matplotlib inline
import matplotlib.pyplot as plt
from matplotlib import ticker
from matplotlib.colors import LogNorm
import matplotlib

f,ax = plt.subplots(1,1,figsize=(7,5))
extent=[fs.min(), fs.max(), es.min(), es.max()]

ax.set_xlim(extent[0], extent[1])
ax.set_ylim(extent[2], extent[3])
ax.set_xlabel(r"true anomaly (f)")
ax.set_ylabel(r"eccentricity")

im = ax.imshow(res, norm=LogNorm(), vmax=1., vmin=1.e-16, aspect='auto', origin="lower", interpolation='nearest', cmap="RdYlGn_r", extent=extent)
cb = plt.colorbar(im, ax=ax)
cb.solids.set_rasterized(True)
cb.set_label("Relative Error")

"""
Explanation: Accuracy
As a test of accuracy and demonstration of issues related to the last section, let's test the numerical stability by intializing particles with small eccentricities and true anomalies, computing their orbital elements back, and comparing the relative error.  We choose the inclination and node longitude randomly:
End of explanation
"""


def simulation(par):
    e,theta = par
    e = 10**e
    theta = 10**theta
    sim = rebound.Simulation()
    sim.add(m=1.)
    a = 1.
    inc = random.random()*np.pi
    Omega = random.random()*2*np.pi
    omega = random.random()*2*np.pi
    sim.add(m=0.,a=a,e=e,inc=inc,Omega=Omega, theta=theta)
    o=sim.calculate_orbits()[0]
    if o.theta < 0:
        o.theta += 2*np.pi
    err = max(np.fabs(o.e-e)/e, np.fabs(o.theta-theta)/theta)
    return err

random.seed(1)
N = 100
es = np.linspace(-16.,-1.,N)
thetas = np.linspace(-16.,-1.,N)
params = [(e,theta) for e in es for theta in thetas]

pool=rebound.InterruptiblePool()
res = pool.map(simulation, params)
res = np.array(res).reshape(N,N)
res = np.nan_to_num(res)

f,ax = plt.subplots(1,1,figsize=(7,5))
extent=[thetas.min(), thetas.max(), es.min(), es.max()]

ax.set_xlim(extent[0], extent[1])
ax.set_ylim(extent[2], extent[3])
ax.set_xlabel(r"true longitude (\theta)")
ax.set_ylabel(r"eccentricity")

im = ax.imshow(res, norm=LogNorm(), vmax=1., vmin=1.e-16, aspect='auto', origin="lower", interpolation='nearest', cmap="RdYlGn_r", extent=extent)
cb = plt.colorbar(im, ax=ax)
cb.solids.set_rasterized(True)
cb.set_label("Relative Error")

"""
Explanation: We see that the behavior is poor, which is physically due to $f$ becoming poorly defined at low $e$.  If instead we initialize the orbits with the true longitude $\theta$ as discussed above, we get much better results:
End of explanation
"""


sim.add(a=-0.2, e=1.4)
sim.status()

"""
Explanation: Hyperbolic & Parabolic Orbits
REBOUND can also handle hyperbolic orbits, which have negative $a$ and $e>1$:
End of explanation
"""


sim = rebound.Simulation()
sim.add(m=1.)
q = 0.1
a=-1.e14
e=1.+q/np.fabs(a)
sim.add(a=a, e=e)
print(sim.calculate_orbits()[0])

"""
Explanation: Currently there is no support for exactly parabolic orbits, but we can get a close approximation by passing a nearby hyperbolic orbit where we can specify the pericenter = $|a|(e-1)$ with $a$ and $e$.  For example, for a 0.1 AU pericenter,
End of explanation
"""


sim = rebound.Simulation()
sim.add(m=1.)
sim.add(a=1.,inc=np.pi,e=0.1, Omega=0., pomega=1.)
print(sim.calculate_orbits()[0])

"""
Explanation: Retrograde Orbits
Orbital elements can be counterintuitive for retrograde orbits, but REBOUND tries to sort them out consistently.  This can lead to some initially surprising results.  For example,
End of explanation
"""


