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/(40.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 = 40
        self.center = (self.maxIterations/2.0,self.maxIterations/2.0,0.0)
        self.sizeLimit = self.maxIterations/2.0
        self.tolerance = 1e-6

        self.electronOneSamples = [0]*self.maxIterations
        self.electronTwoSamples = [0]*self.maxIterations
        for x in xrange(self.maxIterations):
            self.electronOneSamples[x] = [0]*self.maxIterations
            self.electronTwoSamples[x] = [0]*self.maxIterations
            for y in xrange(self.maxIterations):
                theta = atan2(y-self.center[1],x-self.center[0])
                r = sqrt((x-self.center[0])**2+(y-self.center[1])**2)
                self.electronOneSamples[x][y] = electronOne.getTransformedSpinor(r,theta)
                self.electronTwoSamples[x][y] = electronTwo.getTransformedSpinor(r,theta)

        self.distances = ones((self.maxIterations,self.maxIterations,self.maxIterations,self.maxIterations))
               
        print "Finished electron samples with resolution of:", self.maxIterations, "at", time.ctime()

    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 setupSamples(self):
        for x in xrange(self.maxIterations):
            for y in xrange(self.maxIterations):
                theta = atan2(y-self.center[1],x-self.center[0])
                r = sqrt((x-self.center[0])**2+(y-self.center[1])**2)
                self.electronOneSamples[x][y] = electronOne.getTransformedSpinor(r,theta)
                self.electronTwoSamples[x][y] = electronTwo.getTransformedSpinor(r,theta)

        print "Samples set up at", time.ctime()

    def getMeanVelocity(self):
        meanVel = 0.0
        for x in xrange(int(self.maxIterations/2.0)+1.0):
            meanVel += electronOne.getDensity(x,0)*mag(electronOne.getVelocityRelLattice(x,int(self.maxIterations/2.0)))
            
        return meanVel

    def getRelativisticDistance(self,x1,y1,x2,y2):
        theta1 = atan2(y1-self.center[1],x1-self.center[0])
        r1 = sqrt((x1-self.center[0])**2+(y1-self.center[1])**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 setupDistances(self):
        for x1 in xrange(self.maxIterations):
            for y1 in xrange(self.maxIterations):
                for x2 in xrange(x1+1):
                    for y2 in xrange(y1+1):
                        self.distances[x1][y1][x2][y2] = self.getRelativisticDistance(x1,y1,x2,y2)
                        self.distances[x2][y2][x1][y1] = self.distances[x1][y1][x2][y2]
                
        print "Distances set up at", time.ctime()
                        
    def integrateCoulombic(self):
        self.coulombic = 0.0
        for x1 in xrange(self.maxIterations):
            for x2 in xrange(self.maxIterations):
                for y1 in xrange(self.maxIterations):
                    for y2 in xrange(self.maxIterations):
                        distance = self.distances[x1][y1][x2][y2]
                        if distance != 0.0:
                            oneOverDist = 1/distance
                        else:
                            oneOverDist = 1.0
                        part1 = dot(conjugate(self.electronOneSamples[x1][y1]),self.electronOneSamples[x1][y1])
                        part2 = dot(conjugate(self.electronTwoSamples[x2][y2]),self.electronTwoSamples[x2][y2])
                        self.coulombic += part1*part2*oneOverDist
        return self.coulombic.real

        
    def integrateExchange(self):
        self.exchange = 0.0
        for x1 in xrange(self.maxIterations):
            for x2 in xrange(self.maxIterations):
                for y1 in xrange(self.maxIterations):
                    for y2 in xrange(self.maxIterations):
                        distance = self.distances[x1][y1][x2][y2]
                        if distance != 0.0:
                            oneOverDist = 1/distance
                        else:
                            oneOverDist = 1.0
                        part1 = dot(conjugate(self.electronOneSamples[x1][y1]),self.electronOneSamples[x2][y2])
                        part2 = dot(conjugate(self.electronTwoSamples[x2][y2]),self.electronTwoSamples[x1][y1])
                        self.exchange += part1*part2*oneOverDist
        return self.exchange.real



##    def integrateExchangeRadial2(self, theta1, theta2, r1):
##        integ = integrate.quadrature(lambda r2: dot(conjugate(),electronTwo.getTransformedSpinor(r2,theta2)).real*r1*r2/sqrt(r1**2+r2**2-2*r1*r2*cos(theta2-theta1))*dot(conjugate(electronOne.getTransformedSpinor(r2,theta2)),electronTwo.getTransformedSpinor(r1,theta1)).real, 0, self.sizeLimit, tol=self.tolerance, maxiter=self.maxIterations, vec_func=False)
##        print integ
##        return integ[0]
##
##    def integrateExchangeRadial1(self, theta1, theta2):
##        return integrate.quadrature(lambda r1: self.integrateExchangeRadial2(theta1, theta2, r1), 0, self.sizeLimit, tol=self.tolerance, maxiter=self.maxIterations, vec_func=False)[0]
##
##    def integrateExchangeTheta2(self, theta1):
##        return integrate.quadrature(lambda theta2: self.integrateExchangeRadial1(theta1, theta2), 0, 2*pi, tol=self.tolerance, maxiter=self.maxIterations, vec_func=False)[0]
##
##    def integrateExchangeTheta1(self):
##        return integrate.quadrature(lambda theta1: self.integrateExchangeTheta2(theta1), 0, 2*pi, tol=self.tolerance, maxiter=self.maxIterations, vec_func=False)[0]
##
##    def integrateExchange(self):
##        return self.integrateExchangeTheta1()


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,21,1):
    velScale = float(v/8.0)

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

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