"""
    Numba implementation of nbody_opt.  Has same output as original program.
    
"""

from itertools import combinations
from numba import jit, int32, float64
import numpy as np


@jit("void(f8[:], f8[:], f8, f8, f8, f8, f8, f8)")
def update_vs(v1, v2, dt, dx, dy, dz, m1, m2):
    dx_dy_dz = ((dx**2 + dy**2 + dz**2) ** (-1.5))
    value_to_add_1 = m1 * dt * dx_dy_dz
    value_to_add_2 = m2 * dt * dx_dy_dz
    v1[0] -= dx * value_to_add_2
    v1[1] -= dy * value_to_add_2
    v1[2] -= dz * value_to_add_2
    v2[0] += dx * value_to_add_1
    v2[1] += dy * value_to_add_1
    v2[2] += dz * value_to_add_1

@jit("void(f8[:], f8, f8, f8, f8)")
def update_rs(r, dt, vx, vy, vz):
    r[0] += dt * vx
    r[1] += dt * vy
    r[2] += dt * vz

@jit("f8[:](f8[:], f8[:])")
def vec_deltas(x,y):
    return np.subtract(x,y)


@jit("void(i4, char, i4, f4)")
def report_energy(loops, reference, iterations, dt = 0.01):
    '''
        Combined version of the functions offset_momentum, report_energy, and advance.
        Start by offsetting the momentum. Then, perform the nested loop from the original 
        nbody function, but change it to one single loop.  In each iteration, either 
        perform the advance or report energy based on the iteration number.
    '''
    PI = 3.14159265358979323
    SOLAR_MASS = 4 * PI * PI
    DAYS_PER_YEAR = 365.24

    BODIES = {
    'sun': (np.array([0.0, 0.0, 0.0]), np.array([0.0, 0.0, 0.0]), SOLAR_MASS),

    'jupiter': (np.array([4.84143144246472090e+00,
                 -1.16032004402742839e+00,
                 -1.03622044471123109e-01]),
                np.array([1.66007664274403694e-03 * DAYS_PER_YEAR,
                 7.69901118419740425e-03 * DAYS_PER_YEAR,
                 -6.90460016972063023e-05 * DAYS_PER_YEAR]),
                9.54791938424326609e-04 * SOLAR_MASS),

    'saturn': (np.array([8.34336671824457987e+00,
                4.12479856412430479e+00,
                -4.03523417114321381e-01]),
               np.array([-2.76742510726862411e-03 * DAYS_PER_YEAR,
                4.99852801234917238e-03 * DAYS_PER_YEAR,
                2.30417297573763929e-05 * DAYS_PER_YEAR]),
               2.85885980666130812e-04 * SOLAR_MASS),

    'uranus': (np.array([1.28943695621391310e+01,
                -1.51111514016986312e+01,
                -2.23307578892655734e-01]),
               np.array([2.96460137564761618e-03 * DAYS_PER_YEAR,
                2.37847173959480950e-03 * DAYS_PER_YEAR,
                -2.96589568540237556e-05 * DAYS_PER_YEAR]),
               4.36624404335156298e-05 * SOLAR_MASS),

    'neptune': (np.array([1.53796971148509165e+01,
                 -2.59193146099879641e+01,
                 1.79258772950371181e-01]),
                np.array([2.68067772490389322e-03 * DAYS_PER_YEAR,
                 1.62824170038242295e-03 * DAYS_PER_YEAR,
                 -9.51592254519715870e-05 * DAYS_PER_YEAR]),
                5.15138902046611451e-05 * SOLAR_MASS)}

    
    #offset momentum
    px=py=pz=0.0
    
    
    for body in BODIES.keys():
        (r, [vx, vy, vz], m) = BODIES[body]
        px -= vx * m
        py -= vy * m
        pz -= vz * m
        
    (r, v, m) = BODIES[reference]
    v[0] = px / m
    v[1] = py / m
    v[2] = pz / m

    #change the nested loop from nbody function to one loop
    for i in range(loops*iterations):
        e = 0.0
        seenit = set()
        
        #use the built-in combinations func to iterate through the keys
        for body1, body2 in combinations(BODIES.keys(), 2):
            (x1, v1, m1) = BODIES[body1]
            (x2, v2, m2) = BODIES[body2]
            (dx, dy, dz) = vec_deltas(x1, x2)
            #depending on the value of i, either advance or report energy
            if i % iterations == 0:
                #step for report energy
                e -= (m1 * m2) / ((dx * dx + dy * dy + dz * dz) ** 0.5)
            else:
                #advance
                update_vs(v1, v2, dt, dx, dy, dz, m1, m2)
            seenit.add(body1)
        
        #again, check whether its an advance or report_energy step based on i
        for body in BODIES.keys():
            (r, [vx, vy, vz], m) = BODIES[body]
            if i % iterations == 0:
                e += m * (vx * vx + vy * vy + vz * vz) / 2.
            else:
                update_rs(r, dt, vx, vy, vz)
        #report energy
        if i % iterations == 0:
            print(e)



if __name__ == '__main__':
    report_energy(100, 'sun', 20000)

