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


class waveFunction:
    def __init__(self,n,m,elec):
        self.n = n
        self.m = m
        self.elec = elec

        self.c = 3e8

        self.tonorm = 1.0
        self.omega = .5
        self.vSign = +1.0
    
        self.tolerance = 1e-6
        self.maxIterations = 50

        self.gamma0 = array([[0.0,0.0,0.0,-1j],[0.0,0.0,1j,0.0],[0.0,-1j,0.0,0.0],[1j,0.0,0.0,0.0]])
        self.gamma1 = array([[1j,0.0,0.0,0.0],[0.0,-1j,0.0,0.0],[0.0,0.0,1j,0.0],[0.0,0.0,0.0,-1j]])
        self.gamma2 = array([[0.0,0.0,0.0,1j],[0.0,0.0,-1j,0.0],[0.0,-1j,0.0,0.0],[1j,0.0,0.0,0.0]])
        self.gamma3 = array([[0.0,-1j,0.0,0.0],[-1j,0.0,0.0,0.0],[0.0,0.0,0.0,-1j],[0.0,0.0,-1j,0.0]])

    def returnThetaWaveFunction(self,theta):
        return 1/sqrt(2*pi)*e**(-1j*self.m*theta)

    def returnRadialWaveFunction(self,r):
        return sqrt(2.0)/(self.n)*e**(-r**2/(2.0*self.n**2))

    def returnWaveFunction(self,r,theta):
        return self.tonorm*self.returnRadialWaveFunction(r)*self.returnThetaWaveFunction(theta)

    def returnSpinor(self,r,theta):
        waveFunction = self.returnWaveFunction(r,theta)

        if self.m != 0:
            arg = self.m*theta/2.0
            coeffs = [(-sin(arg)+1j*cos(arg)),(-sin(arg)-1j*cos(arg))]
        else:
            arg = -(self.elec-1.0/2.0)*2.0
            coeffs = [arg,arg]
        
        return 1/sqrt(4.0)*array([coeffs[0]*waveFunction,coeffs[1]*waveFunction,coeffs[0]*waveFunction,coeffs[1]*waveFunction])

    def getProbCurrent(self,r,theta):
        derivTheta = self.tonorm*derivative(self.returnThetaWaveFunction,theta,n=1,dx=self.tolerance,order=5)
        
        part1theta = conjugate(self.returnWaveFunction(r,theta))*derivTheta*self.returnRadialWaveFunction(r)
        part2theta = self.returnWaveFunction(r,theta)*conjugate(derivTheta*self.returnRadialWaveFunction(r))

        thetaCurrent = (1/(2j)*(part1theta-part2theta)).real

        return vector(-thetaCurrent*sin(self.m*theta),thetaCurrent*cos(self.m*theta),0.0)

    def getDensity(self,r,theta):
        return (dot(conjugate(self.returnSpinor(r,theta)),self.returnSpinor(r,theta))).real

    def getProbVel(self,r,theta):
        return self.getProbCurrent(r,theta)/self.getDensity(r,theta)

    def getOmega(self):
        omega = self.getProbVel(1.0,0.0)
        self.omega = mag(omega)
        self.vSign = sign(omega[1])       

    def getRelVel(self,r,theta):
        vel = self.vSign*self.omega*r*vector(-sin(theta),cos(theta),0.0)
        return vel/sqrt(1+self.omega**2*r**2/self.c**2)

    def getMeanSpeed(self):
        return mag(self.getRelVel(1.0,0.0))

    def getTransformation(self,r,theta):
        vel = self.getRelVel(r,theta)
        speedOverc = mag(vel)/self.c
        nVector = norm(vel)
        
        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 getTransformedSpinorRadialFirst(self,r,theta):
        return self.getTransformedSpinor(r,theta)

    def getTransformedSpinorThetaFirst(self,theta,r):
        return self.getTransformedSpinor(r,theta)

    def integrateTheta(self,r):
        return integrate.quad(lambda theta: dot(conjugate(self.getTransformedSpinor(r,theta)),self.getTransformedSpinor(r,theta))*r, 0, 2*pi, full_output=1, epsrel=self.tolerance, limit=self.maxIterations)[0]

    def integrateRadial(self):
        return integrate.quad(lambda r: self.integrateTheta(r), 0, Inf, full_output=1, epsrel=self.tolerance, limit=self.maxIterations)[0]

    def normalize(self):
        norm = self.tonorm*self.integrateRadial()
        self.tonorm = 1/sqrt(norm)

    def returnDifferentialSpinor(self,r,theta):
        spinorTheta = self.getTransformedSpinorThetaFirst
        spinorRadial = self.getTransformedSpinorRadialFirst

