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



class waveFunction:
    def __init__(self,n,l,m,elec):
        self.n = n
        self.l = l
        self.m = m
        self.tonorm = 1.0
        self.elec = elec
        self.omega = 1.0
        self.alpha = 0.0
        self.vSign = sign(self.omega)
        self.tolerance = 5e-2
        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]])
        self.gamma0 = array([[1.0,0.0,0.0,0.0],[0.0,1.0,0.0,0.0],[0.0,0.0,-1.0,0.0],[0.0,0.0,0.0,-1.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 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,order=5)

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

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

    def getDensity(self,r):
        return self.tonorm**2*(conjugate(self.radialWaveFunction(r))*self.radialWaveFunction(r)).real
    
    def setOmega(self):
        self.omega = self.getProbCurrent(1.0)[1]/self.getDensity(1.0)
        self.vSign = sign(self.omega)

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

    def returnSpinor(self,r,theta):
        waveFunction = self.returnWaveFunction(r,theta)
        
        elec1coeff = float((self.elec+1)%2)
        elec2coeff = float(self.elec)

        alpha = self.vSign

        spin = -alpha*self.gamma1*sin(theta)+alpha*self.gamma2*cos(theta)
        electrons = 1.0/sqrt(2)*array([elec1coeff,elec2coeff,elec1coeff,elec2coeff])
        coeffs = matrixmultiply(spin,electrons)

        return array([coeffs[0]*waveFunction,coeffs[1]*waveFunction,coeffs[2]*waveFunction,coeffs[3]*waveFunction])


    def getTransformation(self,r,theta):
        speedOverc = self.getSpeed(r)/self.c
        nVector = -self.vSign*vector(-sin(theta),cos(theta),0.0)
        
        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,r,theta):
        return matrixmultiply(self.getTransformation(r,theta),self.returnSpinor(r,theta))

    def getTransformedSpinorRev(self,theta,r):
        return matrixmultiply(self.getTransformation(r,theta),self.returnSpinor(r,theta))

    def returnDifferentialSpinor(self,r,theta):
        spinorRadial = self.getTransformedSpinor
        spinorTheta = self.getTransformedSpinorRev

        radialDifferential = -derivative(spinorRadial,r,n=1,order=5,args=(theta,))*r
        thetaDifferential = -derivative(spinorTheta,theta,n=1,order=5,args=(r,))

        xDifferential = radialDifferential*cos(theta)-thetaDifferential*sin(theta)
        yDifferential = thetaDifferential*cos(theta)+radialDifferential*sin(theta)
       
        return matrixmultiply(self.gamma1,xDifferential)+matrixmultiply(self.gamma2,yDifferential)


class integrations:
    def __init__(self,wavefunctionOne,wavefunctionTwo):
        self.wavefunctionOne = wavefunctionOne
        self.wavefunctionTwo = wavefunctionTwo

        self.omegaOne = self.wavefunctionOne.omega
        self.omegaTwo = self.wavefunctionTwo.omega

        self.maxIterations = 50
        self.tolerance = 5e-2
        self.maxSize = 50.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]])
        self.gamma0 = array([[1.0,0.0,0.0,0.0],[0.0,1.0,0.0,0.0],[0.0,0.0,-1.0,0.0],[0.0,0.0,0.0,-1.0]])

    def integrateKineticTheta(self,r):
        return integrate.quad(lambda theta: (1.0j*dot(matrixmultiply(conjugate(self.wavefunctionOne.getTransformedSpinor(r,theta)),self.gamma0),self.wavefunctionTwo.returnDifferentialSpinor(r,theta))).real, 0, 2*pi, epsrel=self.tolerance, limit=self.maxIterations)[0]
 
    def integrateKineticRadial(self):
        return integrate.quad(lambda r: self.integrateKineticTheta(r), 0, self.maxSize, epsrel=self.tolerance, limit=self.maxIterations)[0]
        
    def integrateKinetic(self):
        return self.integrateKineticRadial()


class basis:
    def __init__(self,maxn):
        self.maxn = maxn
        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)) or l-m == 0:
                        pass
                    else:
                        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-m<0:
                        elec = 1
                    if l-m>0:
                        elec = 0

                    if l-1 == abs(int(l-m)) or m-l == 0:
                        print "Skipped mode:", n+1,l,l-m
                    else:
                        print "Radial QN", n+1
                        print "Angular QN", l
                        print "Magnetic QN", l-m
                        print "Elec index", elec
                        self.basis[i] = waveFunction(n+1,l,l-m,elec)
                        self.basis[i].normalize()
                        self.basis[i].setOmega()
                        print "Normalization was:", self.basis[i].tonorm, "with omega:", self.basis[i].omega
                        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.KEoverlaps = ones((self.maxNbasis,self.maxNbasis),Float)

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

    def short(self):
        print "0 0 kinetic overlap is:",self.integrations[0][0].integrateKinetic()

maxn = 3

Basis = basis(maxn)

Basis.short()
