from scipy import *
from numpy import *
import time

anaught = 1

class basisWaveFunction:
    def __init__(self,tonorm,n,l,m,s):
        self.n = n
        self.l = l
        self.m = m
        self.s = s
        self.anaught = 1.0
        self.tonorm = tonorm

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

    def sphericalWavefunction(self,theta,phi):
        return special.sph_harm(self.m,self.l,theta,phi)

    def fullWavefunction(self,r,theta,phi):
        return self.tonorm*self.radialWavefunction(r)*self.sphericalWavefunction(theta,phi)


class integrateTwoWavefunction:
    def __init__(self,basisFunction1,basisFunction2):
        self.basisFunction1 = basisFunction1
        self.basisFunction2 = basisFunction2        
        
        self.eCharge = -1.0
        self.epsilonNaught = 1.0
        self.Z = 1.0

        self.resolution = 10**-6


    def integrateSeparatelyRadial(self):
        return integrate.quad(lambda r: conjugate(self.basisFunction1.radialWavefunction(r))*r**2*self.basisFunction2.radialWavefunction(r), 0, Inf)[0]

    def integrateSeparatelyPhiTheta(self):
        if self.basisFunction1.m == self.basisFunction2.m and self.basisFunction1.l == self.basisFunction2.l:
            return integrate.dblquad(lambda theta, phi: conjugate(self.basisFunction1.sphericalWavefunction(theta,phi))*sin(phi)*self.basisFunction2.sphericalWavefunction(theta,phi), 0, pi, lambda theta: 0, lambda theta: 2*pi)[0] 
        else:
            return 0.0
    
    def integrateToNorm(self):
        return self.basisFunction1.tonorm*self.basisFunction2.tonorm*self.integrateSeparatelyPhiTheta()*self.integrateSeparatelyRadial()



    def laplacianRadial(self,r):
        return r**2*derivative(self.basisFunction2.radialWavefunction,r,dx=.03,n=2,order=5)+2*r*derivative(self.basisFunction2.radialWavefunction,r,dx=.03,n=1,order=5)

    def laplacianAngular(self,theta,phi):
        return -(self.basisFunction2.l+1)*self.basisFunction2.l*self.basisFunction2.sphericalWavefunction(theta,phi)

    def integrateSeparatelyRadialKinetic(self):
        return integrate.quad(lambda r: conjugate(self.basisFunction1.radialWavefunction(r))*self.laplacianRadial(r), 0, Inf)[0]

    def integrateSeparatelyPhiThetaKinetic(self):
        if self.basisFunction1.m == self.basisFunction2.m and self.basisFunction1.l == self.basisFunction2.l:
            return integrate.dblquad(lambda theta, phi: conjugate(self.basisFunction1.sphericalWavefunction(theta,phi))*sin(phi)*self.laplacianAngular(theta,phi), 0, pi, lambda theta: 0, lambda theta: 2*pi)[0] 
        else:
            return 0.0

    def kineticIntegral(self):
        if abs(self.integrateSeparatelyPhiTheta())<10**-12:
            return 0.0
        else:
            return -(1/2.0)*self.basisFunction1.tonorm*self.basisFunction2.tonorm*(self.integrateSeparatelyRadial()*self.integrateSeparatelyPhiThetaKinetic()+self.integrateSeparatelyPhiTheta()*self.integrateSeparatelyRadialKinetic())


    def integrateSeparatelyRadialNuclear(self):
        return self.basisFunction1.tonorm*self.basisFunction2.tonorm*integrate.quad(lambda r: conjugate(self.basisFunction1.radialWavefunction(r))*r*self.basisFunction2.radialWavefunction(r), 0, Inf)[0]

    def nuclearIntegral(self):
        if abs(self.integrateSeparatelyPhiTheta())<10**-12:
            return 0.0
        else:
            return -self.integrateSeparatelyPhiTheta()*self.integrateSeparatelyRadialNuclear()


    def integrateSeparatelyPhi12Coulombic(self,r1,r2,theta1,theta2):
        return integrate.dblquad(lambda phi1, phi2: conjugate(self.basisFunction1.fullWavefunction(r1,theta1,phi1))*self.basisFunction1.fullWavefunction(r1,theta1,phi1)*r1**2*r2**2*sin(phi1)*sin(phi2)/sqrt(r1**2+r2**2-2*r1*r2*(sin(phi1)*sin(phi2)*cos(theta2-theta1)+cos(phi1)*cos(phi2)))*conjugate(self.basisFunction2.fullWavefunction(r2,theta2,phi2))*self.basisFunction2.fullWavefunction(r2,theta2,phi2), 0, pi, lambda phi1: 0, lambda phi1: pi)[0]

    def integrateSeparatelyTheta12Coulombic(self,r1,r2):
        integ = integrate.dblquad(lambda theta1, theta2: self.integrateSeparatelyPhi12Coulombic(r1,r2,theta1,theta2), 0, 2*pi, lambda theta1: 0, lambda theta2: 2*pi)[0]
        print "Finished integrateSeparatelyTheta12Coulombic", time.ctime()
        return integ

    def integrateSeparatelyRadial2Coulombic(self,r1):
        integ = integrate.quad(lambda r2: self.integrateSeparatelyTheta12Coulombic(r1,r2), 0, r1)[0] + integrate.quad(lambda r2: self.integrateSeparatelyTheta12Coulombic(r1,r2), r1, Inf, limit=10)[0]
        print "Finished integrateSeparatelyRadial2Coulombic", time.ctime()
        return integ

    def integrateSeparatelyRadial1Coulombic(self):
        print "Began integrateSeparatelyRadial1Coulombic", time.ctime()        
        integ = integrate.quad(lambda r1: self.integrateSeparatelyRadial2Coulombic(r1), 0, Inf, limit=10)[0]
        print "Finished integrateSeparatelyRadial1Coulombic", time.ctime()
        return integ



