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



class waveFunction:
    def __init__(self,n,l,m,tonorm):
        self.n = n
        self.l = l
        self.m = m
        self.tonorm = tonorm
        self.scale = 10.0

    def radialWaveFunction(self,r):
        return special.genlaguerre(self.n-self.l-1,2*self.l+1)(2.0*r/(self.n))*e**(-r/(self.n))*(2.0*r/(self.n))**self.l

    def thetaWaveFunction(self,theta):
        return special.sph_harm(self.m,self.l,theta,pi/2)

    def returnWaveFunction(self,r,theta):
        return self.tonorm*self.radialWaveFunction(r)*self.thetaWaveFunction(theta)

    def integrateThetaSeparately(self):
        return integrate.quad(lambda theta: (conjugate(self.thetaWaveFunction(theta))*self.thetaWaveFunction(theta)).real, 0, 2*pi)[0]

    def integrateRadialSeparately(self):
        return integrate.quad(lambda r: conjugate(self.radialWaveFunction(r))*r*self.radialWaveFunction(r), 0, Inf)[0]
    
    def normalize(self):
        self.tonorm = 1/sqrt(self.integrateRadialSeparately()*self.integrateThetaSeparately())

    def getProbCurrent(self,r):
        derivT = derivative(self.thetaWaveFunction,0.0,n=1)

        if r != 0.0:
            part1theta = conjugate(self.radialWaveFunction(r)*self.tonorm)*self.tonorm*derivT*self.radialWaveFunction(r)/r
            part2theta = self.radialWaveFunction(r)*self.tonorm*conjugate(self.tonorm*derivT*self.radialWaveFunction(r))/r
        else:
            part1theta = 0
            part2theta = 0

        current = vector(0.0,(1/(2j)*(part1theta-part2theta)).real,0.0)
        return current

    def setNormalization(self,tonorm):
        self.tonorm = tonorm*self.tonorm



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

        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:
                        print "Radial QN", n+1
                        print "Angular QN", l
                        print "Magnetic QN", m-l
                        count += 1
        print "Total modes:", count
        self.maxNbasis = count
        
        i = 0
        self.basis = [0]*self.maxNbasis
        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:
                        self.basis[i] = waveFunction(n+1,l,m-l,1.0)
                        self.basis[i].normalize()
                        print "Normalization was:", self.basis[i].tonorm, "for mode:", n+1,l,m-l
                        i += 1
                    
        print "Basis finished initialization:", time.ctime()

    def getFullWaveFunction(self,r,theta):
        self.fullWavefunction = 0.0
        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:
                        self.fullWavefunction += self.tonorm*self.coeffs[i]*self.basis[i].returnWaveFunction(r,theta)
                        i += 1
        return self.fullWavefunction

    def getTotalProbCurrent(self,r):
        self.current = vector(0.0,0.0,0.0)
        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:
                        self.current += self.tonorm*self.coeffs[i]*self.basis[i].getProbCurrent(r)
                        i += 1
        return self.current       

    def normalize(self):
        norm = 0.0
        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:
                        norm += (conjugate(self.coeffs[i])*self.coeffs[i])
                        i += 1
                    
        self.tonorm = 1/sqrt(norm)


       
