#!/net/python/bin/python

import numpy as num
import pylab as p
from time import time

#define vector magnitude and dot product functions
dot = lambda a,b: num.sum(a*b)
mag = lambda x: num.sqrt(dot(x,x))

def rand_unit_vector(D=3,N=1):
    """
    returns N random unit vectors of dimension D
    """
    if N == 1:
        x = num.random.random(D)
    else:
        x = num.random.random([D,N])
    return (x/num.sqrt(num.sum(x**2,0)) ).T

def surface_area(L):
    """
    returns the surface area/volume of an n-d box with dimensions
    given by L
    """
    D = len(L)
    if D == 1:
        return 0
    elif D == 2:
        return 2*(L[0]+L[1])
    elif D == 3:
        return 2*(L[0]*L[1] + L[1]*L[2] + L[2]*L[0])
    else:
        print "Surface Area: not defined for %i dimensions" % D
        return 1


def plothist(x,y,fcolor='#CCCCFF'):
    """
    plots a histogram of values in x and y
    """
    binsize = x[1]-x[0]
    x = sum([[xi,xi] for xi in x],[]) + 2*[x[-1]+binsize]
    y = [0] + sum([[yi,yi] for yi in y],[]) + [0]
    p.fill(x,y,fc=fcolor)
    p.xlabel('E')
    p.ylabel('f(E)dE')

def maxwellian(E,kT,D):
    Gamma = ['nan', num.sqrt(num.pi), 1, num.sqrt(num.pi)/2, 1]
    x = 2.0*E/kT
    return 2.0/( kT * 2**(0.5*D) * Gamma[D] ) * x**(0.5*D-1) * num.exp(-0.5*x)

def findrange(kT,D,pct=1.0):
    dE = 0.001
    E=0
    integral = 0
    while True:
        integral += maxwellian(E,kT,D)*dE
        E += dE
        if integral >= 1.0:
            break
    return 1.1*E

def plot_maxwellian(kT=0.5, Emin = 0,Emax = 10.0,D=3,style="--r"):
    x = num.linspace(Emin,Emax,100)
    y = maxwellian(x,kT,D)
    p.plot(x,y,style)
    
def rms(x1,x2):
    assert len(x1) == len(x2)
    sq = [(x1[i]-x2[i])**2 for i in range(len(x1))]
    msq = num.sum(sq)/len(x1)
    return num.sqrt(msq)

