import numpy as np
import matplotlib.pyplot as plt
from itertools import combinations

G = 1
q_n = 10 # number of bodies
N = 1000
tau = 10000
dt = tau/float(N-1)
t = 0

class Body:
    '''Class representing a 2D particle'''
    def __init__(self, x, y, vx, vy, ax, ay, m):

        self.r = np.array([x, y], dtype=float)
        self.v = np.array([vx, vy], dtype=float)
        self.a = np.array([ax, ay], dtype=float)
        self.m = m


    def acc(self, other):
        '''
        newtons law of universal gravitation
        and newtons 2nd law gives a = F/m
        '''
        d = (other.r - self.r)

        if not np.all(d):

            tmp = 0

        else:

            tmp = ((G * other.m)/(np.linalg.norm(d))**2)*(d/np.linalg.norm(d))

        self.a += tmp


    def leapfrog(self):
        '''leapfrog integrator, for solving the movement
        and velocity of a given particle with an acceleration.
        '''
        v_half = self.v * dt/2
        r_next = self.r + v_half * dt
        v_next = v_half + self.a * dt

        self.r = r_next
        self.v = v_next

        return r_next, v_next

particles = [Body(np.random.uniform(-1e3, 1e3), np.random.uniform(-1e3, 1e3), np.random.uniform(-1e1, 1e1), np.random.uniform(-1e1, 1e1), 0, 0, 1) for i in range(q_n)]

'''
Listor som innehåller q_n antal listor,
varje element har N (time step) element
'''
xdata = [np.zeros([N,1]) for i in range(q_n)]
ydata = [np.zeros([N,1]) for i in range(q_n)]
vxdata = [np.zeros([N,1]) for i in range(q_n)]
vydata = [np.zeros([N,1]) for i in range(q_n)]

# computation loop
for i in range(N):

    pairs = combinations(range(q_n), 2)

    for j, k in pairs:

        Body.acc(particles[j], particles[k])

    for n in range(q_n):

        tmp1, tmp2 = Body.leapfrog(particles[n])
        xdata[n][i] = tmp1[0]
        ydata[n][i] = tmp1[1]
        vxdata[n][i] = tmp2[0]
        vydata[n][i] = tmp2[1]

for i in range(q_n):
    plt.plot(xdata[i], ydata[i])

plt.show()