from visual import *
from time import clock
from visual.graph import *

#Program Specifications:
#a) define and use functions
#b) define and use classes with associated methods
#c) use multi-dimensional lists and nested loops
#     or Calculate statistical averages
#d) create and use effective input and output
#e) design test cases

#UNITS:
#Temperature in K
#Mass in kg
#Distances in m
#Veloctiy in m/s
#Energy in J
#Time in s
#Momentum in kg*m/s

#Introduction.
#This program is designed to take a a system of atoms and correlate their velocity
#with their temperature as a function of time. Let there be a system, where a cube
#with volume L^3 has a thin sliding box of mass m. On either side of the partition is
#N/2 number of atoms of the same gas. The gases on one side have a different
#temperature and therefore a different velocity, which results in more collisions
#with the sliding partition than the atoms on the other side. Through the motion of
#the slider the energy from one side is transfered for one gase to the other,
#until eventually the gases have the same temperature.

print#36
print"                        _______________________________"
print"                                Atoms and Slider"
print"                        ==============================="
print"                             For N number of In a Box"
print"              By:  12-17-01. PROJECT2"#40
print#41

#Analyze the Problem.
#through analysis it was determined:
# (1/2)m (v_rms)^2=(3/2) k*T
# where the (3) is from the degrees of freedom.