class Box:
    def __init__(self,N=200,D=3,L=[10.0,10.0,10.0],d=1.0,m=1.0,E=0.5):
        """
        N = number of particles
        D = number of dimensions
        L = length of each dimension
        d = diameter of each particle
        m = mass of each particle
        E = energy per particle
    
        """
        assert len(L) == D
        self.x = num.zeros([N,D])        #position vectors
        self.p = rand_unit_vector(D,N)   #momentum vectors
        self.m = 1.0*m                   #mass
        self.d = 1.0*d                   #particle diameter

        
        self.L = 1.0*num.array(L)        #box dimensions
        self.P = 0                       #Pressure
        self.D = D                       #spacial dimensions
        self.N = N                       #number of particles
        self.E = E                       #energy per particle
        self.t = 0.0                     #cumulative time
        self.SA = surface_area(L)        #surface area
        self.n = self.N/num.prod(self.L) #number density
        

        for i in range(self.D):
            self.x[:,i] = L[i] * num.random.random(N)

        #All particles start with same initial energy
        #  e = p^2/2m
        #  E = Ne
        self.p *= num.sqrt(2.0*m*E)

    def V(self):
        return 1.0*num.prod(self.L)

    def dt(self):
        """
        returns the dt value which will make the fastest
        particle travel 0.4*d units
        """
        v_max = num.max( num.sqrt( num.sum(self.p**2,1) ) ) / self.m
        x_max = 0.4*self.d
        return x_max/v_max

    def particle_energy(self):
        """
        returns an array of energies of each particle
        in the system
        """
        return num.sum(self.p**2,1)/(2*self.m)

    def total_energy(self):
        """
        computes the total energy of the system
        """
        return num.sum( self.particle_energy() )

    def timestep(self,disp=0):
        """
        progresses system by one step, taking into account
        collisions between particles and with the wall

        updates self.P and self.t
        """
        col = 0 #counter for number of collisions
        wcol = 0 #counter for number of wall collisions

        dp = 0 #counts the net momentum imparted on the walls of
              #the box by all wall collisions
        
        dt = self.dt()
        self.t += dt
        
        self.x += self.p * dt/self.m
        
        for i in range(self.N):
            #check for wall collisions
            for j in range(self.D):
                if self.x[i,j] > self.L[j] and self.p[i,j]>0:
                    wcol += 1
                    self.x[i,j] = 2*self.L[j]-self.x[i,j]
                    self.p[i,j] = - self.p[i,j]
                    dp += 2*abs(self.p[i,j])
                if self.x[i,j] < self.d/2 and self.p[i,j]<0:
                    wcol += 1
                    self.x[i,j] = -self.x[i,j]
                    self.p[i,j] = -self.p[i,j]
                    dp += 2*abs(self.p[i,j])
            #check for 2-particle collisions
            for j in range(i):
                #R is the relative position
                R = self.x[i] - self.x[j]
                if mag(R) > self.d: continue
                
                #V is the relative velocity
                #V = self.p[i]/self.m[i] - self.p[j]/self.m[j]
                V = (self.p[i] - self.p[j])/self.m
                if dot(R,V) <= 0:continue

                col += 1

                #c is the velocity of the center of mass
                #c = (self.p[i] + self.p[j])/(self.m[i]+self.m[j])
                c = (self.p[i] + self.p[j])/(2*self.m)

                #mu is the reduced mass
                #mu = self.m[i]*self.m[j]/(self.m[i]+self.m[j])
                mu = self.m/2

                #Collision reflects V over R
                V = 2*R*dot(V,R)/(mag(R)**2) - V

                #recompute momenta
                #self.p[i] = self.m[i]*c + mu*V
                #self.p[j] = self.m[j]*c - mu*V
                self.p[i] = self.m*c + mu*V
                self.p[j] = self.m*c - mu*V

        #update pressure
        self.P = ( (self.t-dt)*self.P + dp/self.SA )/self.t
        return col,wcol

    def EnergyDistribution(self,numbins=None,binsize=None):
        """
        returns a distribution of particle energies,
        broken up into bins (to be used in a histogram)
        """
        E = self.particle_energy()
        max_E = max(E)
        min_E = min(E)
        
        if not numbins:
            if not binsize:
                numbins = self.N/10
                binsize = (max_E-min_E)/numbins
            else:
                numbins = num.ceil( (max_E-min_E)/binsize )
        else:
            if binsize:
                min_E = 0
                max_E = numbins*binsize
            else:
                binsize = (max_E-min_E)/numbins

        x = [min_E + i*binsize for i in range(numbins)]
        y = [0 for xi in x]
        for Ei in E:
            y[min( int((Ei-min_E)/binsize), numbins-1 )] += 1

        y = [yi*1.0/(binsize*self.N) for yi in y]

        return x,y
        
        
class stack:
    def __init__(self,size):
        self.size = size
        self.L = []
    def add(self,item):
        if len(self.L) < self.size:
            self.L.append(item)
        else:
            for i in range(self.size-1):
                self.L[i] = self.L[i+1]
            self.L[-1] = item
    def tolist(self):
        return self.L
    def __getitem__(self,i):
        return self.L[i]
    def __setitem__(self,i,item):
        self.L[i] = item
    def mean(self):
        return sum(self.L)/len(self.L)

class mean_var:
    def __init__(self):
        self.reset()
        
    def add(self,item):
        if self.N == 0:
            self.val = item
        else:
            self.val = (self.N*self.val+item)/(self.N+1)
            self.N += 1
            
    def reset(self):
        self.N = 0
        self.val = 0


def VisualTest(N=100,d=0.3,kT=0.7,L=[2.0,2.0],m=1.0,tmax=3.0,numbins=50):
    D = len(L)
    dE = findrange(kT,D) * 1.0/numbins        #binsize for plotting
    
    B = Box(N=N,D=D,L=L,d=d,m=m,E=kT)

    p.plot([],[])
    
    y_stack = stack(10)
    #P_list = []
    #t_list = []
    #wcol_list = []
    while B.t < tmax:
        col,wcol = B.timestep()
        x,y = map(num.array,B.EnergyDistribution(numbins,dE))
        y_stack.add (y)
        #P_list.append(B.P)
        #t_list.append(B.t)
        #wcol_list.append(wcol)
        maxw = maxwellian(x,kT,D)
        ymean = y_stack.mean()
        p.cla()
        plothist(x,ymean)
        RMS = rms(ymean,maxw)*numbins*dE
        print B.t,col,RMS
        if RMS < 0.3: break
    
    plothist(x,ymean)
    xlim = p.xlim()
    ylim = p.ylim()
    plot_maxwellian(kT,0,numbins*dE,D=D)
    
    desc_string = 'N = %i\nBox = %s\nd = %.1f\nkT = %.1f\nm = %.1f\ntau = %.2f'%(N,'x'.join(map(str,L)),d,kT,m,B.t)
    p.text(xlim[0] + 0.7*(xlim[1]-xlim[0]), ylim[0] + 0.6*(ylim[1]-ylim[0]),desc_string)


