#	NBody in N^2 complexity
import time
import random
import math
import sys
from graphics import *
import psyco
psyco.full()

class cBody:
  def __init__(self,fm,ix,iy,fx,fy):
    self.fMass=fm
    self.fXpos=ix
    self.fYpos=iy
    self.fXvel=fx
    self.fYvel=fy
     
fG=1.0
fMaxMass=5.0
fVmax=300000.0

iXmax=600000000
iYmax=600000000

 
def zPlot(oaGalaxy):
    global oScreen
    oScreen.clear();      	
    for i in oGalaxy:
      v = int((i.fMass/fMaxMass)*(255*255*255))
      c = ((v>>16)%255, (v>>8)%255, v%255)
      c=(255,0,0)
      oScreen.plot((int(i.fXpos)/1000000,int(i.fYpos)/1000000), c)
    oScreen.update()

def fForce(a, b, oGalaxy):
    dx=abs(a.fXpos-b.fXpos)
    dy=abs(a.fYpos-b.fYpos)
    r=math.sqrt((dx*dx)+(dy*dy))
    if r==0:
      a.fMass+=b.fMass
      a.fXvel+=b.fXvel
      a.fYvel+=b.fYvel
      oGalaxy.remove(b)
    else:
      a.fXvel+=((fG*a.fMass*b.fMass)/(r*r))*((b.fXpos-a.fXpos)/r)/a.fMass
      a.fYvel+=((fG*a.fMass*b.fMass)/(r*r))*((b.fYpos-a.fYpos)/r)/a.fMass

def move(oGalaxy):
    for i in oGalaxy:
        for j in oGalaxy:
            if(i!=j):
                fForce(i,j, oGalaxy)
            
    for i in oGalaxy:
        if i.fXvel>fVmax: i.fXvel=fVmax
        if i.fXvel<-fVmax: i.fXvel=-fVmax
        if i.fYvel>fVmax: i.fYvel=fVmax
        if i.fYvel<-fVmax: i.fYvel=-fVmax
        i.fXpos+=i.fXvel
        i.fYpos+=i.fYvel
        
        if not 0<i.fXpos<iXmax:
          oGalaxy.remove(i)
        else:
          if not 0<i.fYpos<iYmax: oGalaxy.remove(i)

  
iN=1000
iTimeStep=1000

if len(sys.argv)>1:
  if len(sys.argv)<>3:
    print("Parameters <bodies> <timesteps>")
    sys.exit(-1);

  iN=eval(sys.argv[1])
  iTimeStep=eval(sys.argv[2])

oScreen=graphicsScreen(iXmax/1000000,iYmax/1000000);
    
oGalaxy = []


def random_body():
    return cBody(random.random()*pow(10,random.randint(1,fMaxMass)), \
                 random.randint(1,iXmax), \
                 random.randint(1,iYmax), \
                 random.random()*3-1.5, \
                 random.random()*3-1.5)

def bigbang_body():
    fMass=random.random()*fMaxMass
    fXpos=iXmax/2+random.random() #Otherwise all will juse merge into one
    fYpos=iYmax/2+random.random()
    angle = random.random()*math.pi*2
    fXvel=math.sin(angle)*fVmax
    fYvel=math.cos(angle)*fVmax

    return cBody(fMass,fXpos,fYpos,fXvel,fYvel)

for i in range(iN):
  	  oGalaxy.append(bigbang_body())

start=time.time()
zPlot(oGalaxy);
for i in range(iTimeStep):
  print i
  if not oGalaxy: #If all bodies have left known space
    iTimeStep=i
    break
  move(oGalaxy)
  zPlot(oGalaxy);
stop=time.time()

print("Simulated "+str(iN)+" bodies for "+str(iTimeStep)+" timesteps in "+str(stop-start)+" seconds");

