from numpy import *
from scipy import *
from visual import *
import time


class planeWave:
    def __init__(self,nX,nY):
        self.nX = nX
        self.nY = nY
       
        self.tonorm = 1.0

        self.c = 300000000
        self.m = 1.0
        self.E = sqrt(self.nX**2+self.nY**2+self.m**2)
        
        self.tolerance = 1e-6
        self.maxIterations = 50
        self.size = 10.0
        
        self.gamma0 = array([[0.0,0.0,1.0,0.0],[0.0,0.0,0.0,1.0],[1.0,0.0,0.0,0.0],[0.0,1.0,0.0,0.0]])
        self.gamma1 = array([[0.0,0.0,0.0,1.0],[0.0,0.0,1.0,0.0],[0.0,-1.0,0.0,0.0],[-1.0,0.0,0.0,0.0]])
        self.gamma2 = array([[0.0,0.0,0.0,-1.0j],[0.0,0.0,1.0j,0.0],[0.0,1.0j,0.0,0.0],[-1.0j,0.0,0.0,0.0]])
        self.gamma3 = array([[0.0,0.0,1.0,0.0],[0.0,0.0,0.0,-1.0],[-1.0,0.0,0.0,0.0],[0.0,1.0,0.0,0.0]])

    def xWaveFunction(self,x):
        return e**(-1j*self.nX*x)

    def yWaveFunction(self,y):
        return e**(-1j*self.nY*y)

    def returnWaveFunction(self,x,y):
        return self.tonorm*self.xWaveFunction(x)*self.yWaveFunction(y)

    def integrateXSeparately(self):
        return integrate.quad(lambda x: (conjugate(self.xWaveFunction(x))*self.xWaveFunction(x)).real, 0, self.size, epsrel=self.tolerance, limit=self.maxIterations)[0]

    def integrateYSeparately(self):
        return integrate.quad(lambda y: (conjugate(self.yWaveFunction(y))*self.yWaveFunction(y)).real, 0, self.size, epsrel=self.tolerance, limit=self.maxIterations)[0]
   
    def normalize(self):
        norm = self.integrateXSeparately()*self.integrateYSeparately()
        self.tonorm = 1/sqrt(norm)

    def returnSpinor(self,x,y):
        waveFunction = self.returnWaveFunction(x,y)
        coeffs = array([sqrt((self.E+self.m)/(2.0*self.m)),1/sqrt(2*self.m*(self.m+self.E))])
        return array([coeffs[0]*waveFunction,coeffs[0]*waveFunction,coeffs[1]*waveFunction,coeffs[1]*waveFunction])

    def getProbCurrent(self,x,y):
        derivX = self.tonorm*derivative(self.xWaveFunction,x,dx=self.size/self.maxIterations,n=1,order=5)
        derivY = self.tonorm*derivative(self.yWaveFunction,y,dx=self.size/self.maxIterations,n=1,order=5)
        
        part1X = conjugate(self.returnWaveFunction(x,y))*derivX*self.yWaveFunction(y)
        part2X = self.returnWaveFunction(x,y)*conjugate(derivX)*self.xWaveFunction(x)
        Xcurrent = (1/(2j)*(part1X-part2X)).real
        
        part1Y = conjugate(self.returnWaveFunction(x,y))*derivY*self.xWaveFunction(x)
        part2Y = self.returnWaveFunction(x,y)*conjugate(derivY)*self.yWaveFunction(y)
        Ycurrent = (1/(2j)*(part1Y-part2Y)).real
        
        return vector(Xcurrent,Ycurrent,0.0)

    def getDensity(self,x,y):
        return (dot(conjugate(self.returnSpinor(x,y)),self.returnSpinor(x,y))).real

    def getProbVel(self,x,y):
        return self.getProbCurrent(x,y)/self.getDensity(x,y)

    def getRelVel(self,x,y):
        vel = self.getProbVel(x,y)
        return vel*sqrt(1-mag(vel)**2/self.c**2)
 
    def getTransformation(self,x,y):
        relvel = self.getRelVel(x,y)
        speedOverc = mag(relvel)/self.c
        nVector = norm(relvel)
        
        phi = arctanh(speedOverc)
        coshExp = cosh(phi/2)
        sinhCon = sinh(phi/2)

        return array([[coshExp,0.0,-nVector[2]*sinhCon,-(nVector[0]-1j*nVector[1])*sinhCon],[0.0,coshExp,-(nVector[0]+1j*nVector[1])*sinhCon,nVector[2]*sinhCon],[-nVector[2]*sinhCon,-(nVector[0]-1j*nVector[1])*sinhCon,coshExp,0.0],[-(nVector[0]+1j*nVector[1])*sinhCon,nVector[2]*sinhCon,0.0,coshExp]])

    def getTransformedSpinor(self,x,y):
        return array(matrixmultiply(self.getTransformation(x,y),self.returnSpinor(x,y)))

    def getTransformedSpinorX(self,x,y):
        return self.getTransformedSpinor(x,y)

    def getTransformedSpinorY(self,y,x):
        return self.getTransformedSpinor(x,y)

    def returnDifferentialSpinor(self,x,y):
        spinorX = self.getTransformedSpinorX
        spinorY = self.getTransformedSpinorY

        xDifferential = derivative(spinorX,x,n=1,dx=self.size/self.maxIterations,order=5,args=(y,))
        yDifferential = derivative(spinorY,y,n=1,dx=self.size/self.maxIterations,order=5,args=(x,))

        return array(matrixmultiply(self.gamma1,xDifferential)+matrixmultiply(self.gamma2,yDifferential))