##    def integrateFullyRadial(self,function,theta,phi):
##        return integrate.quad(lambda r: self.basisFunction1.fullWavefunction(r,theta,phi)*r**2*sin(phi)*function*conjugate(self.fullWavefunction(r,theta,phi)), 0, Inf)[0]
##
##    def integrateFullyTheta(self,function,phi):
##        return integrate.quad(lambda theta: self.integrateFullyRadial(function,theta,phi), 0, 2*pi)[0]
##       
##    def integrateFullyPhi(self,function):
##        return integrate.quad(lambda phi: self.integrateFullyTheta(function,phi), 0, pi)[0]


    


      

print '\n-----------------------------'
print "Calculation Instance Started"
print time.ctime()

maxn = 5                                                                                        # level n of max order in radial quantum number, not number of basis functions
pqN = ones(maxn)
maxL = ones(maxn)
for n in range(maxn):
    pqN[n] = n + 1
    maxL[n] = n
print "Principal Quantum Numbers:", pqN
print "Maximum Angular Momentum Quantum Numbers:", maxL

sumModesN = 0
for n in range(maxn):
    sumModesN += (n+1)**2                                                                      # sum up modes from quantum magnetic number
print "Total Number of Modes:", sumModesN*2                                                         # sum of all modes, this is the size of our matrix in its base dimension,
print ""                                                                                          # although it should be in block diagonal form.

basis = [None]*maxn
realBasis = [None]*sumModesN

nuclearIntegral = [None]*sumModesN



##normc = [[1.0], [1.0, 1.0], [1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0]]
normc = [[2.0], [0.35355339059327379, 0.20412414523193154], [0.12830005981991685, 0.045360921162651453, 0.020286020648339485], [0.062500000000000028, 0.016137430609197572, 0.004658474953124562, 0.0017607380306844094], [0.035777087639996638, 0.0073029674334022174, 0.0015936381457791913, 0.00039840953644479778, 0.00013280317881493256]]


print "\nSetting Up Basis"
i = 0
for n in range(maxn):
    basis[n]=[0]*(n+1)
    for l in range(n+1):
        basis[n][l]=[0]*(2*l+1)
        for m in range(2*l+1):
            basis[n][l][m]=[0]*2
            for s in range(2):