def Pressure(N=100,d=0.3,kT=0.7,L=[2.0,2.0],m=1.0,tmax=3.0,numbins=50):
    D = len(L)
    dE = findrange(kT,D) * 1.0/numbins        #binsize for plotting
    
    B = Box(N=N,D=D,L=L,d=d,m=m,E=kT)
    
    y_stack = stack(10)
    P_list = []
    while B.t < tmax:
        col,wcol = B.timestep()
        x,y = map(num.array,B.EnergyDistribution(numbins,dE))
        y_stack.add (y)
        P_list.append(B.P)
        maxw = maxwellian(x,kT,D)
        ymean = y_stack.mean()
        RMS = rms(ymean,maxw)*numbins*dE
        print B.t,col,RMS
        if RMS < 0.3: break

    P_list = P_list[5:]
    P = num.mean(P_list)
    dP = num.sqrt(num.sum([(xi-P)**2 for xi in P_list])/len(P_list))

    V = B.V()
    n = N/V
    print 'kT =',kT
    print 'n =',n
    print 'P =', P,'+/=',dP
    print 'nkT =',n*kT
    return N,V,kT,P,dP


def FindTau(N=100,d=0.3,kT=0.7,L=[2.0,2.0],m=1.0,tmax=3.0,dE=0.1):
    D = len(L)
    numbins = N/2
    dE = findrange(kT,D) * 1.0/numbins 
    B = Box(N=N,D=D,L=L,d=d,m=m,E=kT)
    y_stack = stack(10)
    while 1:
        x,y = map(num.array,B.EnergyDistribution(numbins,dE))
        y_stack.add (y)
        col,wcol = B.timestep()
        maxw = maxwellian(x,kT,D)
        ymean = y_stack.mean()
        RMS = rms(ymean,maxw)*numbins*dE
        print B.t,col,RMS
        if RMS < 0.3: break
    
    return B.t
            

if __name__ == "__main__":
    #--------------------------------
    N = 100         #number of particles
    d = 0.2         #particle diameter
    kT = 0.5        #energy per particle
    L = [3.0,3.0]   #size and dimension of box
    m = 1.0         #particle mass
    
    tmax = 5.       #iteration time (simulation units)

    numbins = 50    #num. bins for plotting
    #--------------------------------
    LLL = []
    for kT in (0.3,0.5,0.7):
        for N in (100,120,140):
            for L in ( [2.8,2.8],[3.0,3.0],[3.2,3.2] ):
                LLL.append(Pressure(N=N,d=d,kT=kT,L=L,m=m,tmax=tmax,numbins=numbins))
    print LLL
    #VisualTest(N=N,d=d,kT=kT,L=L,m=m,tmax=tmax,numbins=numbins)
    #p.show()


    """
    mlist = [0.5+0.1*i for i in range(31)]
    tlist = []
    for m in mlist:
        print '---------------'
        print 'm = ',m
        print '---------------'
        t = []
        for i in range(7):
            t.append(FindTau(N=N,d=d,kT=kT,L=L,m=m,tmax=tmax))
        tlist.append(t)
    
    print mlist
    print tlist
    """
    

    """
    dlist = [0.2+0.01*i for i in range(31)]
    tlist = []
    for d in dlist:
        print '---------------'
        print 'd = ',d
        print '---------------'
        t = []
        for i in range(7):
            t.append(FindTau(N=N,d=d,kT=kT,L=L,m=m,tmax=tmax))
        tlist.append(t)
    
    print dlist
    print tlist
    """


    """
    Nlist = [100+10*i for i in range(11)]
    tlist = []
    for N in Nlist:
        print '---------------'
        print 'N = ',N
        print '---------------'
        t = []
        for i in range(7):
            t.append(FindTau(N=N,d=d,kT=kT,L=L,m=m,tmax=tmax))
        tlist.append(t)
    
    print Nlist
    print tlist
    """