class integrations:
    def __init__(self,spinorOne,spinorTwo):
        self.spinorOne = spinorOne
        self.spinorTwo = spinorTwo

        self.maxIterations = 50
        self.tolerance = 5e-2
        self.size = 10.0

        self.center = vector(self.size/2.0,self.size/2.0,0.0)
        
        self.c = 300000000
        
        self.gamma0 = array([[0.0,0.0,1.0,0.0],[0.0,0.0,0.0,1.0],[1.0,0.0,0.0,0.0],[0.0,1.0,0.0,0.0]])

    def getDistanceNuclear(self,x,y):
        dist = mag(vector(x,y,0.0)-self.center)
        if dist != 0.0:
            return dist
        else:
            return self.tolerance

    def integrateOverlapX(self,y):
        return integrate.quad(lambda x: (dot(conjugate(self.spinorOne.getTransformedSpinor(x,y)),self.spinorTwo.getTransformedSpinor(x,y))).real, 0, self.size, epsrel=self.tolerance, limit=self.maxIterations)[0]

    def integrateOverlapY(self):
        return integrate.quad(lambda y: self.integrateOverlapX(y), 0, self.size, epsrel=self.tolerance, limit=self.maxIterations)[0]

    def integrateOverlap(self):
        return self.integrateOverlapY()


    def integrandForNuclear(self,x,y):
        result = (dot(conjugate(self.spinorOne.getTransformedSpinor(x,y)),matrixmultiply(self.gamma0,self.spinorTwo.getTransformedSpinor(x,y)))).real/self.getDistanceNuclear(x,y)
        return result

    def integrateNuclearX(self,y):
        return integrate.quad(lambda x: self.integrandForNuclear(x,y), 0, self.size, points=[self.center[0]], epsrel=self.tolerance, limit=self.maxIterations)[0]

    def integrateNuclearY(self):
        return integrate.quad(lambda y: self.integrateNuclearX(y), 0, self.size, points=[self.center[1]], epsrel=self.tolerance, limit=self.maxIterations)[0]

    def integrateNuclear(self):
        return self.integrateNuclearY()


    def integrateKineticX(self,y):
        return integrate.quad(lambda x: (1j*dot(conjugate(self.spinorOne.getTransformedSpinor(x,y)),matrixmultiply(self.gamma0,self.spinorTwo.returnDifferentialSpinor(x,y)))).real, 0, self.size, epsrel=self.tolerance, limit=self.maxIterations)[0]
       
    def integrateKineticY(self):
        return integrate.quad(lambda y: self.integrateKineticX(y), 0, self.size, epsrel=self.tolerance, limit=self.maxIterations)[0]
        
    def integrateKinetic(self):
        return self.integrateKineticY()