##                print n+1
##                print l
##                print m-l
##                print s
##                print ""
                basis[n][l][m][s] = basisWaveFunction(normc[n][l],n+1,l,m-l,s)
            realBasis[i] = basisWaveFunction(normc[n][l],n+1,l,m-l,s)
            i += 1

##print "\nRunning Normalization Check"
##
##norms = [None]*maxn
##for n in range(maxn):
##    norms[n]=[0]*(n+1)
##    for l in range(n+1):
##        print "Principal QN", n+1
##        print "Angular Momentum QN", l
##        norms[n][l] = 1/sqrt(integrateTwoWavefunction(basis[n][l][0][0],basis[n][l][0][0]).integrateToNorm())
##        print "Normalization Constant", norms[n][l]
##        print ""
##
##dump = open("dumpnorms", mode='w')
##dump.write(str(norms))
##dump.close()
##
##print "\nCalculating Nuclear Integrals"
##            
##nuclearIntegral = [None]*sumModesN
##for i in range(sumModesN):
##    nuclearIntegral[i]=[0]*sumModesN
##    for j in range(sumModesN):
##        Nint = integrateTwoWavefunction(realBasis[i],realBasis[j]).nuclearIntegral()
##        if abs(Nint)>10**-12:
##            nuclearIntegral[i][j] = Nint
##        else:
##            nuclearIntegral[i][j] = 0.0
##        print "Nuclear Integral",i,j,"is",nuclearIntegral[i][j]
##
##dump = open("dumpnuclear", mode='w')
##dump.write(str(nuclearIntegral))
##dump.close()
##
##print "\nCalculating Kinetic Integrals"
##            
##kineticIntegral = [None]*sumModesN
##for i in range(sumModesN):
##    kineticIntegral[i]=[0]*sumModesN
##    for j in range(sumModesN):
##        KEint = integrateTwoWavefunction(realBasis[i],realBasis[j]).kineticIntegral()
##        if abs(KEint)>10**-12:
##            kineticIntegral[i][j] = KEint
##        else:
##            kineticIntegral[i][j] = 0.0
##        print "Kinetic Integral",i,j,"is",kineticIntegral[i][j]
##
##dump = open("dumpkinetic", mode='w')
##dump.write(str(kineticIntegral))
##dump.close()

print "\nCalculating Coulombic Overlap Integrals"
            
coulombicIntegral = [None]*sumModesN
for i in range(sumModesN):
    coulombicIntegral[i]=[0]*sumModesN
    for j in range(sumModesN):
        COint = integrateTwoWavefunction(realBasis[i],realBasis[j]).integrateSeparatelyRadial1Coulombic()
        if abs(COint)>10**-12:
            coulombicIntegral[i][j] = COint
        else:
            coulombicIntegral[i][j] = 0.0
        print "Coulombic Overlap Integral",i,j,"is",coulombicIntegral[i][j]


            

print "\nCalculation Instance Ended"
print time.ctime()
print '-----------------------------'




##from visual.graph import *
##oGraph = gdisplay(x=200, y=400, width=600, height=400, title='Probability vs. Radial Distance', xtitle='Radial Distance', ytitle='Probability', xmax=40.0, xmin=0.0, ymax=0.8, ymin=0.00, foreground=color.black, background=color.white)    # Orbital graph window preparation
##oCurve1 = gcurve(color=color.black)
##oCurve2 = gcurve(color=color.black)
##oCurve3 = gcurve(color=color.black)
##
##tographn3l0 = basisWaveFunction(3,0)
##tographn3l1 = basisWaveFunction(3,1)
##tographn3l2 = basisWaveFunction(3,2)
##
##for x in arange(0.0, 40.0+0.01+0.001, 0.1):
##    oCurve1.plot(pos=(x,(tographn3l0.normedFullRadialWavefunctionDensity(x,3,0))))
##
##for x in arange(0.0, 40.0+0.01+0.001, 0.1):
##    oCurve2.plot(pos=(x,(tographn3l1.normedFullRadialWavefunctionDensity(x,3,1))))  
##
##for x in arange(0.0, 40.0+0.01+0.001, 0.1):
##    oCurve3.plot(pos=(x,(tographn3l2.normedFullRadialWavefunctionDensity(x,3,2))))




                      
