'''
Name: Sajjan Singh Mehta
Recitation Assignment #3b
10/10/2007
PHYS 113-002
Recitation Class #2
'''

from visual import *
from random import *

# Constants
G = 6.67e-11
yr = 3.15e7
dt = 0.0002 * yr
solarmass = 1.9891e30   # Mass of the Sun to use as a base mass
au = 1.5e11             # Value of AU (astronomical units) in meters

# Define the solar bodies
star = []

star.append(sphere(m = 2 * solarmass, radius = 6.955e8 * 25))
star[0].pos = vector(0, 0, 0)
star[0].p = star[0].m * vector(0, -1.49e4, 0)
star[0].color = color.yellow
star[0].trail = curve(color = color.yellow)

star.append(sphere(m = solarmass, radius = 6.955e8 * 15))
star[1].pos = vector(au, 0, 0)
star[1].p = star[1].m * vector(0, 2.98e4, 0)
star[1].color = color.blue
star[1].trail = curve(color = color.blue)

# Randomly color the stars - for fun
'''
Removed as the colors for the objects were specifically defined
in the instructions

for i in star:
    c = (random(), random(), random())
    i.color = c
    i.trail = curve(color = c)
'''

# Calculates the gravitational force in the direction of bodyY
def gForce(bodyX, bodyY):
    r = bodyX.pos - bodyY.pos
    return r * ((-G * bodyX.m * bodyY.m) / mag(r)**3)

#scene.range = 9e11
scene.autoscale = 0

t = 0

while(t < 15 * yr):
    rate(1000)

    # Updating the positions of the stars first helps to overcome
    # the issue that results when updating all values at once
    for myStar in star:
        myStar.pos += dt * myStar.p / myStar.m
        myStar.trail.append(myStar.pos)

    # Update the net force on the star and then updates its momentum
    for myStar in star:
        myStar.F = vector(0, 0, 0)
        for otherStar in star:
            if(myStar != otherStar):
                myStar.F += gForce(myStar, otherStar)
        myStar.p += myStar.F * dt

    # Calculate the total momentum
    pTotal = vector(0, 0, 0)
    for myStar in star:
        pTotal += myStar.p
    print mag(pTotal)
    
    t += dt