#! /bin/env python
#
#
#  Flavor 2:  using  python numpy arrays, simple way

import sys,os,math
import numpy as np

def read_table(file):
    """ return the regular part of a table as a matrix"""
    f = open(file)              # open file
    lines = f.readlines()       # get list of all lines
    f.close()                   # close
    t = []                      # prepare empty list of rows
    nc = 0                      # number of columns
    for line in lines:
        if line[0] == '#': continue      # skip comments
        words = line.split()
        if len(words) == 0: continue     # skip blank lines
        if nc == 0: nc = len(words)      # remeber # columns
        row=[]                           # empty row
        for w in words:
            row.append(float(w))         # fill the row
        t.append(np.array(row))          # add row to table
    return t

class nbody:
    def __init__(self,file):
        f = open(file)
        lines = f.readlines()
        f.close()
        self.n = 0
        self.t = 0.0
        m=[]
        x=[]
        vx=[]
        ax=[]
        y=[]
        vy=[]
        ay=[]
        nw = 0
        for line in lines:
            if line[0] == '#': continue            # skip comment lines
            words = line.split()
            if len(words) == 0: continue           # skip blank lines
            if nw == 0: nw = len(words)
            if nw != 5: continue                   # enforce M,X,Y,VX,VY table with 5 columns
            m.append(float(words[0]))
            x.append(float(words[1]))
            y.append(float(words[2]))
            vx.append(float(words[3]))
            vy.append(float(words[4]))
            ax.append(0.0)
            ay.append(0.0)
            self.n = self.n + 1
        self.m = np.array(m)
        self.x = np.array(x)
        self.y = np.array(y)
        self.vx = np.array(vx)
        self.vy = np.array(vy)
        self.ax = np.array(ax)
        self.ay = np.array(ay)
    def force(self,eps=0.0):
        """compute new accellerations based on current position, 
        with given softening"""
        eps2 = eps*eps
        for i in range(0,self.n):
            self.ax[i] = 0.0
            self.ay[i] = 0.0
            xi=self.x[i]
            yi=self.y[i]
            for j in range(0,self.n):
                if i==j: continue
                dx = self.x[j]-xi
                dy = self.y[j]-yi
                p = 1.0/math.sqrt(dx*dx+dy*dy+eps2)
                p = self.m[j]*p*p*p;
                self.ax[i] = self.ax[i] + p*dx
                self.ay[i] = self.ay[i] + p*dy
    def euler(self,dt):                
        """take a single forward Euler step"""
        for i in range(0,self.n):
            self.x[i]  = self.x[i]  + dt*self.vx[i]
            self.y[i]  = self.y[i]  + dt*self.vy[i]
            self.vx[i] = self.vx[i] + dt*self.ax[i]
            self.vy[i] = self.vy[i] + dt*self.ay[i]
        self.t = self.t + dt
    def list(self):
        for i in range(0,self.n):
            print "%d %f  %f %f  %f %f" % (i+1,self.m[i],self.x[i],self.y[i],self.vx[i],self.vy[i])

def bench(nstep=100, tstop=1.0, eps=0.1, file='pyth2d.mpv'):
    a = nbody(file)
    step = tstop/nstep
    print "step=",step
    a.force(eps)
    a.list()
    for i in range(0,nstep):
        a.euler(step)
        a.force(eps)
    a.list()