from visual import *
from visual.graph import *
from __future__ import division
from random import *

#body 1,  a 10g sphere
mass1 = sphere(pos=(0,0.001,0),radius=0.5,color=color.red)
mass1.m=0.01


#define gravitational field:
Fgrav=vector(0,mass1.m*(-980),0)
#initial velocity
scene.range=vector(200,200,100)

#define environment
gulv = box(pos=vector(20,0,0),length=400,width=300,height=0.001,color=color.green)
#coordinate system
xaxis=arrow(pos=(0,0,0),axis=(30,0,0),color=color.blue)
yaxis=arrow(pos=(0,0,0),axis=(0,30,0),color=color.red)
zaxis=arrow(pos=(0,0,0),axis=(0,0,30),color=color.white)
#how many events this time?
countMax=1000
statsON=1
#let's make a histogram of the results:
graphwindow=gdisplay(xtitle='range x',ytitle='N',ymax=countMax/2)
distanceX=ghistogram(bins=arange(1,150,5),color=color.red,accumulate=1,average=1)
dxlist=[]
dzlist=[]
dylist=[]

graphwindow2=gdisplay(xtitle='range X',ytitle='range Z',xmax=150,ymax=30,ymin=-30)
hits=gdots(gdisplay=graphwindow2, color=color.white)


#look at 100 events:
distXSum=0.0
distZSum=0.0

counter=0
while counter<countMax:
    #but first, try to give spread in initial velocity
    vel=vector(358./sqrt(2),358./sqrt(2),0)
    deltaV=vector(gauss(0,14./sqrt(2.)),gauss(0,22./sqrt(2.)),gauss(0,3))
    #intialize position of mass1
    mass1.pos=(0,0.001,0)
    #print deltaV
    vel= vel+deltaV
    #print vel
    #momentum is then
    mass1.p=mass1.m*vel
    mass1.trail=curve(color=color.red)
    
    #define time step
    deltat=0.01
    #start at time 0
    t=0

    while mass1.pos.y>gulv.pos.y:
        mass1.p = mass1.p + Fgrav*deltat
        mass1.pos = mass1.pos + (mass1.p/mass1.m)*deltat
        mass1.trail.append(pos=mass1.pos)
        #print counter
        t=t+deltat
        
    else:
        dxlist.append(mass1.pos.x)
        #print mass1.pos.x
        distanceX.plot(data=dxlist)
        
        distXSum = distXSum+mass1.pos.x
        distZSum = distZSum+mass1.pos.z
        hits.plot(pos=(mass1.pos.x,mass1.pos.z))
        counter=counter+1