import rebound
import numpy as np

"""
Explanation: Using Variational Equations With the Chain Rule
For a complete introduction to variational equations, please read the paper by Rein and Tamayo (2016).
Variational equations can be used to calculate derivatives in an $N$-body simulation. More specifically, given a set of initial conditions $\alpha_i$ and a set of variables at the end of the simulation $v_k$, we can calculate all first order derivatives
$$\frac{\partial v_k}{\partial \alpha_i}$$
as well as all second order derivates
$$\frac{\partial^2 v_k}{\partial \alpha_i\partial \alpha_j}$$
For this tutorial, we work with a two planet system. 
We first chose the semi-major axis $a$ of the outer planet as an initial condition (this is our $\alpha_i$). At the end of the simulation we output the velocity of the star in the $x$ direction (this is our $v_k$). 
To do that, let us first import REBOUND and numpy.
End of explanation
"""


def calculate_vx(a):
    sim = rebound.Simulation()
    sim.add(m=1.) # star
    sim.add(primary=sim.particles[0],m=1e-3, a=1) # inner planet
    sim.add(primary=sim.particles[0],m=1e-3, a=a) # outer planet
    sim.integrate(2.*np.pi*10.) # integrate for ~10 orbits
    return sim.particles[0].vx  # return star's velocity in the x direction

calculate_vx(a=1.5) # initial semi-major axis of the outer planet is 1.5

"""
Explanation: The following function takes $a$ as a parameter, then integrates the two planet system and returns the velocity of the star at the end of the simulation.
End of explanation
"""


calculate_vx(a=1.51) # initial semi-major axis of the outer planet is 1.51

"""
Explanation: If we run the simulation again, with a different initial $a$, we get a different velocity:
End of explanation
"""


def calculate_vx_derivative(a):
    sim = rebound.Simulation()
    sim.add(m=1.) # star
    sim.add(primary=sim.particles[0],m=1e-3, a=1) # inner planet
    sim.add(primary=sim.particles[0],m=1e-3, a=a) # outer planet
    
    v1 = sim.add_variation()    # add a set of variational particles
    v1.vary(2,"a")              # initialize the variational particles 
    
    sim.integrate(2.*np.pi*10.) # integrate for ~10 orbits
    return sim.particles[0].vx, v1.particles[0].vx # return star's velocity and its derivative

"""
Explanation: We could now run many different simulations to map out the parameter space. This is a very simple example of a typical use case: the fitting of a radial velocity datapoint. 
However, we can be smarter than simple running an almost identical simulation over and over again by using variational equations. These will allow us to calculate the derivate of the stellar velocity at the end of the simulation. We can take derivative with respect to any of the initial conditions, i.e. a particle's mass, semi-major axis, x-coordinate, etc. Here, we want to take the derivative with respect to the semi-major axis of the outer planet. The following function does exactly that:
End of explanation
"""


calculate_vx_derivative(a=1.5) 

"""
Explanation: Note the two new functions. sim.add_variation() adds a set of variational particles to the simulation. All variational particles are by default initialized to zero. We use the vary() function to initialize them to a variation that we are interested in. Here, we initialize the variational particles corresponding to a change in the semi-major axis, $a$, of the particle with index 2 (the outer planet).
End of explanation
"""


a0=1.5
va0, dva0 = calculate_vx_derivative(a=a0) 
def v(a):
    return va0 + (a-a0)*dva0
print(v(1.51))

"""
Explanation: We can use the derivative to construct a Taylor series expansion of the velocity around $a_0=1.5$:
$$v(a) \approx v(a_0) + (a-a_0) \frac{\partial v}{\partial a}$$
End of explanation
"""


def calculate_vx_derivative_2ndorder(a):
    sim = rebound.Simulation()
    sim.add(m=1.) # star
    sim.add(primary=sim.particles[0],m=1e-3, a=1) # inner planet
    sim.add(primary=sim.particles[0],m=1e-3, a=a) # outer planet
    
    v1 = sim.add_variation()   
    v1.vary(2,"a")             
    
    # The following lines add and initialize second order variational particles
    v2 = sim.add_variation(order=2, first_order=v1)   
    v2.vary(2,"a")  
    
    sim.integrate(2.*np.pi*10.) # integrate for ~10 orbits
    # return star's velocity and its first and second derivatives
    return sim.particles[0].vx, v1.particles[0].vx, v2.particles[0].vx 

