from visual import *
from visual.graph import *

# mission to Mars
#  Jon Marr, 8-Jan-2015

launchvel=vector(0,0,0)

rocketmass=1e6

Einput=0.5*rocketmass*mag(launchvel)**2
print "energy input (J)=",Einput

scaler=1000 #factor that radii are increased by to make bodies visible
scalerm=scaler # factor for moons' radii

dt=60 #(time steps of 1 minute)

dot=sphere(pos=(5e11,5e11,0),radius=1,color=color.black)
dot2=sphere(pos=(-5e11,-5e11,0),radius=1,color=color.black)

G=6.67e-11 # Newton's Grav. constant in mks units  

# masses
sunm=2e30 # mass of Sun
earthm=5.97e24 # mass of Earth in kg
venusm=0.8150*earthm
marsm=0.1074*earthm

#distances
earthd=1.5e11 # earth's orbital radius in m
venusd=0.7233*earthd
marsd=1.5237*earthd

#initial orbital speeds
vvenus=sqrt(G*sunm/venusd)
vearth=sqrt(G*sunm/earthd)
vmars=sqrt(G*sunm/marsd)

#radii
sunr=7e8
earthr=6.37e6
venusr=0.949*earthr
marsr=0.533*earthr

sun = sphere(pos=(0,0,0), radius=sunr*scaler/50, color=color.yellow)
venus = sphere(pos=(0,venusd,0),radius=venusr*scaler, color=color.green)
earth = sphere(pos=(-earthd,0,0), radius=earthr*scaler, color=color.blue)
mars = sphere(pos=(0,-marsd,0),radius=marsr*scaler, color=color.red)

yaxis = box(pos=(0,0,0), length = 1e9, width = 1e12, height = 1e9, color=color.green)

#speeds
sun.v = vector(0,0,0)
venus.v = vector(-vvenus,0,0)
earth.v = vector(0,-vearth,0)
mars.v = vector(vmars,0,0)

#spaceship, starting on Earth
ship = sphere(pos=earth.pos-(0,1.001*earthr,0), radius=marsr*scaler, color=color.white)
ship.v = earth.v+launchvel

sun.m=sunm
venus.m=venusm
earth.m=earthm 
mars.m=marsm

t=0 # t is time in seconds
thr=0  # time in hours (to use for ease of calculating end of while loop)
while mag(ship.pos) < 2*mag(mars.pos)and mag(ship.pos-earth.pos)>=earthr and mag(ship.pos-mars.pos)>(marsr*100):

    #rate(5000)
    t=t+dt
    thr=t/3600 # time in hours = time in seconds/(60*60)

    # first calculate force between bodies
    # fes = force on Earth due to Sun
    # s=sun, v=venus, e=earth, ma=mars
    fvs= -(G*sun.m*venus.m/mag(venus.pos-sun.pos)**3.0)*(venus.pos-sun.pos)
    fsv = fvs
    fes = -(G*sun.m*earth.m/mag(earth.pos-sun.pos)**3.0)*(earth.pos-sun.pos)
    fse = -fes
    fmas = -(G*sun.m*mars.m/mag(mars.pos-sun.pos)**3.0)*(mars.pos-sun.pos)
    fsma = -fmas

    #add spaceship.  First calc accel (note, m of ship cancels, not needed)
    accshs = -(G*sun.m/mag(ship.pos-sun.pos)**3.0)*(ship.pos-sun.pos)
    accshv = -(G*venus.m/mag(ship.pos-venus.pos)**3.0)*(ship.pos-venus.pos)
    accshe = -(G*earth.m/mag(ship.pos-earth.pos)**3.0)*(ship.pos-earth.pos)
    accshma = -(G*mars.m/mag(ship.pos-mars.pos)**3.0)*(ship.pos-mars.pos)
    accshtot = accshs+accshv+accshe+accshma

    # now do v update using Net force on each body
    venus.v = venus.v+(fvs/venus.m)*dt
    earth.v = earth.v + (fes/earth.m)*dt
    mars.v = mars.v+(fmas/mars.m)*dt
    ship.v = ship.v + accshtot*dt
    
    # pos. update for each body
    venus.pos = venus.pos + venus.v*dt
    earth.pos = earth.pos + earth.v*dt
    mars.pos = mars.pos + mars.v*dt
    ship.pos = ship.pos + ship.v*dt
    
    #check if ship is close to a planet or moon...if so, shorten time step
    delse=ship.pos-earth.pos
    delsma=ship.pos-mars.pos
    delsv=ship.pos-venus.pos
    if mag(delse) < 5e7 or mag(delsma) < 5e7 or mag(delsv) < 5e7:
        dt = 1
    else:
        dt = 60
    if mag(ship.pos-earth.pos)<earthr:
        print "CRASH to EARTH"
    if mag(ship.pos-mars.pos)<=(marsr*100):
        print "REACH MARS"
    if mag(ship.pos) >= 2*mag(mars.pos):
        print "Ship lost in space."
print "time in seconds",t
print "trip time (years)=",t/3.16e7