from visual import *

scene.autoscale = 1
G = 6.67e-11

#This program models a binary system

#Define Coordinate Planes
#xyplane = box(pos=(0,0,0), length=1e12, height=1e12, width=1,color = color.cyan)
#yzplane = box(pos=(0,0,0), length=1, height=1e12, width=1e12,color = color.cyan)
#xzplane = box(pos=(0,0,0), length=1e12, height=1, width=1e12,color = color.cyan)


#Define two objects, the earth and the sun
earth = sphere(pos = (1.5e11,0,0), radius = 4e9, color = color.red)
earth.mass = 5.97e24
earth.p = vector(0, 29.79e3, 0) * earth.mass

sun = sphere(pos = vector(0,0,0), radius = 12e9, color = color.yellow)
sun.mass = 2e30
sun.p = vector(0,-1,0) * sun.mass

asteroid = sphere (pos = (4e11,0,0), radius = 4e9, color = color.white)
asteroid.mass = 5.97e6
inclination = 0.1
asteroid.p = vector(0, cos(inclination)*sqrt(G*sun.mass/(asteroid.pos.x)), sin(inclination)*sqrt(G*sun.mass/(asteroid.pos.x))) * asteroid.mass

SunEarthVector = arrow(pos=(0,0,0), axis=(5,0,0), shaftwidth=4e8)
EarthAsteroidVector = arrow(pos=earth.pos, axis=(5,0,0), shaftwidth=4e8)
#SunAsteroidVector = arrow(pos=(0,0,0), axis=(5,0,0), shaftwidth=4e8)
LAsteroidVector = arrow(pos=(0,0,0), axis=(5,0,0), shaftwidth=4e8)
pAsteroidVector = arrow(pos=asteroid.pos, axis=(5,0,0), shaftwidth=4e8)
TempVector = arrow(pos=(0,0,0), axis=(5,0,0), shaftwidth=4e8,color=color.yellow)




print sqrt(G*sun.mass/(asteroid.pos.x))
com = sphere(pos = ((earth.mass*earth.pos+sun.mass*sun.pos)/(sun.mass+earth.mass)), radius = 2e5, color = color.blue)
#This will display the orbital trace for each of the stars
for a in [earth, sun,com,asteroid]:
  a.orbit = curve(color=a.color, radius = 2e9)

dt = 864*2
scene.autoscale = 0
while 1:
  rate(4000)
#Calculate the distance between the stars as a vector
  distES = sun.pos - earth.pos
  distAS = sun.pos - asteroid.pos
# F = Gmm rhat/r^3
  forceES = 6.7e-11 * earth.mass * sun.mass * distES / mag(distES)**3
  forceAS = 6.7e-11 * asteroid.mass * sun.mass * distAS / mag(distAS)**3
#Change the momentum of the earth and the sun
  earth.p = earth.p + forceES*dt
  sun.p = sun.p - forceES*dt
  asteroid.p = asteroid.p + forceAS*dt
#  com.pos = ((earth.mass*earth.pos+sun.mass*sun.pos)/(sun.mass+earth.mass))
#  com.orbit.append(pos=com.pos)

#update the orbit
  for a in [earth, sun,asteroid]:  
    a.pos = a.pos + a.p/a.mass * dt
    a.orbit.append(pos=a.pos)

#update the vectors
  SunEarthVector.axis = -distES
  EarthAsteroidVector.axis = asteroid.pos-earth.pos
  EarthAsteroidVector.pos = earth.pos
  TempVector.axis = SunEarthVector.axis+EarthAsteroidVector.axis
  LAsteroidVector.axis = cross(-distAS,asteroid.p)/1e11
#  SunAsteroidVector.axis = -distAS
  pAsteroidVector.axis = asteroid.p
  pAsteroidVector.pos = asteroid.pos