"""
Explanation: Compare this value with the explicitly calculate one above. They are almost the same! But we can do even better, by using second order variational equations to calculate second order derivatives.
End of explanation
"""


a0=1.5
va0, dva0, ddva0 = calculate_vx_derivative_2ndorder(a=a0) 
def v(a):
    return va0 + (a-a0)*dva0 + 0.5*(a-a0)**2*ddva0
print(v(1.51))

"""
Explanation: Using a Taylor series expansion to second order gives a better estimate of v(1.51).
End of explanation
"""


def calculate_w_derivative(a):
    sim = rebound.Simulation()
    sim.add(m=1.) # star
    sim.add(primary=sim.particles[0],m=1e-3, a=1) # inner planet
    sim.add(primary=sim.particles[0],m=1e-3, a=a) # outer planet
    
    v1 = sim.add_variation()    # add a set of variational particles
    v1.vary(2,"a")              # initialize the variational particles 
    
    sim.integrate(2.*np.pi*10.) # integrate for ~10 orbits
    
    c = 1.02 # some constant
    w = (sim.particles[0].vx-c)**2
    dwda = 2.*v1.particles[0].vx * (sim.particles[0].vx-c)
    
    return w, dwda # return w and its derivative
calculate_w_derivative(1.5)

"""
Explanation: Now that we know how to calculate first and second order derivates of positions and velocities of particles, we can simply use the chain rule to calculate more complicated derivates. For example, instead of the velocity $v_x$, you might be interested in the quantity $w\equiv(v_x - c)^2$ where $c$ is a constant. This is something that typically appears in a $\chi^2$ fit. The chain rule gives us:
$$ \frac{\partial w}{\partial a} = 2 \cdot (v_x-c)\cdot \frac{\partial v_x}{\partial a}$$
The variational equations provide the $\frac{\partial v_x}{\partial a}$ part, the ordinary particles provide $v_x$.
End of explanation
"""


def calculate_vx_derivative_h():
    h, k = 0.1, 0.2
    e = float(np.sqrt(h**2+k**2))
    omega = np.arctan2(k,h)
    sim = rebound.Simulation()
    sim.add(m=1.) # star
    sim.add(primary=sim.particles[0],m=1e-3, a=1) # inner planet
    sim.add(primary=sim.particles[0],m=1e-3, a=1.5, e=e, omega=omega) # outer planet
    
    v1 = sim.add_variation()   
    dpde     = rebound.Particle(simulation=sim, particle=sim.particles[2], variation="e")
    dpdomega = rebound.Particle(simulation=sim, particle=sim.particles[2], m=1e-3, a=1.5, e=e, omega=omega, variation="omega")
    v1.particles[2] = h/e * dpde - k/(e*e) * dpdomega
        
    sim.integrate(2.*np.pi*10.) # integrate for ~10 orbits
    # return star's velocity and its first derivatives
    return sim.particles[0].vx, v1.particles[0].vx
calculate_vx_derivative_h()

"""
Explanation: Similarly, you can also use the chain rule to vary initial conditions of particles in a way that is not supported by REBOUND by default. For example, suppose you want to work in some fancy coordinate system, using $h\equiv e\sin(\omega)$ and $k\equiv e \cos(\omega)$ variables instead of $e$ and $\omega$. You might want to do that because $h$ and $k$ variables are often better behaved near $e\sim0$. In that case the chain rule gives us:
$$\frac{\partial p(e(h, k), \omega(h, k))}{\partial h} = \frac{\partial p}{\partial e}\frac{\partial  e}{\partial h} + \frac{\partial p}{\partial \omega}\frac{\partial \omega}{\partial h}$$
where $p$ is any of the particles initial coordinates. In our case the derivates of $e$ and $\omega$ with respect to $h$ are:
$$\frac{\partial \omega}{\partial h} = -\frac{k}{e^2}\quad\text{and}\quad \frac{\partial e}{\partial h} = \frac{h}{e}$$
With REBOUND, you can easily implement this. The following function calculates the derivate of the star's velocity with respect to the outer planet's $h$ variable.
End of explanation
"""


