from scipy import *
from numpy import *
import time

maxNlinear = 4                                                              # number of gaussian basis functions
silent = 0

def stringnormsarraytoarray(lineread, maxNbasis):
    dummy2 = lineread.replace(',','')
    dummy3 = dummy2.replace(']','')
    dummy4 = dummy3.replace('[','')
    dummyFinal = dummy4.split()

    norms = [None]*maxNbasis
    for i in range(maxNbasis):
        norms[i] = float(dummyFinal[i])
            
    return norms

def stringsquarearraytoarray(lineread, sizeSquare):
    dummy2 = lineread.replace(',','')
    dummy3 = dummy2.replace(']','')
    dummy4 = dummy3.replace('[','')
    dummyFinal = dummy4.split()

    target = [None]*sizeSquaresN
    m = 0
    for i in range(sizeSquare):
        target[i] = [0]*sizeSquare
        for j in range(sizeSquare):
            if abs(float(dummyFinal[m]))>10**-12:
                target[i][j] = float(dummyFinal[m])
            else:
                target[i][j] = 0.0
            m += 1

    return target


class basisWaveFunction:
    def __init__(self,sigma,nx,ny):
        self.nx = nx
        self.ny = ny
        self.sigma = sigma

    def xWavefunction(self,x):
        return sqrt(1/(sqrt(pi)*2**self.nx*factorial(self.nx)))*sqrt(1/self.sigma)*exp(-x**2/(2.0*self.sigma**2))*special.hermite(self.nx)(x/self.sigma)

    def yWavefunction(self,y):
        return sqrt(1/(sqrt(pi)*2**self.ny*factorial(self.ny)))*sqrt(1/self.sigma)*exp(-y**2/(2.0*self.sigma**2))*special.hermite(self.ny)(y/self.sigma)

    def fullWavefunction(self,x,y):
        return self.xWavefunction(x)*self.yWavefunction(y)



