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

from visual import *

# Constants
G = 6.67e-11
yr = 3.15e7
dt = 0.002 * yr

# Define the solar system bodies
solarBodies = []

# Define the properties of the Sun
# Momentum is taken into account for this simulation and if zoomed into the
# sun, a red trail can be seen displaying the minor but present gravitational
# effects of the Earth and Jupiter.
sun = sphere(pos = vector(0, 0, 0), radius = 6.955e8 * 75, m = 2e30,
             color = color.yellow, trail = curve(color = color.red))
sun.p = sun.m * vector(0, 0, 0)
solarBodies.append(sun)

# Define the properties of the Earth
earth = sphere(pos = vector(1.5e11, 0, 0), radius = 6.371e4 * 100000, m = 6e24,
               color = color.blue, trail = curve(color = color.blue))
earth.p = earth.m * vector(0, 2.98e4, 0)
solarBodies.append(earth)

# Define the properties of Jupiter
jupiter = sphere(pos = vector(7.8e11, 0, 0), radius = 7.149e6 * 2500, m = 2e27, 
                 color = (0.95, 0.95, 0.7), trail = curve(color = (0.95, 0.95, 0.7)))
jupiter.p = jupiter.m * vector(0, 1.3e4)
solarBodies.append(jupiter)

# 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)

# Sets the scene properties
scene.range = 9e11
scene.autoscale = 0

t = 0

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

    '''
    Original code for updating the position/momentum of the solar
    bodies.  Removed to allow for consideration of all forces.
    I am keeping this commented in case it is needed for grading
    purposes.
    
    # Update Earth's momentum, position, and trail
    r = earth.pos - sun.pos
    earth.F = ((-G * sun.m * earth.m) / mag(r)**3) * r
    r = earth.pos - jupiter.pos
    earth.F += ((-G * jupiter.m * earth.m) / mag(r)**3) * r
    earth.p += earth.F * dt
    earth.pos += dt * earth.p / earth.m
    earth.trail.append(earth.pos)

    # Update Jupiter's momentum, position and trail
    r = jupiter.pos - sun.pos
    jupiter.F = ((-G * sun.m * jupiter.m) / mag(r)**3) * r
    jupiter.p += jupiter.F * dt
    jupiter.pos += dt * jupiter.p / jupiter.m
    jupiter.trail.append(jupiter.pos)
    '''

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

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

    t += dt