q=1
while q==1:
    # A model of an ideal gas with hard-sphere collisions
    # Program uses Numeric Python arrays for high speed computations

    class wall:
        def __init__(self,pos,mass,length,width,height,velocity):
            self.pos=pos
            self.mass=mass
            self.length=length
            self.width=width
            self.height=height
            self.velocity=velocity
            self.p=mass*velocity
            self.makewall()
        #We need a Method that defines the Momentum of the Wall.
        def makewall(self):
            self.wall=box(pos=self.pos,length=self.length,height=self.height,width=self.width)
        def movewall(self):
            self.wall.pos=self.pos

    def cont():
            print"Program is done."
            run=1
            q=0
            while run==1:
                control=input("1 to continue, 0 to exit")
                if control==1:
                    q=1
                    run=0
                    return q
                elif control==0:
                    q=0
                    run=0
                    return q
                else:
                    q=0
                    run=1
                    return q
                    
    win=500

    Natoms = 50  # change this to have more or fewer atoms

    # Typical values
    L = 1. # container is a cube L on a side
    gray = (0.7,0.7,0.7) # color of edges of container
    Raxes = 0.005 # radius of lines drawn on edges of cube
    Matom = 4E-3/6E23 # helium mass
    Ratom = 0.03 # wildly exaggerated size of helium atom
    k = 1.4E-23 # Boltzmann constant
    #T = 300. # around room temperature
    dt = 1E-5

    T=input("Enter the temperature of the atoms left of the partition(Kelvin):")
    print "Initial Average KE(J):", 3/2*k*T
    scene = display(title="Gas", width=win, height=win, x=0, y=0,
                    range=L, center=(L/2.,L/2.,L/2.))

    deltav = 100. # binning for v histogram]
        
    vdist = gdisplay(x=0, y=win, ymax = Natoms*deltav/1000.,
                 width=win, height=win/2., xtitle='v', ytitle='dN')
    theory = gcurve(color=color.cyan)
    observation = ghistogram(bins=arange(0.,3000.,deltav),
                            accumulate=1, average=1, color=color.red)
    xtwin = gdisplay(x=0, y=0, width=win, height=win/2, title='Partition Position vs. t', xtitle='t(s)', ytitle='x(m)', ymax=L)
    xt=gcurve(color=color.green)
    KEtwin = gdisplay(x=0, y=0, width=win, height=win/2, title='Mean KE of Gas vs. t', xtitle='t(s)', ytitle='KE(J)')
    KEt=gcurve(color=color.red)    

    dv = 10.
    for v in arange(0.,3001.+dv,dv): # theoretical prediction
        theory.plot(pos=(v,
            (deltav/dv)*Natoms*4.*pi*((Matom/(2.*pi*k*T))**1.5)
                         *exp((-0.5*Matom*v**2)/(k*T))*v**2*dv))

    xaxis = curve(pos=[(0,0,0), (L,0,0)], color=gray, radius=Raxes)
    yaxis = curve(pos=[(0,0,0), (0,L,0)], color=gray, radius=Raxes)
    zaxis = curve(pos=[(0,0,0), (0,0,L)], color=gray, radius=Raxes)
    xaxis2 = curve(pos=[(L,L,L), (0,L,L), (0,0,L), (L,0,L)], color=gray, radius=Raxes)
    yaxis2 = curve(pos=[(L,L,L), (L,0,L), (L,0,0), (L,L,0)], color=gray, radius=Raxes)
    zaxis2 = curve(pos=[(L,L,L), (L,L,0), (0,L,0), (0,L,L)], color=gray, radius=Raxes)

    h=L
    w=L
    l=.03
    m=1e-18
    v=vector(0.,0.,0.)
        
    pos=vector(L/2,L/2,L/2)
    wall=wall(pos,m,l,h,w,v)

    Atoms = []
    colors = [color.red, color.green, color.blue,
              color.yellow, color.cyan, color.magenta]
    poslist = []
    plist = []
    mlist = []
    rlist = []

    for i in range(Natoms):
        Lmin = 1.1*Ratom
        Lmax = L-Lmin
        x = Lmin+((Lmax-Lmin)*random())/2-(l)/2-.1*Ratom
        y = Lmin+(Lmax-Lmin)*random()
        z = Lmin+(Lmax-Lmin)*random()
        r = Ratom
        Atoms = Atoms+[sphere(pos=(x,y,z), radius=r, color=colors[i % 6])]
        mass = Matom*r**3/Ratom**3
        pavg = sqrt(2.*mass*1.5*k*T) # average kinetic energy p**2/(2mass) = (3/2)kT
        theta = pi*random()
        phi = 2*pi*random()
        px = pavg*sin(theta)*cos(phi)
        py = pavg*sin(theta)*sin(phi)
        pz = pavg*cos(theta)
        poslist.append((x,y,z))
        plist.append((px,py,pz))
        mlist.append(mass)
        rlist.append(r)

    pos = array(poslist)
    p = array(plist)
    m = array(mlist)
    m.shape = (Natoms,1) # Numeric Python: (1 by Natoms) vs. (Natoms by 1)
    radius = array(rlist)

    t = 0.0
    Nsteps = 0
    pos = pos+(p/m)*(dt/2.) # initial half-step
    time = clock()

    while 1:
        observation.plot(data=mag(p/m))

        # Update all positions
        pos = pos+(p/m)*dt

        r = pos-pos[:,NewAxis] # all pairs of atom-to-atom vectors
        rmag = sqrt(add.reduce(r*r,-1)) # atom-to-atom scalar distances
        hit = less_equal(rmag,radius+radius[:,NewAxis])-identity(Natoms)
        hitlist = sort(nonzero(hit.flat)).tolist() # i,j encoded as i*Natoms+j

        # If any collisions took place:
        for ij in hitlist:
            i, j = divmod(ij,Natoms) # decode atom pair
            hitlist.remove(j*Natoms+i) # remove symmetric j,i pair from list
            ptot = p[i]+p[j]
            mi = m[i,0]
            mj = m[j,0]
            vi = p[i]/mi
            vj = p[j]/mj
            ri = Atoms[i].radius
            rj = Atoms[j].radius
            a = mag(vj-vi)**2
            if a == 0: continue # exactly same velocities
            b = 2*dot(pos[i]-pos[j],vj-vi)
            c = mag(pos[i]-pos[j])**2-(ri+rj)**2
            d = b**2-4.*a*c
            if d < 0: continue # something wrong; ignore this rare case
            deltat = (-b+sqrt(d))/(2.*a) # t-deltat is when they made contact
            pos[i] = pos[i]-(p[i]/mi)*deltat # back up to contact configuration
            pos[j] = pos[j]-(p[j]/mj)*deltat
            mtot = mi+mj
            pcmi = p[i]-ptot*mi/mtot # transform momenta to cm frame
            pcmj = p[j]-ptot*mj/mtot
            rrel = norm(pos[j]-pos[i])
            pcmi = pcmi-2*dot(pcmi,rrel)*rrel # bounce in cm frame
            pcmj = pcmj-2*dot(pcmj,rrel)*rrel
            p[i] = pcmi+ptot*mi/mtot # transform momenta back to lab frame
            p[j] = pcmj+ptot*mj/mtot
            pos[i] = pos[i]+(p[i]/mi)*deltat # move forward deltat in time
            pos[j] = pos[j]+(p[j]/mj)*deltat
     
        # Bounce off walls
        outside = less_equal(pos,Ratom) # walls closest to origin
        p1 = p*outside
        p = p-p1+abs(p1) # force p component inward
        outside = greater_equal(pos,L-Ratom) # walls farther from origin
        p1 = p*outside
        p = p-p1-abs(p1) # force p component inward

        #Make Atoms bounce off of wall.
        M99=wall.mass+Matom
        m91=Matom-wall.mass
        m92=wall.mass-Matom

        outside = greater_equal(pos[:,0],wall.pos.x-Ratom*1.1) # walls farther from origin
        pw=[]
        dpw=[]
        for q in range(Natoms):
            pw.append([wall.p.x])
            dpw.append([0., 0., 0.])
        pw=array(pw)
        dpw=array(dpw)
        dp1f = ((((m91/M99)*p[:,0])+2*(wall.mass/M99)*wall.p.x*(Matom/wall.mass))-p[:,0])*outside
        dpw = ((((m92/M99)*wall.p.x)+2*(Matom/wall.mass)*p[:,0]*(wall.mass/Matom))-pw[:,0])*outside
        #print dpw
        cow=0.
        for q in range(Natoms): # force p component inward
            cow=cow+dpw[q]
        wall.p.x=wall.p.x+cow
        #print
        #print wall.p
        #print
        #print dpw
        #print
        #print dp1f
        p[:,0] = dp1f+p[:,0]
        #We need to make an array that stores the momentum imparted to the wall.
        
        # Update positions of display objects
        KEgas=(p[0,0]**2+p[0,1]**2+p[0,2]**2)/(2*Matom)
        #print "Step:",t
        for i in range(Natoms):
            Atoms[i].pos = pos[i]
            wall.pos.x=wall.pos.x+(wall.p.x/wall.mass)*dt
            wall.movewall()
            if (i!=0):
                KEgas=KEgas+(p[i,0]**2+p[i,1]**2+p[i,2]**2)/(2*Matom)
                #print KEgas
        if wall.pos.x>=L:
            break
        xt.plot(pos=(t, wall.pos.x))
        KEt.plot(pos=(t, KEgas/Natoms))
        
        Nsteps = Nsteps+1
        t = t+dt

        if Nsteps == 50:
            print '%3.1f seconds for %d steps with %d Atoms' % (clock()-time, Nsteps, Natoms)

    print "Final Average KE(J):", KEgas/Natom
    print "Final Temperature(K):", (2/3)*(1/k)*(KEgas/Natoms)
    q=cont()
    #rate(30)