class integrateTwoWavefunction:
    def __init__(self,basisFunction1,basisFunction2):
        self.basisFunction1 = basisFunction1
        self.basisFunction2 = basisFunction2
        self.lim = 7.0

    def integrateSeparatelyX(self):
        return integrate.quad(lambda x: conjugate(self.basisFunction1.xWavefunction(x))*self.basisFunction2.xWavefunction(x), -Inf, Inf)[0]

    def integrateSeparatelyY(self):
        return integrate.quad(lambda y: conjugate(self.basisFunction1.yWavefunction(y))*self.basisFunction2.yWavefunction(y), -Inf, Inf)[0]

    def integrateToNorm(self):
        return self.integrateSeparatelyX()*self.integrateSeparatelyY()



    def integrateSeparatelyXforRestoringPotential(self):
        return integrate.quad(lambda x: conjugate(self.basisFunction1.xWavefunction(x))*(x**2)*self.basisFunction2.xWavefunction(x), -Inf, Inf)[0]

    def integrateSeparatelyYforRestoringPotential(self):
        return integrate.quad(lambda y: conjugate(self.basisFunction1.yWavefunction(y))*(y**2)*self.basisFunction2.yWavefunction(y), -Inf, Inf)[0]
    
    def integrateRestoringPotential(self):
        return (1/2.0)*(self.integrateSeparatelyY()*self.integrateSeparatelyXforRestoringPotential()+self.integrateSeparatelyX()*self.integrateSeparatelyYforRestoringPotential())




    def laplacianX(self,x):
        return derivative(self.basisFunction2.xWavefunction,x,dx=.0003,n=2,order=5)

    def laplacianY(self,y):
        return derivative(self.basisFunction2.yWavefunction,y,dx=.0003,n=2,order=5)

    def integrateSeparatelyXkinetic(self):
        return integrate.quad(lambda x: conjugate(self.basisFunction1.xWavefunction(x))*self.laplacianX(x), -Inf, Inf)[0]
            
    def integrateSeparatelyYkinetic(self):
        return integrate.quad(lambda y: conjugate(self.basisFunction1.yWavefunction(y))*self.laplacianY(y), -Inf, Inf)[0]

    def integrateKinetic(self):
        return -(1/2.0)*(self.integrateSeparatelyY()*self.integrateSeparatelyXkinetic()+self.integrateSeparatelyX()*self.integrateSeparatelyYkinetic())


    def integrateSeparatelyX2forCoulombicRepulsion(self,y1,y2,x1):
        part1 = integrate.quadrature(lambda x2: conjugate(self.basisFunction1.fullWavefunction(x1,y1))*self.basisFunction1.fullWavefunction(x1,y1)*1/sqrt((x1-x2)**2+(y1-y2)**2)*conjugate(self.basisFunction2.fullWavefunction(x2,y2))*self.basisFunction2.fullWavefunction(x2,y2), -self.lim, x1, tol=5e-2, maxiter=16, vec_func=False)[0]
        part2 = integrate.quadrature(lambda x2: conjugate(self.basisFunction1.fullWavefunction(x1,y1))*self.basisFunction1.fullWavefunction(x1,y1)*1/sqrt((x1-x2)**2+(y1-y2)**2)*conjugate(self.basisFunction2.fullWavefunction(x2,y2))*self.basisFunction2.fullWavefunction(x2,y2), x1, self.lim, tol=5e-2, maxiter=16, vec_func=False)[0]
        return part1+part2

    def integrateSeparatelyX1forCoulombicRepulsion(self,y1,y2):    
        return integrate.quadrature(lambda x1: self.integrateSeparatelyX2forCoulombicRepulsion(y1,y2,x1), -self.lim, self.lim, tol=5e-2, maxiter=16, vec_func=False)[0]

    def integrateSeparatelyY2forCoulombicRepulsion(self,y1):
        return integrate.quadrature(lambda y2: self.integrateSeparatelyX1forCoulombicRepulsion(y1,y2), -self.lim, y1, tol=5e-2, maxiter=16, vec_func=False)[0] + integrate.quadrature(lambda y2: self.integrateSeparatelyX1forCoulombicRepulsion(y1,y2), y1, self.lim, tol=5e-2, maxiter=16, vec_func=False)[0]

    def integrateSeparatelyY1forCoulombicRepulsion(self):
        return integrate.quadrature(lambda y1: self.integrateSeparatelyY2forCoulombicRepulsion(y1), -self.lim, self.lim, tol=5e-2, maxiter=16, vec_func=False)[0]

    def integrateCoulombicRepulsion(self):
        return self.integrateSeparatelyY1forCoulombicRepulsion()


    def integrateSeparatelyX2forCoulombicExchange(self,y1,y2,x1):
        part1 = integrate.quadrature(lambda x2: conjugate(self.basisFunction1.fullWavefunction(x1,y1)*self.basisFunction2.fullWavefunction(x2,y2))*1/sqrt((x1-x2)**2+(y1-y2)**2)*self.basisFunction2.fullWavefunction(x1,y1)*self.basisFunction1.fullWavefunction(x2,y2), -self.lim, x1, tol=5e-2, maxiter=16, vec_func=False)[0]
        part2 = integrate.quadrature(lambda x2: conjugate(self.basisFunction1.fullWavefunction(x1,y1)*self.basisFunction2.fullWavefunction(x2,y2))*1/sqrt((x1-x2)**2+(y1-y2)**2)*self.basisFunction2.fullWavefunction(x1,y1)*self.basisFunction1.fullWavefunction(x2,y2), x1, self.lim, tol=5e-2, maxiter=16, vec_func=False)[0]
        return part1+part2

    def integrateSeparatelyX1forCoulombicExchange(self,y1,y2):    
        return integrate.quadrature(lambda x1: self.integrateSeparatelyX2forCoulombicExchange(y1,y2,x1), -self.lim, self.lim, tol=5e-2, maxiter=16, vec_func=False)[0]

    def integrateSeparatelyY2forCoulombicExchange(self,y1):
        return integrate.quadrature(lambda y2: self.integrateSeparatelyX1forCoulombicExchange(y1,y2), -self.lim, y1, tol=5e-2, maxiter=16, vec_func=False)[0] + integrate.quadrature(lambda y2: self.integrateSeparatelyX1forCoulombicExchange(y1,y2), y1, self.lim, tol=5e-2, maxiter=16, vec_func=False)[0]

    def integrateSeparatelyY1forCoulombicExchange(self):
        return integrate.quadrature(lambda y1: self.integrateSeparatelyY2forCoulombicExchange(y1), -self.lim, self.lim, tol=5e-2, maxiter=16, vec_func=False)[0]

    def integrateCoulombicExchange(self):
        return self.integrateSeparatelyY1forCoulombicExchange()


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

maxNbasis = maxNlinear**2