##        thetaDifferential = -derivative(spinorTheta,theta,n=1,dx=self.tolerance,order=5,args=(r,))
        thetaDifferential = [0.0,0.0,0.0,0.0]
        radialDifferential = derivative(spinorRadial,r,n=1,dx=self.tolerance,order=5,args=(theta,))
        
        self.gammaT = - self.gamma1*sin(self.m*theta) + self.gamma2*cos(self.m*theta)
        self.gammaR = self.gamma1*cos(self.m*theta) + self.gamma2*sin(self.m*theta)
        
        differentialSpinor = array(matrixmultiply(self.gammaT,thetaDifferential)+matrixmultiply(self.gammaR,radialDifferential))

##        print "differential spinor", matrixmultiply(self.gamma0,differentialSpinor)
##        print "conjugate spinor", conjugate(self.getTransformedSpinor(theta))
##        print "product", (dot(conjugate(self.getTransformedSpinor(theta)),matrixmultiply(self.gamma0,differentialSpinor))).real

        return differentialSpinor


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

        self.maxIterations = 50
        self.tolerance = 1e-6
        self.radius = 1.0
        self.c = 3e8
        
        self.gamma0 = array([[0.0,0.0,0.0,-1j],[0.0,0.0,1j,0.0],[0.0,-1j,0.0,0.0],[1j,0.0,0.0,0.0]])

    def integrateOverlapTheta(self,r):
        return integrate.quad(lambda theta: (dot(conjugate(self.spinorOne.getTransformedSpinor(r,theta)),self.spinorTwo.getTransformedSpinor(r,theta))).real*r, 0, 2*pi, epsrel=self.tolerance, limit=self.maxIterations)[0]

    def integrateOverlapRadial(self):
        return integrate.quad(lambda r: self.integrateOverlapTheta(r), 0, Inf, epsrel=self.tolerance, limit=self.maxIterations)[0]

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


    def integratePotentialTheta(self,r):
        return integrate.quad(lambda theta: (dot(conjugate(self.spinorOne.getTransformedSpinor(r,theta)),self.spinorTwo.getTransformedSpinor(r,theta))).real*r**3, 0, 2*pi, epsrel=self.tolerance, limit=self.maxIterations)[0]

    def integratePotentialRadial(self):
        return integrate.quad(lambda r: self.integratePotentialTheta(r), 0, Inf, epsrel=self.tolerance, limit=self.maxIterations)[0]

    def integratePotential(self):
        return self.integratePotentialRadial()


    def integrateKineticTheta(self,r):
        return integrate.quad(lambda theta: (dot(conjugate(self.spinorOne.getTransformedSpinor(r,theta)),matrixmultiply(self.gamma0,self.spinorTwo.returnDifferentialSpinor(r,theta)))).real*r, 0, 2*pi, epsrel=self.tolerance, limit=self.maxIterations)[0]

    def integrateKineticRadial(self):
        return integrate.quad(lambda r: self.integrateKineticTheta(r), 0, Inf, epsrel=self.tolerance, limit=self.maxIterations)[0]

    def integrateKinetic(self):
        return self.integrateKineticRadial()




 

    def getRelativisticDistance(self,theta):
        E1speed = self.spinorOne.speed
        gamma = 1/sqrt(1-E1speed**2/self.c**2)

        posE1relE2inL = self.radius*vector(cos(theta)-1.0,sin(theta)*gamma,0.0)
        dist = mag(posE1relE2inL)
        
        if dist == 0.0:
            return self.tolerance
        else:
            return dist

    def integrandForExchange(self,theta):
        return (dot(conjugate(self.spinorOne.getTransformedSpinor(0.0)),self.spinorOne.getTransformedSpinor(theta))/self.getRelativisticDistance(theta)*dot(conjugate(self.spinorTwo.getTransformedSpinor(theta)),self.spinorTwo.getTransformedSpinor(0.0))).real

    def integrateExchangeTheta(self):
        return integrate.quad(lambda theta: self.integrandForExchange(theta), 0, 2*pi, full_output=1, points=[0,2*pi], epsrel=self.tolerance, limit=self.maxIterations)[0]

    def integrateExchange(self):
        return 2*pi*self.radius**2*self.integrateExchangeTheta()
    

    def integrandForCoulombic(self,theta):
        return (dot(conjugate(self.spinorTwo.getTransformedSpinor(0.0)),self.spinorTwo.getTransformedSpinor(0.0))/self.getRelativisticDistance(theta)*dot(conjugate(self.spinorOne.getTransformedSpinor(theta)),self.spinorOne.getTransformedSpinor(theta))).real

    def integrateCoulombicTheta(self):
        return integrate.quad(lambda theta: self.integrandForCoulombic(theta), 0, 2*pi, full_output=1, points=[0,2*pi], epsrel=self.tolerance, limit=self.maxIterations)[0]   

    def integrateCoulombic(self):
        return 2*pi*self.radius**2*self.integrateCoulombicTheta()
    


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

        count = 0
        for elec in xrange(2):
            for n in xrange(self.maxn):
                for m in xrange(n*2+1):
                    count += 1

        self.maxNbasis = count
        print "Number of modes:", count
        
        i = 0
        self.basis = [0]*self.maxNbasis
        for elec in xrange(2):
            for n in xrange(self.maxn):
                for m in xrange(n*2+1):
                    print "Radial QN", n+1
                    print "Azimuthal QN", m-n
                    print "Electron Number", elec
                  
                    self.basis[i] = waveFunction(n+1,m-n,elec)
                    self.basis[i].getOmega()
                    
