## Model 4: Collisions and Interactive Matter
## Author: M. Haley
from __future__ import division ## treat integers as real numbers in division
from visual import *
scene.autoscale = 0 ## disable autozoom
scene.width = 800
scene.height = 800

## CONSTANTS

## set probe interaction parameters
probeRange = 2
probeStrength = 1
probeVelocity = vector(0,0,0)
## set Morse interaction parameters
EM = 2
req = 2
alpha = .5

## OBJECTS AND INITIAL VALUES

## define probe
probe = sphere(pos=(-25,0,0), radius=.3, color=color.red)
probe.m = 1000 ## probe mass
probe.p = probe.m*probeVelocity 
## define particles
nTheta = 6
nPhi = 5
nPart = (nTheta-1)*nPhi
particles = zeros(nPart, sphere)
clusterRadius = 2
mPart = 1 ## mass of each particle
## array the particles spherically
position = vector(0,0,0)
for i in range(0,nTheta-1):
    theta = pi/nTheta*(i+1)
    for j in range(0,nPhi):
        phi = 2*pi/nPhi*j
        position.x = clusterRadius*sin(theta)*cos(phi) + random.uniform(-.2,.2)
        position.y = clusterRadius*sin(theta)*sin(phi) + random.uniform(-.2,.2)
        position.z = clusterRadius*cos(theta) + random.uniform(-.2,.2)
        particles[i*nPhi+j] = sphere(pos=position, radius=.2, color=color.yellow)
        particles[i*nPhi+j].m = mPart
        particles[i*nPhi+j].p = vector(0,0,0)
positionUpdates = zeros(nPart, vector)
numEaten = 0
## define center of mass
CoM = pyramid(pos=vector(0,0,0),size=(.3,.3,.3), axis=vector(0,1,0),color=color.green)
## set up time values
deltat = .01
t = 0 ## start counting time at zero

## CALCULATIONS

## print out initial state of the system
UTotal = 0
KTotal = 0
pTotal = vector(0,0,0)
for i in range(0,nPart):
    for j in range(i+1,nPart):
        r = particles[i].pos-particles[j].pos
        rmag = mag(r)
        UMorse = EM*(1-e**(-alpha*(rmag-req)))**2-EM
        UTotal = UTotal + UMorse
    KTotal = KTotal + mag(particles[i].p)**2/(2*mPart)
    pTotal = pTotal + particles[i].p
KProbe = mag(probe.p)**2/(2*probe.m)
print 'INITIAL STATE'
print 'Cluster K_i =', KTotal
print 'Cluster U_i =', UTotal
print 'Total cluster energy =', KTotal+UTotal
print 'Probe energy =', KProbe
print 'TOTAL SYSTEM ENERGY =', KTotal+UTotal+KProbe
print 'Total cluster momentum =', pTotal
print 'Probe momentum =', probe.p
print 'Cluster plus probe momentum =', pTotal + probe.p
print '------------------------------------------'
print 'FINAL STATE...(calculating)'

## animation loop
while t < 10:
    for i in range(0,nPart):
        ## calculate the force on the ith particle due to every other particle
        Fnet = vector(0,0,0)
        for j in range(0,nPart):
            r = particles[i].pos-particles[j].pos
            rmag = mag(r)
            rhat = r/mag(r) 
            #FMorse = 
            if i!=j: Fnet = Fnet + FMorse  ## no self-interaction!
        ## calculate probe force on ith particle
        rP = particles[i].pos-probe.pos
        rPmag = mag(rP)
        rPhat = rP/rPmag
        #Fprobe = 
        Fnet = Fnet + Fprobe
        ## momentum principle for ith particle and probe
        particles[i].p = particles[i].p + Fnet*deltat
        probe.p = probe.p + (-Fprobe)*deltat ## minus sign for equal and opposite force
        ## calculate position change of ith particle
        positionUpdates[i] = particles[i].p/particles[i].m*deltat

#        ## INELASTIC COLLISION
#        ## -------------------
#        grabDistance = 2
#        if rPmag < grabDistance:
#            probe.p = probe.p + particles[i].p
#            probe.radius = probe.radius + .05
#            probe.m = probe.m +mPart
#            particles[i].pos = vector(random.randint(-100,-50),random.randint(-100,100),random.randint(-100,100))
#            particles[i].p = vector(0,0,0)
#            particles[i].color = color.black
#            numEaten = numEaten +1
#        ## -------------------
    
    ## end of ith atom loop
    ## update all atomic positions simultaneously
    for i in range(0,nPart):
        particles[i].pos = particles[i].pos + positionUpdates[i]
    ## update position of probe
    probe.pos = probe.pos + probe.p/probe.m*deltat
    
    ## CENTER OF MASS
    CoM.p = vector(0,0,0)
    ## sum the momentum of all atoms
    for i in range(0,nPart):
        CoM.p = CoM.p + particles[i].p
    #VCM = 
    CoM.pos = CoM.pos + VCM*deltat
    
    ## update time   
    t = t + deltat
    rate(100) ## animation speed
## end of time loop
    
## print out final state of the system
UTotal = 0
KTotal = 0
pTotal = vector(0,0,0)
for i in range(0,nPart):
    for j in range(i+1,nPart):
        r = particles[i].pos-particles[j].pos
        rmag = mag(r)
        UMorse = EM*(1-e**(-alpha*(rmag-req)))**2-EM
        UTotal = UTotal + UMorse
    KTotal = KTotal + mag(particles[i].p)**2/(2*mPart)
    pTotal = pTotal + particles[i].p
KProbe = mag(probe.p)**2/(2*probe.m)
print 'Cluster K_f =', KTotal
print 'Cluster U_f =', UTotal
print 'Total cluster energy =', KTotal+UTotal
print 'Probe energy =', KProbe
print 'TOTAL SYSTEM ENERGY =', KTotal+UTotal+KProbe
print 'Total cluster momentum =', pTotal
print 'Probe momentum =', probe.p
print 'Cluster plus probe momentum =', pTotal + probe.p