import rebound
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import os
import seaborn as sns
from matplotlib.lines import Line2D

d2r = np.pi/180.
sim = rebound.Simulation()
sim.add(m=1.7, x=0.)
sim.add(primary=sim.particles[0], m=1.e-10, a=0.56, e=0., inc=14.9*d2r, Omega=183.5*d2r)
sim.add(primary=sim.particles[0], m=1.e-10, a=20., e=0., inc=14.9*d2r, Omega=183.5*d2r)
sim.add(primary=sim.particles[0], m=1.e-10, a=22.6, e=0., inc=14.9*d2r, Omega=183.5*d2r)
sim.add(primary=sim.particles[0],
              m=0.2,
              a=207.,
              e=0.32,
              inc=49*d2r,
              omega=18.*d2r,
              f=180.*d2r,
              Omega=47.*d2r)
sim.move_to_com()
npart = 3
labels=['Inner disc','Inner companion','Outer disc']
colours=[col_list[1],col_list[4],col_list[7]]


# Now integrate through time
torb = 2.*np.pi*np.sqrt((sim.particles[(npart+1)].a**3/(sim.particles[(npart+1)].m + sim.particles[0].m)))
orbitinyrs = torb/(2*np.pi)
max_orb = 2000
tmax = torb*max_orb
N_out = 5000
a_out = np.zeros((N_out,npart+1))
e_out = np.zeros((N_out,npart+1))
inc_out = np.zeros((N_out,npart+1))
times = np.linspace(0,tmax,N_out)
unitLs = np.zeros((npart+1,3))
relative_misalignments = np.zeros((N_out,npart))
for i, time in enumerate(times):
    sim.integrate(time)
    # Calculate the properties with respect to the primary
    for j,p in enumerate(sim.particles[1:]):
        orbits = p.calculate_orbit()
        pos = [p.x - sim.particles[0].x, p.y - sim.particles[0].y, p.z - sim.particles[0].z]
        vel = [p.vx - sim.particles[0].vx, p.vy - sim.particles[0].vy, p.vz - sim.particles[0].vz]
        L = np.cross(pos,vel)
        unitL = L / np.linalg.norm(L)
        unitLs[j] = unitL
        a_out[i,j] = orbits.a
        e_out[i,j] = orbits.e
        inc_out[i,j] = orbits.inc
    relative_misalignments[i,0] = np.arccos(np.dot(unitLs[0,:],unitLs[3,:]))/d2r
    relative_misalignments[i,1] = np.arccos(np.dot(unitLs[1,:],unitLs[3,:]))/d2r
    relative_misalignments[i,2] = np.arccos(np.dot(unitLs[2,:],unitLs[3,:]))/d2r

# Plot those results
sns.set_style("ticks", {
        "font.family": "serif",
        "font.serif": ["Times", "Palatino", "serif"],
        "font.size" : 16
    })
plt.clf()
fig, ax = plt.subplots(nrows = 1, ncols = 1, sharex = False, sharey = False, figsize=(10,4))
ax1 = ax.twiny()

for i in range(npart):
    ax.plot(times/torb*orbitinyrs,relative_misalignments[:,i],label=labels[i],color=colours[i])
    
ax.set_xlabel(r"Time (years)")
ax.set_ylabel(r"Relative misalignment (degrees)")


ax.set_xlim([0,2000*orbitinyrs])
ax.set_xticklabels([0,r"$0.5 \times 10^6$",r"$1 \times 10^6$",r"$1.5 \times 10^6$", r"$2 \times 10^6$",r"$2.5 \times 10^6$",r"$3 \times 10^6$",r"$3.5 \times 10^6$",r"$4 \times 10^6$"])
ax.set_ylim([0,90])
ax.legend(loc='lower left')

# Twin axes
ax1.set_xticks([0,500,1000,1500,2000])
ax1.set_xlabel("Time (binary orbits)")

# Add a custom legend for the RH plot
custom_lines = [Line2D([0], [0], color='k', lw=2),
                Line2D([0], [0], color='k', lw=2,linestyle='--')]