##    def getRelativisticDistance(self,r1,r2,theta1,theta2):
##        velElectronOne = self.spinorOne.getVelocity(r1,theta1)
##        positionElectronTwoRelativeToElectronOneInLattice = vector(r1*cos(theta1),r1*sin(theta1),0.0)-vector(r2*cos(theta2),r2*sin(theta2),0.0)
##        gamma = 1/sqrt(1-dot(velElectronOne,velElectronOne)/self.c**2)
##        
##        if mag(velElectronOne) == 0.0:
##            positionElectronTwoRelativeToElectronOne = positionElectronTwoRelativeToElectronOneInLattice
##        else:
##            positionElectronTwoRelativeToElectronOne = positionElectronTwoRelativeToElectronOneInLattice+(gamma-1)/dot(velElectronOne,velElectronOne)*dot(velElectronOne,positionElectronTwoRelativeToElectronOneInLattice)*velElectronOne
##        dist = mag(positionElectronTwoRelativeToElectronOne)
##        if dist == 0.0:
##            return self.tolerance
##        else:
##            return dist
##
##
##    def integrandForExchange(self,r1,r2,theta1,theta2):
##        return (dot(matrixmultiply(conjugate(self.spinorOne.getTransformedSpinor(r1,theta1)),self.gamma0),self.spinorOne.getTransformedSpinor(r2,theta2))*r1*r2/self.getRelativisticDistance(r1,r2,theta1,theta2)*dot(matrixmultiply(conjugate(self.spinorTwo.getTransformedSpinor(r2,theta2)),self.gamma0),self.spinorTwo.getTransformedSpinor(r1,theta1))).real
##
##    def integrateExchangeTheta2(self,r1,r2,theta1):
##        return integrate.quadrature(lambda theta2: self.integrandForExchange(r1,r2,theta1,theta2), 0, 2*pi, tol=self.tolerance, maxiter=self.maxIterations, vec_func=False)[0]   
##
##    def integrateExchangeTheta1(self,r1,r2):
##        return integrate.quadrature(lambda theta1: self.integrateExchangeTheta2(r1,r2,theta1), 0, 2*pi, tol=self.tolerance, maxiter=self.maxIterations, vec_func=False)[0]
##
##    def integrateExchangeRadial2(self,r1):
##        return integrate.quadrature(lambda r2: self.integrateExchangeTheta1(r1,r2), 0, self.maxSize, tol=self.tolerance, maxiter=self.maxIterations, vec_func=False)[0]
##    
##    def integrateExchangeRadial1(self):
##        return integrate.quadrature(lambda r1: self.integrateExchangeRadial2(r1), 0, self.maxSize, tol=self.tolerance, maxiter=self.maxIterations, vec_func=False)[0]
##
##    def integrateExchange(self):
##        return self.integrateExchangeRadial1()
##
##
##    def integrandForCoulombic(self,r1,r2,theta1,theta2):
##        return (dot(matrixmultiply(conjugate(self.spinorOne.getTransformedSpinor(r1,theta1)),self.gamma0),self.spinorOne.getTransformedSpinor(r1,theta1))*r1*r2/self.getRelativisticDistance(r1,r2,theta1,theta2)*dot(matrixmultiply(conjugate(self.spinorTwo.getTransformedSpinor(r2,theta2)),self.gamma0),self.spinorTwo.getTransformedSpinor(r2,theta2))).real
##
##    def integrateCoulombicTheta2(self,r1,r2,theta1):
##        return integrate.quadrature(lambda theta2: self.integrandForCoulombic(r1,r2,theta1,theta2), 0, 2*pi, tol=self.tolerance, maxiter=self.maxIterations, vec_func=False)[0]   
##
##    def integrateCoulombicTheta1(self,r1,r2):
##        return integrate.quadrature(lambda theta1: self.integrateCoulombicTheta2(r1,r2,theta1), 0, 2*pi, tol=self.tolerance, maxiter=self.maxIterations, vec_func=False)[0]   
##
##    def integrateCoulombicRadial2(self,r1):
##        return integrate.quadrature(lambda r2: self.integrateCoulombicTheta1(r1,r2), 0, self.maxSize, tol=self.tolerance, maxiter=self.maxIterations, vec_func=False)[0]
##    
##    def integrateCoulombicRadial1(self):
##        return integrate.quadrature(lambda r1: self.integrateCoulombicRadial2(r1), 0, self.maxSize, tol=self.tolerance, maxiter=self.maxIterations, vec_func=False)[0]
##
##    def integrateCoulombic(self):
##        return self.integrateCoulombicRadial1()




