# Import modules
from __future__ import division
from visual import *
from visual.graph import *
from math import *

# Define constants
G = 6.67300*10**-11
mass_star_1 = 2e30 #kg
mass_star_2 = 1e30 #kg
period = 15*24*60*60 #s
radius_star_1 = 695500000
radius_star_2 = 347750000
seperation = radius_star_1*10
dt = 30

# Some calculations
central_mass = mass_star_1+mass_star_2
distance_star_1 = (mass_star_2*seperation)/(mass_star_1+mass_star_2)
distance_star_2 = (mass_star_1*seperation)/(mass_star_1+mass_star_2)
omega = (2*math.pi)/period
velocity_star_1=vector(0,-omega*distance_star_1,0)
velocity_star_2=vector(0,omega*distance_star_2,0)

# Define planets, arrows, etc
star_1 = sphere(pos=(-distance_star_1,0,0),mass=mass_star_1,radius=radius_star_1,color=color.blue, velocity=velocity_star_1)
star_2 = sphere(pos=(distance_star_2,0,0),mass=mass_star_2,radius=radius_star_2,velocity=velocity_star_2)
trail = curve(color=color.white)
trail_2 = curve(color=color.red)

star_2.p = star_2.mass*star_2.velocity
star_1.p = star_1.mass*star_1.velocity

while(1==1):
    rate(500)
    trail.append(pos=star_1.pos)
    trail_2.append(pos=star_2.pos)

    # Update the force
    distance = mag(star_1.pos-star_2.pos)
    f_gravity_dir = norm(star_1.pos-star_2.pos)
    f_gravity_mag = G*((star_2.mass*star_1.mass)/(distance**2))
    f_gravity = f_gravity_dir*f_gravity_mag

    #print f_gravity

    # Update position and momentum
    star_1.p = star_1.p + -f_gravity*dt
    star_1.pos += (star_1.p/star_1.mass)*dt
    star_2.p = star_2.p + f_gravity*dt
    star_2.pos += (star_2.p/star_2.mass)*dt