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 = 40.0/(3.0/2.0)

    def radialWaveFunction(self,r):
        return e**(-r*self.scale/self.n)*(2.0*r*self.scale/self.n)**self.l

    def thetaWaveFunction(self,theta):
        return e**(-1j*self.m*theta)

    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,theta):
        derivR = derivative(self.radialWaveFunction,r,n=1)
        derivT = derivative(self.thetaWaveFunction,theta,n=1)

        part1r = conjugate(self.returnWaveFunction(r,theta))*self.tonorm*derivR*self.thetaWaveFunction(theta)
        part2r = self.returnWaveFunction(r,theta)*conjugate(self.tonorm*derivR*self.thetaWaveFunction(theta))

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

        current = vector((1/(2j)*(part1r-part2r)).real,(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):
                    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):
                    self.basis[i] = waveFunction(n+1,l,m-l,1.0)
                    self.basis[i].normalize()
                    print "Normalization was:", self.basis[i].tonorm
                    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):
                    self.fullWavefunction += self.tonorm*self.coeffs[i]*self.basis[i].returnWaveFunction(r,theta)
                    i += 1
        return self.fullWavefunction

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

    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,theta):
        spinor = self.returnSpinor(r,theta)
        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)
        self.density = (dot(conjugate(spinor),spinorTrans)).real
        return self.density

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

    def getElecProbCurrent(self,r,theta):
        self.current = self.basis.getTotalProbCurrent(r,theta)
        return self.current

    def getVelocity(self,r,theta):
        density = self.getDensity(r,theta)
        if density != 0.0:
            return self.velScale*vector(-self.getElecProbCurrent(r,theta)[1]*r*sin(theta),self.getElecProbCurrent(r,theta)[1]*r*cos(theta),0.0)/density
        else:
            return vector(0.0,0.0,0.0)

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

    def getVelocityRelLattice(self,r,theta):
        velocity = self.getVelocity(r,theta)
        velocityMag = mag(velocity)
        self.relVelocity = velocity/sqrt(1+velocityMag**2)
        return self.relVelocity

    def getTransformation(self,r,theta):
        velRelLattice = self.getVelocityRelLattice(r,theta)
        nVector = norm(velRelLattice)

        phi = sign(self.getElecProbCurrent(r,theta)[1])*arctanh(mag(velRelLattice))
        coshExp = cosh(phi/2)
        sinhCon = sinh(phi/2)

        B = [[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]]
        return B

    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 = 8
        self.center = (self.maxIterations/2.0,self.maxIterations/2.0,0.0)
        self.sizeLimit = self.maxIterations/2.0
        self.tolerance = 1
             
        print "Using resolution of:", self.maxIterations


    def integrateXSeparatelyForNorm(self,electron,y):
        return integrate.quad(lambda x: electron.getDensity(sqrt(x**2+y**2),atan2(y,x)), -self.sizeLimit, self.sizeLimit, epsrel=self.tolerance, limit=self.maxIterations)[0]

    def integrateYSeparatelyForNorm(self,electron):
        return integrate.quad(lambda y: self.integrateXSeparatelyForNorm(electron,y), -self.sizeLimit, self.sizeLimit, epsrel=self.tolerance, limit=self.maxIterations)[0]

    def normalizeDensities(self):
        elecOneDensity = self.integrateYSeparatelyForNorm(self.electronOne)
        elecTwoDensity = self.integrateYSeparatelyForNorm(self.electronTwo)

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


    def integrateXSeparatelyForMeanVelocity(self,y):
        return integrate.quad(lambda x: self.electronOne.getDensity(sqrt(x**2+y**2),atan2(y,x))*mag(self.electronOne.getVelocityRelLattice(x,y)), -self.sizeLimit, self.sizeLimit, epsrel=self.tolerance, limit=self.maxIterations)[0]

    def integrateYSeparatelyForMeanVelocity(self):
        return integrate.quad(lambda y: self.integrateXSeparatelyForMeanVelocity(y), -self.sizeLimit, self.sizeLimit, epsrel=self.tolerance, limit=self.maxIterations)[0]

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


    def getRelativisticDistance(self,x1,y1,x2,y2):
        theta1 = atan2(y1,x1)
        r1 = sqrt(x1**2+y1**2)
        
        velElectronOne = self.electronOne.getVelocityRelLattice(r1,theta1)
        positionElectronTwoRelativeToElectronOneInLattice = vector(x1,y1,0.0)-vector(x2,y2,0.0)
        gamma = 1/sqrt(1-dot(velElectronOne,velElectronOne))
        
        if mag(velElectronOne) == 0.0:
            positionElectronTwoRelativeToElectronOne = positionElectronTwoRelativeToElectronOneInLattice
        else:
            positionElectronTwoRelativeToElectronOne = positionElectronTwoRelativeToElectronOneInLattice+(gamma-1)/dot(velElectronOne,velElectronOne)*dot(velElectronOne,positionElectronTwoRelativeToElectronOneInLattice)*velElectronOne
        dist = mag(positionElectronTwoRelativeToElectronOne)

        return dist


    def integrateExchangeY2(self, x1, y1, x2):
        theta1 = atan2(y1,x1)
        r1 = sqrt(x1**2+y1**2)

        return integrate.quad(lambda y2: dot(conjugate(self.electronOne.getTransformedSpinor(r1,theta1)),self.electronOne.getTransformedSpinor(sqrt(x2**2+y2**2),atan2(y2,x2))).real/self.getRelativisticDistance(x1,y1,x2,y2)*dot(conjugate(self.electronTwo.getTransformedSpinor(sqrt(x2**2+y2**2),atan2(y2,x2))),self.electronTwo.getTransformedSpinor(r1,theta1)).real, -self.sizeLimit, self.sizeLimit, points=[y1], epsrel=self.tolerance, limit=self.maxIterations)[0]

    def integrateExchangeX2(self, x1, y1):
        return integrate.quad(lambda x2: self.integrateExchangeY2(x1, y1, x2), -self.sizeLimit, self.sizeLimit, points=[x1], epsrel=self.tolerance, limit=self.maxIterations)[0]
    
    def integrateExchangeY1(self, x1):
        integ = integrate.quad(lambda y1: self.integrateExchangeX2(x1, y1), -self.sizeLimit, self.sizeLimit, epsrel=self.tolerance, limit=self.maxIterations)
        print integ, "in loop Y1 at", time.ctime()
        return integ[0]

    def integrateExchangeX1(self):
        return integrate.quad(lambda x1: self.integrateExchangeY1(x1), -self.sizeLimit, self.sizeLimit, epsrel=self.tolerance, limit=self.maxIterations)[0]

    def integrateExchange(self):
        return self.integrateExchangeX1()


    def integrateCoulombicY2(self, x1, y1, x2):
        theta1 = atan2(y1,x1)
        r1 = sqrt(x1**2+y1**2)

        return integrate.quad(lambda y2: dot(conjugate(self.electronOne.getTransformedSpinor(r1,theta1)),self.electronOne.getTransformedSpinor(r1,theta1)).real/self.getRelativisticDistance(x1,y1,x2,y2)*dot(conjugate(self.electronTwo.getTransformedSpinor(sqrt(x2**2+y2**2),atan2(y2,x2))),self.electronTwo.getTransformedSpinor(sqrt(x2**2+y2**2),atan2(y2,x2))).real, -self.sizeLimit, self.sizeLimit, points=[y1], epsrel=self.tolerance, limit=self.maxIterations)[0]

    def integrateCoulombicX2(self, x1, y1):
        return integrate.quad(lambda x2: self.integrateCoulombicY2(x1, y1, x2), -self.sizeLimit, self.sizeLimit, points=[x1], epsrel=self.tolerance, limit=self.maxIterations)[0]
    
    def integrateCoulombicY1(self, x1):
        integ = integrate.quad(lambda y1: self.integrateCoulombicX2(x1, y1), -self.sizeLimit, self.sizeLimit, epsrel=self.tolerance, limit=self.maxIterations)
        print integ, "in loop Y1 at", time.ctime()
        return integ[0]

    def integrateCoulombicX1(self):
        return integrate.quad(lambda x1: self.integrateCoulombicY1(x1), -self.sizeLimit, self.sizeLimit, epsrel=self.tolerance, limit=self.maxIterations)[0]

    def integrateCoulombic(self):
        return self.integrateCoulombicX1()





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)
                
                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)
                
                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,5,1):
    velScale = float(v*2.5)

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

    integrations.normalizeDensities()
    
    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()