class basis:
    def __init__(self,maxn):
        self.maxn = maxn
        self.tonorm = 1

        count = 0
        for nX in xrange(maxn):
            for nY in xrange(maxn):
                count += 1
        self.maxNbasis = count
        
        print "Total modes:", self.maxNbasis
        
        i = 0
        self.basis = [0]*self.maxNbasis
        for nX in xrange(maxn):
            for nY in xrange(maxn):
                print "Mode X QN", nX
                print "Mode Y QN", nY
               
                self.basis[i] = planeWave(nX,nY)
                self.basis[i].normalize()
                
                print "Normalization was:", self.basis[i].tonorm
                i += 1

        self.integrations = [0]*self.maxNbasis
        for i in xrange(self.maxNbasis):
            self.integrations[i] = [0]*self.maxNbasis
            for j in xrange(self.maxNbasis):       
                self.integrations[i][j] = integrations(self.basis[i],self.basis[j])

        self.overlaps = ones((self.maxNbasis,self.maxNbasis),Float)
        self.Noverlaps = ones((self.maxNbasis,self.maxNbasis),Float)
        self.KEoverlaps = ones((self.maxNbasis,self.maxNbasis),Float)
        self.XCoverlaps = ones((self.maxNbasis,self.maxNbasis),Float)
        self.COoverlaps = ones((self.maxNbasis,self.maxNbasis),Float)

        print "Basis and Integrations Set Up and Normalized at:", time.ctime()

    def fillOverlaps(self):
        for i in xrange(self.maxNbasis):
            for j in xrange(i+1):
                overlap = self.integrations[i][j].integrateOverlap()
                if abs(overlap) < 1e-15:
                    overlap = 0.0
                self.overlaps[i][j] = overlap
                self.overlaps[j][i] = overlap
                print "Overlap",i,j,"is:",overlap,"at:",time.ctime()

    def fillNoverlaps(self):
        for i in xrange(self.maxNbasis):
            for j in xrange(i+1):
                Noverlap = self.integrations[i][j].integrateNuclear()
                if abs(Noverlap) < 1e-15:
                    Noverlap = 0.0
                self.Noverlaps[i][j] = Noverlap
                self.Noverlaps[j][i] = Noverlap
                print "Nuclear",i,j,"is:",Noverlap,"at:",time.ctime()

    def fillKEoverlaps(self):
        for i in xrange(self.maxNbasis):
            for j in xrange(i+1):
                KEoverlap = self.integrations[i][j].integrateKinetic()
                if abs(KEoverlap) < 1e-15:
                    KEoverlap = 0.0
                self.KEoverlaps[i][j] = KEoverlap
                self.KEoverlaps[j][i] = KEoverlap
                print "Kinetic",i,j,"is:",KEoverlap,"at:",time.ctime()

    def fillXCEnergies(self):
        for i in xrange(self.maxNbasis):
            for j in xrange(i+1):
                XCoverlap = self.integrations[i][j].integrateExchange()
                if abs(XCoverlap) < 1e-15:
                    XCoverlap = 0.0
                self.XCoverlaps[i][j] = XCoverlap
                self.XCoverlaps[j][i] = XCoverlap
                print "Exchange Overlap",i,j,"is:",XCoverlap,"at:",time.ctime()

    def fillCOEnergies(self):
        for i in xrange(self.maxNbasis):
            for j in xrange(i+1):
                COoverlap = self.integrations[i][j].integrateCoulombic()
                if abs(COoverlap) < 1e-15:
                    COoverlap = 0.0
                self.COoverlaps[i][j] = COoverlap
                self.COoverlaps[j][i] = COoverlap
                print "Coulombic Overlap",i,j,"is:",COoverlap,"at:",time.ctime()

    def short(self):
        print "0 0 overlap is:", self.integrations[0][0].integrateOverlap()
        print "0 0 nuclear overlap is:",self.integrations[0][0].integrateNuclear()
        print "0 0 kinetic overlap is:",self.integrations[0][0].integrateKinetic()