##                    self.basis[i].normalize()
##                    print "Normalization was:", self.basis[i].tonorm

                    
##                print "Average speed was:", self.basis[i].getMeanSpeed()
                    
                    print ""
                    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.Poverlaps = 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 fillPoverlaps(self):
        for i in xrange(self.maxNbasis):
            for j in xrange(i+1):
                Poverlap = self.integrations[i][j].integratePotential()
                if abs(Poverlap) < 1e-15:
                    Poverlap = 0.0
                self.Poverlaps[i][j] = Poverlap
                self.Poverlaps[j][i] = Poverlap
                print "Nuclear",i,j,"is:",Poverlap,"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):
        for n in xrange(self.maxNbasis/2.0):
            print "Time is:", time.ctime()
            print "Radial QN is:", self.integrations[n][n].spinorOne.n
            print "Azimuthal QN is:", self.integrations[n][n].spinorOne.m
            print n, n, "overlap is:", self.integrations[n][n].integrateOverlap()
            print n, n, "potential overlap is:",self.integrations[n][n].integratePotential()
            print n, n, "kinetic overlap is:",self.integrations[n][n].integrateKinetic()
##            print n, n, "exchange overlap is:",self.integrations[n][n].integrateExchange()
##            print n, n, "coulombic overlap is:",self.integrations[n][n].integrateCoulombic()

maxn = 3

Basis = basis(maxn)

Basis.short()
