#!/usr/bin/env python3


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

solarmass = 1.989e30
au = 1.496e11
gravitationalconstant = 6.67408e-11

# create a sim object
sim = rebound.Simulation()
sim.G = gravitationalconstant

# sun
sim.add(m=solarmass)

# jupiter
x  = -3.413527132484646E+00*au;
y = -4.178200275190616E+00*au;
z =  9.368355203631962E-02*au;
vx = 5.754715843387071E-03*au/(3600*24);
vy = -4.414713409455822E-03*au/(3600*24);
vz = -1.103770928813920E-04*au/(3600*24);
m  = 1e-3*solarmass;

sim.add(m=m, x=x, y=y, z=z, vx=vx, vy=vy, vz=vz)

# mars
x  = 1.831636514321734E-01*au;
y = -1.422003023901739E+00*au;
z = -3.451000894109571E-02*au;
vx = 1.440153096004164E-02*au/(3600*24);
vy = 2.977634949564440E-03*au/(3600*24);
vz = -2.910984559065813E-04*au/(3600*24);
m  = 3.2e-7*solarmass;

sim.add(m=m, x=x, y=y, z=z, vx=vx, vy=vy, vz=vz)


# works only in au coordinates....
#sim.add(["Sun", "Mercury", "Jupiter"])


# now add 10000 inactive tracerparticles between 2 and 4 au
N_testparticle = 10000
a_ini = np.linspace(2*au, 4*au, N_testparticle)
for a in a_ini:
    sim.add(a=a, f=np.random.rand()*2.*np.pi, e=0.5*np.random.rand()) # mass is set to 0 by default, random true anomaly,random eccentricity with max 0.5

# set the integrator to type REB_INTEGRATOR_LEAPFROG
sim.integrator = "leapfrog"

# create time array, let's say 1 orbit, plot 250 times per orbit
oneyearinseconds = 3600*24*365.25
simulationtime = 1e3*oneyearinseconds # simulation time
Nsteps = 1000
times = np.linspace(0, simulationtime, Nsteps)

# set the time step
# we use 1e-2 of an orbit at 2 au
orbit = 2*np.pi*np.sqrt(8*au*au*au/(gravitationalconstant*solarmass));

sim.dt = orbit*1e-2


sim.N_active = 3
sim.move_to_com()
# now integrate
semimajoraxis = np.zeros(N_testparticle)
eccentricity = np.zeros(N_testparticle)
for i, t in enumerate(times):
    print(t/simulationtime, end="\r")
    sim.integrate(t, exact_finish_time=1)
    for j, p in enumerate(sim.particles[3:]):
        semimajoraxis[j] = p.a/au
        eccentricity[j] = p.e
    fig, ax = plt.subplots()
    plt.grid(True)
    ax.scatter(semimajoraxis, eccentricity, s=1, c='r')
    ax.set_xlabel("semi-major axis [au]")
    ax.set_ylabel("eccentricity")
    ax.set_xlim(2,4)
    ax.set_ylim(0,1)
    fig.savefig("orbit"+ ("%05d_%g" % (i, sim.t/oneyearinseconds)) + ".png")
    np.savetxt("orbit"+ ("%05d_%g" % (i, sim.t/oneyearinseconds)) + ".txt", np.transpose([semimajoraxis, eccentricity]))
    plt.close()


# store if we want to run it longer later on
sim.save("checkpoint.bin")
# binned asteroids
semimajoraxis = semimajoraxis[semimajoraxis < 5]
semimajoraxis = semimajoraxis[semimajoraxis > 1]
bins = 700
plt.grid(True)
plt.hist(semimajoraxis, bins=bins, label="after %g years" % (simulationtime/oneyearinseconds))
binlength = (np.max(semimajoraxis) - np.min(semimajoraxis))/bins
plt.title("Asteroids, %d bins" % bins)
plt.xlabel("semi major axis [au]")
plt.ylabel("number of asteroids, binlength= %f au" % binlength)
plt.legend()
plt.savefig("bins.pdf")