##        print "0 0 exchange overlap is:",self.integrations[0][0].integrateExchange()
##        print "0 0 coulombic overlap is:",self.integrations[0][0].integrateCoulombic()
        
        print "1 1 overlap is:", self.integrations[1][1].integrateOverlap()
        print "1 1 nuclear overlap is:",self.integrations[1][1].integrateNuclear()
        print "1 1 kinetic overlap is:",self.integrations[1][1].integrateKinetic()
        
        print "2 2 overlap is:", self.integrations[2][2].integrateOverlap()
        print "2 2 nuclear overlap is:",self.integrations[2][2].integrateNuclear()
        print "2 2 kinetic overlap is:",self.integrations[2][2].integrateKinetic()

        print "3 3 overlap is:", self.integrations[3][3].integrateOverlap()
        print "3 3 nuclear overlap is:",self.integrations[3][3].integrateNuclear()
        print "3 3 kinetic overlap is:",self.integrations[3][3].integrateKinetic()
##        print "3 3 exchange overlap is:",self.integrations[3][3].integrateExchange()
##        print "3 3 coulombic overlap is:",self.integrations[3][3].integrateCoulombic()

        print "4 4 overlap is:", self.integrations[4][4].integrateOverlap()
        print "4 4 nuclear overlap is:",self.integrations[4][4].integrateNuclear()
        print "4 4 kinetic overlap is:",self.integrations[4][4].integrateKinetic()
        

maxn = 8

Basis = basis(maxn)

##Basis.fillOverlaps()
##dump = open("dumpOverlaps", mode='w')
##dump.write(str(Basis.overlaps))
##dump.close()

##Basis.fillNoverlaps()
##dump = open("dumpNoverlaps", mode='w')
##dump.write(str(Basis.Noverlaps))
##dump.close()

##Basis.fillKEoverlaps()
##dump = open("dumpKEoverlaps", mode='w')
##dump.write(str(Basis.KEoverlaps))
##dump.close()

##Basis.fillXCEnergies()
##dump = open("dumpXCoverlaps", mode='w')
##dump.write(str(Basis.XCoverlaps))
##dump.close()
##
##Basis.fillCOEnergies()
##dump = open("dumpCOoverlaps", mode='w')
##dump.write(str(Basis.COoverlaps))
##dump.close()

Basis.short()