print "\nTotal Number of Modes:", maxNbasis                                           # sum of all modes, this is the size of our matrix in its base dimension,
print ""

basis = [None]*maxNbasis
sigma = 1.0
##coeffs = [None]*maxNbasis

print "\nSetting Up Basis\n"
i=0
for nx in range(maxNlinear):
    for ny in range(maxNlinear):
        basis[i] = basisWaveFunction(sigma,nx,ny)
##        if silent == 0:
##            print "Cartesian X QN", nx
##            print "Cartesian Y QN", ny
##            print ""
##            print ""
        i += 1

##print "\nRunning Normalization Check\n"
##i=0
##norms = [None]*maxNbasis
##for nx in range(maxNlinear):
##    for ny in range(maxNlinear):
##        norms[i] = 1/sqrt(integrateTwoWavefunction(basis[i],basis[i]).integrateToNorm())
##        if silent == 0:
##            print "Cartesian X QN", basis[i].nx
##            print "Cartesian Y QN", basis[i].ny
##            print ""
##            print "Normalization Constant", norms[i]
##            print ""
##            print ""
##        i += 1

##print "\nCalculating Potential Integrals\n"
##restoringPotentialIntegral = [None]*maxNbasis
##for i in range(maxNbasis):
##    restoringPotentialIntegral[i]=[0]*maxNbasis
##    for j in range(i+1):
##        RPint = integrateTwoWavefunction(basis[i],basis[j]).integrateRestoringPotential()
##        if abs(RPint)>10**-9:
##            restoringPotentialIntegral[i][j] = RPint
##            restoringPotentialIntegral[j][i] = RPint
##        else:
##            restoringPotentialIntegral[i][j] = 0.0
##            restoringPotentialIntegral[j][i] = 0.0
##        print "Restoring Potential Integral",i,j,"is",restoringPotentialIntegral[i][j]
##
##dump = open("dumprestoringpotential", mode='w')
##dump.write(str(restoringPotentialIntegral))
##dump.close()
##
##print "\nCalculating Kinetic Integrals\n"
##kineticIntegral = [None]*maxNbasis
##for i in range(maxNbasis):
##    kineticIntegral[i]=[0]*maxNbasis
##    for j in range(i+1):
##        KEint = integrateTwoWavefunction(basis[i],basis[j]).integrateKinetic()
##        if abs(KEint)>10**-9:
##            kineticIntegral[i][j] = KEint
##            kineticIntegral[j][i] = KEint
##        else:
##            kineticIntegral[i][j] = 0.0
##            kineticIntegral[j][i] = 0.0
##        if silent == 0:
##            print "Kinetic Integral",i,j,"is",kineticIntegral[i][j]
##
##dump = open("dumpkinetic", mode='w')
##dump.write(str(kineticIntegral))
##dump.close()
##
##print "\nCalculating Coulombic Repulsion Integrals"
##coulombicIntegral = [None]*maxNbasis
##for i in range(maxNbasis):
##    coulombicIntegral[i]=[0]*maxNbasis
##    for j in range(i+1):
##        COint = integrateTwoWavefunction(basis[i],basis[j]).integrateCoulombicRepulsion()
##        if abs(COint)>10**-9:
##            coulombicIntegral[i][j] = COint
##            coulombicIntegral[j][i] = COint
##        else:
##            coulombicIntegral[i][j] = 0.0
##            coulombicIntegral[j][i] = 0.0
##        print "Coulombic Repulsion Integral",i,j,"is",coulombicIntegral[i][j]
##
##dump = open("dumpcoulombic", mode='w')
##dump.write(str(coulombicIntegral))
##dump.close()

print "\nCalculating Coulombic Exchange Integrals"
exchangeIntegral = [None]*maxNbasis
for i in range(maxNbasis):
    exchangeIntegral[i]=[0]*maxNbasis
    for j in range(i+1):
        XCint = integrateTwoWavefunction(basis[i],basis[j]).integrateCoulombicExchange()
        if abs(XCint)>10**-9:
            exchangeIntegral[i][j] = COint
            exchangeIntegral[j][i] = COint
        else:
            exchangeIntegral[i][j] = 0.0
            exchangeIntegral[j][i] = 0.0
        print "Coulombic Exchange Integral",i,j,"is",exchangeIntegral[i][j]

dump = open("dumpexchange", mode='w')
dump.write(str(exchangeIntegral))
dump.close()
            

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




                      
