#!/usr/bin/python3
import turtle
from math import sqrt

class Body:
    def __init__(self):
        self.mass=1
        self.x=self.y=0
        self.vx=self.vy=0
        self.ax=self.ay=0
        self.marker=turtle.Turtle()
        self.marker.speed(0)

def summary(bodies):
    E=0
    mx=my=0
    for a in bodies:
        E+=0.5*a.mass*(a.vx**2+a.vy**2)
        mx+=a.mass*a.vx
        my+=a.mass*a.vy
        for b in bodies:
            if (a==b): continue
            r=sqrt((a.x-b.x)**2+(a.y-b.y)**2)
            E+=-0.5*a.mass*b.mass/r
    print("Total energy:   %f"%E)
    print("Total momentum: (%f,%f)"%(mx,my))
    
bodies=[]
b=Body()
b.mass=10
b.x=-5
bodies.append(b)
b=Body()
b.x=10
b.vy=0.5
bodies.append(b)
b=Body()
b.x=11
b.vy=1.4
b.mass=0.001
bodies.append(b)

# Adjust first body's velocity for zero total momentum
mx=my=0
for i in range(1,len(bodies)):
    mx+=bodies[i].mass*bodies[i].vx
    my+=bodies[i].mass*bodies[i].vy

bodies[0].vx=-mx/bodies[0].mass
bodies[0].vy=-my/bodies[0].mass
    
summary(bodies)

dt=0.02
t=0

turtle.setworldcoordinates(-20,-20,20,20)
turtle.tracer(10,0)
for a in bodies:
    a.marker.penup()
    a.marker.goto(a.x,a.y)
    a.marker.pendown();
    
while(t<300):
    for a in bodies:
        fx=fy=0
        for b in bodies:
            if (a==b): continue
            r=sqrt((a.x-b.x)**2+(a.y-b.y)**2)
            fx+=-a.mass*b.mass*(a.x-b.x)/r**3
            fy+=-a.mass*b.mass*(a.y-b.y)/r**3
        a.vx=a.vx+0.5*(a.ax+fx/a.mass)*dt
        a.vy=a.vy+0.5*(a.ay+fy/a.mass)*dt
        a.ax=fx/a.mass
        a.ay=fy/a.mass
    for a in bodies:
        a.x=a.x+a.vx*dt+0.5*a.ax*dt*dt
        a.y=a.y+a.vy*dt+0.5*a.ay*dt*dt
        a.marker.goto(a.x,a.y)
    t=t+dt

print("Finished")
summary(bodies)
input('Press ENTER to close window')