##class densityPoint:
##    def __init__(self):
##        self.visibility = 1
##        self.point = sphere(visible=self.visibility, radius=.001)
##
##    def setAttributes(self,position,size,color):
##        self.point.pos = position
##        self.point.radius = size
##        self.point.color = color
##
##    def toggleVisibility(self):
##        self.visibility = (1+self.visibility)%2
##
##class currentVector:
##    def __init__(self):
##        self.visibility = 1
##        self.currentArrow = arrow(visible=self.visibility)
##
##    def setAttributes(self,position,direction):
##        self.currentArrow.pos = position
##        self.currentArrow.axis = direction
##
##    def toggleVisibility(self):
##        self.visibility = (1+self.visibility)%2
##
##class visualLattice:
##    def __init__(self,electronOne,sizeOfLattice,center):
##        self.electronOne = electronOne
##        self.sizeOfLattice = sizeOfLattice
##        self.center = center
##        self.scale = sizeOfLattice
##
##        self.pointLattice=[None]*sizeOfLattice
##        self.currentLattice=[None]*sizeOfLattice
##        for x in xrange(sizeOfLattice):
##            self.pointLattice[x]=[None]*sizeOfLattice
##            self.currentLattice[x]=[None]*sizeOfLattice
##            for y in xrange(sizeOfLattice):
##                position = vector(x-sizeOfLattice/2.0,y-sizeOfLattice/2.0,0.0)
##                r = mag(position)
##                theta = atan2(position[1],position[0])
##                currentDirection = self.electronOne.getVelocityRelLattice(r,theta)
##                transformationMatrix = self.electronOne.getTransformation(r,theta)
##                
##                density = self.electronOne.getDensity(r,theta)
##                size = self.scale*density
##                phase = self.electronOne.getPhase(r,theta)
##                colorArg = phase/pi
##                color = ((1+abs(colorArg))/2.0,0.0,(1-abs(colorArg))/2.0)
##                
##                self.pointLattice[x][y] = densityPoint()
##                self.currentLattice[x][y] = currentVector()
##                self.pointLattice[x][y].setAttributes(position,size,color)
##                self.currentLattice[x][y].setAttributes(position,currentDirection)
##
##    def updateLattice(self,electronOne):
##        for x in xrange(self.sizeOfLattice):
##            for y in xrange(self.sizeOfLattice):
##                position = vector(x-sizeOfLattice/2.0,y-sizeOfLattice/2.0,0.0)
##                r = mag(position)
##                theta = atan2(position[1],position[0])
##                currentDirection = self.electronOne.getVelocityRelLattice(r,theta)
##                transformationMatrix = self.electronOne.getTransformation(r,theta)
##                
##                density = self.electronOne.getDensity(r,theta)
##                size = self.scale*density
##                phase = self.electronOne.getPhase(r,theta)
##                colorArg = phase/pi
##                color = ((1+abs(colorArg))/2.0,0.0,(1-abs(colorArg))/2.0)
##                
##                self.pointLattice[x][y].setAttributes(position,size,color)
##                self.currentLattice[x][y].setAttributes(position,currentDirection)


##count = 0
##for n in xrange(maxn):
##    for l in xrange(n+1):
##        for m in xrange(2*l+1):
##            if l-1 == abs(int(m-l)):
##                pass
##            else:
##                count += 1
##
##coeffsOne = [None]*count
##coeffsTwo = [None]*count
##i = 0
##for n in xrange(maxn):
##    for l in xrange(n+1):
##        for m in xrange(2*l+1):
##            if l-1 == abs(int(m-l)):
##                pass
##            else:
##                if m-l==-2:
##                    coeffsOne[i] = 1.0
##                else:
##                    coeffsOne[i] = 0.0
##                if m-l==2:
##                    coeffsTwo[i] = 1.0
##                else:
##                    coeffsTwo[i] = 0.0
##                i += 1
##
##print "Electron One Coefficients", coeffsOne
##print "Electron Two Coefficients", coeffsTwo


##sizeOfLattice = 24
##center = (0.0,0.0,0.0)
##
##scene = display(title='Density Plot', x=0, y=0, width=600, height=600, center=center, background=(0,0,0))
##visualLattice = visualLattice(electronOne, sizeOfLattice, center)
##visualLattice.updateLattice(electronOne)

##results = list()
##
##for v in xrange(0,21,1):
##    velScale = float(v/3000.0)
##
##    integrations.normalize()
##    
##    meanVel = integrations.getMeanVelocity()
##    xcInt = integrations.integrateExchange()
##    coInt = integrations.integrateCoulombic()
##
##    print "Mean Velocity is:", meanVel
##    print "Exchange Integral is:", xcInt
##    print "Coulombic Integral is:", coInt
##
##    results.append([meanVel,xcInt,coInt])
##
##    print "Time was:", time.ctime()
##
##dump = open("dumpAnalysis", mode='w')
##dump.write(str(results))
##dump.close()


##    def getPhase(self,r,theta):
##        return atan2(self.returnSpinor(r,theta)[0].imag,self.returnSpinor(r,theta)[0].real)


