#Simulation of a comet
#
#Radii are exaggerated for clarity; momentum and force vectors are shown in green and red, respectively.
#
#This assumptions of this model are:
#		1. the change in velocity of the Sun is negligible, so it is almost stationary;
#		2. the net force on the comet is the gravitational force of the Sun on the comet
#
#Data is printed to a file


from visual import *

#set the scene so vpython won't autoscale as we watch the animation
#set the range to be just bigger than the radius of Earth's orbit
#scene.autoscale=0
#scene.range=1.8e11

#create the sun and earth objects; set the position of the earth to the right of the sun
sun=sphere(radius=7e9, pos=(0,0,0), color=color.yellow)
comet=sphere(radius=6.4e9, pos=(0.338*1.5e11,0,0), color=color.blue)

#set the mass of the earth and sun
sun.m=2e30
comet.m=6e20

#define the value of the gravitational constant
G=6.7e-11

#set the initial velocity of the sun to move upward with the right value so that it will have a circular orbit
#calculate the momentum of the comet
#note that velocity is a vector and momentum is a vector
comet.v=1.2*vector(0,5.141e4,0)
comet.p=comet.m*comet.v

#make a trail so you can see the comet's orbit
trail=curve(color=comet.color)

#define the force and momentum arrow and a scale so that they won't be too big or too small on the screen
fscale=2e-12
Farrow=arrow(color=color.red)
pscale=3e-19
Parrow=arrow(color=color.green)

#start at t=0
#set dt to be the seconds in one day
t=0
dt=3600*3

#open a file for writing data to
file = open("comet-data.txt", "w");

#if you want to go for one orbit, change 1 to 1*dt, or 2dt for two orbits, etc.
while 1:
	#slow down the animation
	rate(100)
	
	#calculate the distance between the comet and sun
	r=sqrt((sun.x-comet.x)**2+(sun.y-comet.y)**2 + (sun.z-comet.z)**2)

	#calculate the unit vector pointing from the comet to the sun
	rhat=(sun.pos-comet.pos)/r
	
	#calculate the gravitational force on the comet
	Fsp=(G*comet.m*sun.m/r**2)*rhat

	#calculate the final momentum of the comet after the time interval dt
	#this is the MOMENTUM PRINCIPLE
	comet.p=comet.p+Fsp*dt

	#calculate the final position of comet after the time interval dt
	#this is from the definition of velocity
	comet.pos=comet.pos+(comet.p/comet.m)*dt
	
	#add a point to the curve that is the trail	
	trail.append(pos=comet.pos)
	
	#calculate the position and axis of the force arrow
	Farrow.pos=comet.pos
	Farrow.axis=fscale*Fsp

	#calculate the position and axis of the momentum arrow
	Parrow.pos=comet.pos
	Parrow.axis=pscale*comet.p
	
	#calculate the time since the animation began
	#this is only needed if you want to stop the animation at a certain time t
	t=t+dt
	
	comet.K=0.5*comet.m*mag(comet.p/comet.m)**2;
	
	print >> file, comet.pos.x, mag(comet.p/comet.m), comet.K

file.close()