class electron:
    def __init__(self,basis,maxn,elec):
        self.basis = basis
        self.maxn = maxn
        self.elec = elec
        self.tonorm = 1.0
        self.velScale = 1.0
        self.omega = 1.0

    def returnSpinor(self,r,theta):
        waveFunction = self.basis.getFullWaveFunction(r,theta)
        elec1coeff = float((self.elec+1)%2)
        elec2coeff = float(self.elec)
        return [elec1coeff*waveFunction, elec2coeff*waveFunction, elec1coeff*waveFunction, elec2coeff*waveFunction]

    def getDensity(self,r):
        spinor = self.returnSpinor(r,0.0)
        gamma0 = [[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]]
        spinorTrans = matrixmultiply(gamma0,spinor)
        return (dot(conjugate(spinor),spinorTrans)).real

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

    def getElecProbCurrent(self,r):
        return self.basis.getTotalProbCurrent(r)*r

    def getVelocity(self,r):
        density = self.getDensity(r)
        if density != 0.0:
            return self.velScale*vector(0.0,self.omega*r,0.0)/density
        else:
            return vector(0.0,0.0,0.0)

    def setOmega(self):
        self.omega = self.velScale*self.getElecProbCurrent(.4)[1]/self.getDensity(.4)
        print self.omega
        self.omega = self.velScale*self.getElecProbCurrent(.7)[1]/self.getDensity(.7)
        print self.omega

    def setVelScale(self,velScale):
        self.velScale = velScale
        self.setOmega()

    def getSpeed(self,r):
        return abs(r*self.omega)/(sqrt(1+r**2*self.omega**2))

    def getTransformation(self,r,theta):
        speed = self.getSpeed(r)
        vSign = sign(self.omega)
        nVector = vector(-sin(theta),vSign*cos(theta),0.0)

        phi = vSign*arctanh(speed)
        coshExp = cosh(phi/2)
        sinhCon = sinh(phi/2)

        return [[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,r,theta):
        self.spinorRelLattice = matrixmultiply(self.getTransformation(r,theta),self.returnSpinor(r,theta))
        return self.spinorRelLattice


class integrations:
    def __init__(self,electronOne,electronTwo):
        self.electronOne = electronOne
        self.electronTwo = electronTwo

        self.maxIterations = 50
        self.tolerance = 1e-2
        self.offset = self.tolerance
        self.maxSize = 15.0

    def integrateRadialForNorm(self,electron):
        return 2*pi*integrate.quad(lambda r: electron.getDensity(r)*r, 0, Inf, limit=self.maxIterations, epsrel=self.tolerance)[0]

    def normalize(self):
        elecOneDensity = self.integrateRadialForNorm(self.electronOne)
        elecTwoDensity = self.integrateRadialForNorm(self.electronOne)
                
        for i in xrange(self.electronOne.basis.maxNbasis):
            self.electronOne.basis.basis[i].setNormalization(1/sqrt(elecOneDensity))
            self.electronTwo.basis.basis[i].setNormalization(1/sqrt(elecTwoDensity))

        print "Finished normalization of densities at", time.ctime()

    def integrateRadialForMeanVel(self):
        return 2*pi*integrate.quad(lambda r: self.electronOne.getDensity(r)*self.electronOne.getSpeed(r)*r, 0, Inf, limit=self.maxIterations, epsrel=self.tolerance)[0]

    def getMeanVelocity(self):
        return self.integrateRadialForMeanVel()

    def getRelativisticDistance(self,r1,r2,theta2):
        return mag(vector(r2*cos(theta2)-r1,r2*sin(theta2)/sqrt(1-self.electronOne.getSpeed(r1)**2),0.0))



    def integrandForExchange(self,r1,r2,theta2):
        return (dot(conjugate(self.electronOne.getTransformedSpinor(r1,0.0)),self.electronOne.getTransformedSpinor(r2,theta2))*r1*r2/self.getRelativisticDistance(r1,r2,theta2)*dot(conjugate(self.electronTwo.getTransformedSpinor(r2,theta2)),self.electronTwo.getTransformedSpinor(r1,0.0))).real

    def integrateExchangeTheta2(self,r1,r2):
        return integrate.quadrature(lambda theta2: self.integrandForExchange(r1,r2,theta2), 0, 2*pi, tol=self.tolerance, maxiter=self.maxIterations, vec_func=False)[0]   

    def integrateExchangeRadial2(self,r1):
        return integrate.quadrature(lambda r2: self.integrateExchangeTheta2(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 2*pi*self.integrateExchangeRadial1()



    def integrandForCoulombic(self,r1,r2,theta2):
        return (dot(conjugate(self.electronOne.getTransformedSpinor(r1,0.0)),self.electronOne.getTransformedSpinor(r1,0.0))*r1*r2/self.getRelativisticDistance(r1,r2,theta2)*dot(conjugate(self.electronTwo.getTransformedSpinor(r2,theta2)),self.electronTwo.getTransformedSpinor(r2,theta2))).real

    def integrateCoulombicTheta2(self,r1,r2):
        return integrate.quadrature(lambda theta2: self.integrandForCoulombic(r1,r2,theta2), 0, 2*pi, tol=self.tolerance, maxiter=self.maxIterations, vec_func=False)[0]   

    def integrateCoulombicRadial2(self,r1):
        return integrate.quadrature(lambda r2: self.integrateCoulombicTheta2(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 2*pi*self.integrateCoulombicRadial1()









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)




maxn = 3

count = 0
for n in xrange(maxn):
    for l in xrange(n+1):
        for m in xrange(2*l+1):
            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 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

BasisOne = basis(maxn, coeffsOne)
BasisTwo = basis(maxn, coeffsOne)

electronOne = electron(BasisOne, maxn, 0)
electronTwo = electron(BasisTwo, maxn, 1)

electronOne.basis.normalize()
electronTwo.basis.normalize()

integrations = integrations(electronOne, electronTwo)

##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)

    electronOne.setVelScale(velScale)
    electronTwo.setVelScale(velScale)

    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()








