# exclude this line on your local environment if you already have REBOUND installed
!pip install rebound


import numpy as np
import matplotlib.pyplot as plt
import math
import rebound


sim = rebound.Simulation()
sim.add("Sun")
sim.add("Jupiter")
sim.add("Saturn")
sim.add("Uranus")
sim.add("Neptune")

sim.move_to_com()
sim.integrator = "whfast"
sim.t = 0

sim.collision_resolve = "merge"

numberOfActiveParticles = 5 # sun & major planets

aNeptune = sim.particles[4].a


numberOfTestParticles = 400

rayleighScale = 0.03

integrationTime = 1e5

sim.dt = 0.1

a_initial = np.random.uniform(1.82 * aNeptune, 1.86 * aNeptune, numberOfTestParticles)
e_initial = np.random.uniform(0, 1, numberOfTestParticles)

i_initial = np.random.rayleigh(rayleighScale, numberOfTestParticles)
i_initial /= np.max(i_initial); i_initial *= (np.pi / 2) / i_initial.max()  # scale to 0-90 degrees

M_initial = np.random.uniform(-(np.pi), np.pi, numberOfTestParticles)
omega_initial = np.random.uniform(-(np.pi), np.pi, numberOfTestParticles)
Omega_initial = np.random.uniform(-(np.pi), np.pi, numberOfTestParticles)


for (a, e, i, M, omega, Omega) in zip(a_initial, e_initial, i_initial, M_initial, omega_initial, Omega_initial):
    sim.add(
        a = a,
        e = e,
        inc = i,
        M = M, # mean anomaly
        omega = omega, # argument of pericenter
        Omega = Omega # longitude of ascending node
    )
    
sim.N_active = numberOfActiveParticles

initialXs = np.array(list(map(lambda x: x.x, sim.particles[5:])))
initialYs = np.array(list(map(lambda x: x.y, sim.particles[5:])))
initialZs = np.array(list(map(lambda x: x.z, sim.particles[5:])))

plt.figure(figsize=(18, 6))
plt.tight_layout()

plt.subplot(1, 5, 1); plt.xlabel("a"); plt.ylabel("#")
plt.hist(a_initial);

plt.subplot(1, 5, 2); plt.xlabel("a"); plt.ylabel("e")
plt.scatter(a_initial, e_initial);

plt.subplot(1, 5, 4); plt.xlabel("e"); plt.ylabel("#")
plt.hist(e_initial);

plt.subplot(1, 5, 3); plt.xlabel("a"); plt.ylabel("i")
plt.scatter(a_initial, list(map(lambda x: math.degrees(x), i_initial)));

plt.subplot(1, 5, 5); plt.xlabel("i"); plt.ylabel("#")
plt.hist(list(map(lambda x: math.degrees(x), i_initial)));


outputPoints = int(integrationTime / 1000)

times = np.linspace(0., integrationTime, outputPoints)

for i, time in enumerate(times):
    sim.integrate(time)


a_final = np.array(list(map(lambda x: x.a / aNeptune, sim.particles[5:])))
e_final = np.array(list(map(lambda x: x.e, sim.particles[5:])))
i_final = np.array(list(map(lambda x: math.degrees(x.inc), sim.particles[5:])))

plt.figure(figsize=(15, 5))
plt.tight_layout()

plt.subplot(1, 4, 1); plt.xlabel("a / aN"); plt.ylabel("#")
plt.hist(a_final, bins=np.arange(1.80, 1.89, 0.01));
plt.xlim(1.80, 1.88);

plt.subplot(1, 4, 2); plt.xlabel("a / aN"); plt.ylabel("e")
plt.scatter(a_final, e_final);
plt.xlim(1.80, 1.88)
plt.ylim(0, 1)

plt.subplot(1, 4, 3); plt.xlabel("e"); plt.ylabel("#")
plt.hist(e_final, bins=np.arange(0, 1.1, 0.1));
plt.xlim(0, 1)

plt.subplot(1, 4, 4); plt.xlabel("i"); plt.ylabel("#")
plt.hist(i_final);


def orbitMap(xs, ys, zs):
  a = 100 # AU
  
  plt.figure(figsize=(15, 7))
  plt.tight_layout()

  plt.subplot(1, 2, 1); plt.xlabel("x"); plt.ylabel("y")
  plt.scatter(xs, ys);
  plt.xlim(-a, a)
  plt.ylim(-a, a)

  plt.subplot(1, 2, 2); plt.xlabel("x"); plt.ylabel("z")
  plt.scatter(xs, zs);
  plt.xlim(-a, a)
  plt.ylim(-a, a);

orbitMap(initialXs, initialYs, initialZs)

xs = np.array(list(map(lambda x: x.x, sim.particles[5:])))
ys = np.array(list(map(lambda x: x.y, sim.particles[5:])))
zs = np.array(list(map(lambda x: x.z, sim.particles[5:])))

orbitMap(xs, ys, zs)


#rebound.OrbitPlot(sim, lim=70, slices=True, unitlabel="[AU]", color=False, periastron=False);