from scipy import *
from pylab import *
from excor import *

def rhs(u, r, l, Z, vSpline, vxcSpline, en):
    """ Returns right hand side of differential equation """
    valV=interpolate.splev(r, vSpline)
    vxcVal=interpolate.splev(r, vxcSpline)
    return (u[1], 2.*(l*(l+1.)/(2.*(r**2.)) + (-Z+valV)/r + vxcVal - en)*u[0])

def getPsiAtZero(en, l, Z, vSpline, vxcSpline, mesh):
    """ Solves differential equation and returns extrapolated
    wave function at origin """
    uStart = (0, getPsiStartDeriv())
    u=integrate.odeint(rhs, uStart, mesh, args=(l, Z, vSpline, vxcSpline, en))
    return u[-1,0] - (u[-2,0]-u[-1,0])*mesh[-1]/(mesh[-2]-mesh[-1])

def getPsi(en, l, Z, vSpline, vxcSpline, mesh):
    """ Solves differential equation and returns whole wave function"""
    uStart = (0, getPsiStartDeriv())
    u=integrate.odeint(rhs, uStart, mesh, args=(l, Z, vSpline, vxcSpline, en))
    return u[:,0]

def getPsiStartDeriv():
    return 1e-5

def findEnergies(l, Z, vSpline, vxcSpline, eMin, eMax, mesh):
    """ Find all eigenstates for some values of l and Z between
    energies eMin and eMax. Use nonequidistant stepping! Use one with 1/n^2 """
    nr=0;
    enAll=[]
    numMin=sqrt(-0.5*(Z-1.)**2./eMin) 
    numMax=sqrt(-0.5*(Z-1.)**2./eMax)
    numStep=0.1
    if(abs(numMax-numMin)/numStep < 10.): # at least 10 steps
        numStep=abs(numMax-numMin)/10.
    numCurrent=numMin
    enCurrent=(-0.5*(Z-1.)**2./(numCurrent)**2)
   
    oldPsi=getPsiAtZero(enCurrent, l, Z, vSpline, vxcSpline, mesh);
    while numCurrent+numStep<numMax:
        enCurrent=(-0.5*(Z-1.)**2./(numCurrent)**2)
        enNext=(-0.5*(Z-1.)**2./(numCurrent+numStep)**2)
        newPsi=getPsiAtZero(enNext, l, Z, vSpline, vxcSpline, mesh)
        if(oldPsi*newPsi<0):
            en=optimize.brentq(getPsiAtZero, enCurrent, enNext, (l, Z, vSpline, vxcSpline, mesh))
            enAll.append([en, nr])
            nr=nr+1
            numCurrent=numCurrent+numStep
        else:
            numCurrent=numCurrent+numStep
        oldPsi=newPsi
    return enAll

def findEnMinMax(l, Z, numElectrons):
    """ Returns region in which to search for bound states, this is used only in first iteration """
    nMax=1;
    while 1:
        num=2.*nMax*(nMax+1.)*(2.*nMax+1.)/6.
        if num>=numElectrons: break
        else:
            nMax=nMax+1
    eMax=(-0.5*((-1.)**2.)/(nMax**2.))/2.
    eMin=(-0.5*((Z-1.)**2.)/((l+1.)**2.))*1.5
    return (eMin,eMax)

def findAllStates(numElectrons, Z, vSpline, vxcSpline, mesh, oldStates, guessDifferences, oldOldStates=[]):
    """ Finds all energy states for some fixed value of Z, and with extra potential
    vSpline and vxcSpline. Asumes that for fixed value of l lowest state energy state increases
    as l is increased. Returns array of elements of form: (Energy, n_r, l) """
    tempL=0
    allStates=[]
    if(len(oldStates) == 0):
        #Calculating states for the first time
        while 1:
            (eMin, eMax)=findEnMinMax(tempL, Z, numElectrons)
            states=findEnergies(tempL, Z, vSpline, vxcSpline, eMin, eMax, mesh)
            if len(states)==0:
                break
            else:
                for i in states:
                    allStates.append([i[0], i[1], tempL])
            tempL=tempL+1
    else:
        #Using previously found states to more easily compute new states
        #just find the states in the neighbourhood of last states
        for j in range(0, len(oldStates)):
            if(guessDifferences==True):
                eMin=oldStates[j][0] - 2.*abs(oldOldStates[j][0]-oldStates[j][0]) - 0.02 #add the constant 0.02 because low energy states can more easily fluctuate in energy
                eMax=oldStates[j][0] + 1.5*abs(oldOldStates[j][0]-oldStates[j][0]) + 0.02
            else:
                eMin=oldStates[j][0]*1.5
                eMax=oldStates[j][0]*0.5
            states=findEnergies(oldStates[j][2], Z, vSpline, vxcSpline, eMin, eMax, mesh)
            if(len(states)!=1):
                print "Expecting one state!", " Found: ", len(states)
                quit()
            else:
                allStates.append([states[0][0], oldStates[j][1], oldStates[j][2]])
    return allStates
            
def normalize(func, mesh):
    """ Normalizes the wavefunction."""
    square=map(lambda u:(u*u), func)
    norm=integrate.simps(square,mesh)
    funcRet= -func/sqrt(abs(norm))
    return funcRet

def fillUpLowestStates(states, numElectrons):
    """ Takes as input the states that he has found, and then returns Z
    lowest states with their degeneracies
    Takes array of form   (Energy, n_r, l)
    Returns array of form (Energy, n_r, l, degeneracy)   """
    states.sort(lambda x,y: cmp(x[0],y[0]))
    occStates=[]
    num=0
    for i in range(0,len(states)):
        l=states[i][2]
        if(num+2*(2*l+1) <= numElectrons):
            occStates.append([states[i][0], states[i][1], states[i][2], 2*(2*l+1)])
            num+=2*(2*l+1)
        else:
            occStates.append([states[i][0], states[i][1], states[i][2], numElectrons-num])
            break
        if(num==numElectrons): break
    return occStates

def getAdmixDensity(states, Z, vSpline, vxcSpline, mesh, oldDenSpline, admix):
    """ Admixes old and new electron density, input parameters
    are a list of (energy, n_r, l, degeneracy) and others.
    This is threedimensional density, not divided by r
    or something like that."""
    density=zeros(len(mesh), dtype=float)
    sqr=zeros(len(mesh), dtype=float)
    u0=zeros(len(mesh), dtype=float)
    for i in states:
        u=getPsi(i[0], i[2], Z, vSpline, vxcSpline, mesh)
        deg=i[3]
        u0=normalize(u, mesh)
        sqr=(u0*u0)/(4.*pi*mesh*mesh)
        sqr=sqr*deg
        density=density+sqr
    for i in range(0,len(mesh)):
        density[i]=admix*density[i]+(1.-admix)*interpolate.splev(mesh[i],oldDenSpline)
    meshInv=mesh[::-1]
    admDenInv=density[::-1]
    admDenSpline=interpolate.splrep(meshInv, admDenInv, s=0)
    return admDenSpline

def rhsPotential(u, r, densitySpline):
    """ Right hand side (derivative) for potential equation"""
    rho=interpolate.splev(r, densitySpline)
    return (u[1], -4.*pi*r*rho)
    return 0

def findHartreePotential(densitySpline, Z, mesh):
    """ Calculates Hartree potential out of charge density.
    Real potential, one that enters sch equation, is this one divided by r!"""
    meshInv=mesh[::-1]
    uStart=(0, 1)
    uInv=integrate.odeint(rhsPotential, uStart, meshInv, args=(densitySpline,))
    solutionInv=map(lambda x,r:(x+r*(Z-uInv[-1,0])/meshInv[-1]), uInv[:,0], meshInv)
    solutionSpline=interpolate.splrep(meshInv, solutionInv, s=0)
    return solutionSpline

def getRs(rho):
    if (rho<1e-100): return 1e100
    else: return pow(4.*pi*rho/3.,-1./3.)

def findXCPotential(density, mesh, excObj):
    """ Calculates exchange correlation potential in LDA"""
    xcPot=zeros(len(mesh), dtype=float)
    for i in range(0,len(mesh)):
        rho=interpolate.splev(mesh[i], density)
        rs=getRs(rho)
        xcPot[i]=excObj.Vx(rs)+excObj.Vc(rs)
    xcPotInv=xcPot[::-1]
    meshInv=mesh[::-1]
    xcPotSpline=interpolate.splrep(meshInv, xcPotInv, s=0)
    return xcPotSpline

def findTotalEnergy(states, denSpline, harPotSpline, mesh, excObj):
    """ Calculates total energy. """
    bandEnergy=0
    for i in states: bandEnergy=bandEnergy+i[0]*i[3]
    meshInv=mesh[::-1]
    integrandHa=zeros(len(meshInv), dtype=float)
    integrandEx=zeros(len(meshInv), dtype=float)
    for j in range(0, len(meshInv)):
        den=interpolate.splev(meshInv[j],denSpline)
        ehartree=0.5 * interpolate.splev(meshInv[j], harPotSpline) / meshInv[j]
        excvxc=excObj.ExVx(getRs(den))+excObj.EcVc(getRs(den))
        spaceFactor=4.*pi*meshInv[j]**2.
        integrandHa[j]= spaceFactor*den*ehartree
        integrandEx[j]= spaceFactor*den*excvxc
    enHa=integrate.simps(integrandHa, meshInv)
    enEx=integrate.simps(integrandEx, meshInv)
    totalEnergy=bandEnergy-enHa+enEx
    return (totalEnergy, bandEnergy)

def psiSq100(r,Z): return Z**3.*(1./pi)*exp(-2.*r*Z)
def psiSq200(r,Z): return Z**3.*(1./(4.*2.*pi))*(1-r*Z/2.)**2.*exp(-1.*r*Z)
def psiSq210(r,Z): return Z**3.*(1./(4.*24.*pi))*(r*Z)**2.*exp(-1.*r*Z)
def psiSq300(r,Z): return Z**3.*(1./(4.*27.*pi))*(2.-(4./3.)*r*Z+(4./27.)*(r*Z)**2.)**2.*exp(-(2./3.)*r*Z)
def psiSq310(r,Z): return Z**3.*(1./(4.*6.*27.**2.*pi))*(8.*r*Z-(4./3.)*(r*Z)**2.)**2.*exp(-(2./3.)*r*Z)
def psiSq320(r,Z): return Z**3.*(1./(4.*30.*81.**2.*pi))*(4.*(r*Z)**2.)**2.*exp(-(2./3.)*r*Z)


def guessStartDensity(Z, mesh):
    """ Returns tweaked starting density. Use hydrogenlike wavefunctions. Works for closed shells only."""
    if(Z==2):
        dens=array(map(lambda r:2.*psiSq100(r,Z-0.), mesh))
    elif (Z==4):
        dens=array(map(lambda r:2.*psiSq100(r,Z-1.) + 2.*psiSq200(r,Z-1.), mesh))
    elif (Z==10):
        dens=array(map(lambda r:2.*psiSq100(r,Z-1.) + 2.*psiSq200(r,Z-3.) + 6.*psiSq210(r,Z-5.), mesh))
    elif (Z==12):
        dens=array(map(lambda r:2.*psiSq100(r,Z-1.) + 2.*psiSq200(r,Z-3.) + 6.*psiSq210(r,Z-5.) + 2.*psiSq300(r,Z-9.), mesh))
    elif (Z==18):
        dens=array(map(lambda r:2.*psiSq100(r,Z-1.) + 2.*psiSq200(r,Z-3.) + 6.*psiSq210(r,Z-5.) + 2.*psiSq300(r,Z-9.) + 6.*psiSq310(r,Z-11.), mesh))
    else:
        print "Works only for Z=2, 4, 10, 12, 18 !"
        quit()
    meshInv=mesh[::-1]
    densInv=dens[::-1]
    denSpline=interpolate.splrep(meshInv, densInv, s=0)
    return denSpline
    
def runSelfConsistentLoop(Z, numElectrons, mesh, maxSteps, relPrecision):
    """ Runs LDA selfconsist loop."""
    allDensities=[]
    allMeshInv=[]
    meshInv=mesh[::-1]
    denSpline=guessStartDensity(Z,mesh) #start by guessing density. the same as in Z=1 atom, but with Z electrons
    admix=0.4 # use small admix because then energies in subsequent iterations dont change much and it is easier to find them
    excObj=ExchangeCorrelation()
    states=[]
    allStates=[]
    for i in range (0, maxSteps):
        harSpline=findHartreePotential(denSpline, Z, mesh) #find Hartree potential from density
        vxcSpline=findXCPotential(denSpline, mesh, excObj) #find exchange-corr potential from density
        if(i-2>=0):
            if(len(allStates[i-2])==len(allStates[i-1])):  #Use oldstates-2 to speed up the algorithm. Can use old differences between iterations to guess range for next iteration
                states=findAllStates(numElectrons, Z, harSpline, vxcSpline, mesh, allStates[i-1], True, allStates[i-2])
        else:
            states=findAllStates(numElectrons, Z, harSpline, vxcSpline, mesh, states, False) #find states in this potential
        states=fillUpLowestStates(states, numElectrons) #fills up Z lowest states
        allStates.append(states)
        allEnergies=findTotalEnergy(states, denSpline, harSpline, mesh, excObj) #calculate total energy

        tempd=map(lambda x:4.*pi*x**2.*interpolate.splev(x, denSpline),meshInv) #calculate radial density
        charge=integrate.simps(tempd, meshInv)
        allDensities.append(tempd) #save it for plotting at the end
        allMeshInv.append(meshInv)
        print "Step = ", i+1, " total en = ", allEnergies[0], " band en = ", allEnergies[1], "charge  = ", charge

        if(abs((charge-Z)/Z)>0.05): #
            print "Lost some states in the process..."
            quit()

        if(i>0):
            if(abs((allEnergies[0]-oldEnergy)/oldEnergy) < relPrecision): # if precision achieved stop
                print
                print "Filled states are:"
                for j in states:
                    print "Orbital energy = ", j[0], " n_r = ", j[1], " l = ", j[2], " degeneracy = ", j[3]
                print
                print "Total energy for Z = ", Z, " is = ", allEnergies[0], " Hartree"
                print
                for i in range(len(allDensities)):
                    plot(meshInv, allDensities[i], '-o')
#                plot(transpose(allMeshInv), transpose(allDensities),"o-")
                show()
                return 1
        oldEnergy=allEnergies[0] #save last total energy for comparison in next step
        oldDenSpline=map(lambda x:x,denSpline) #save old density
        denSpline=getAdmixDensity(states, Z, harSpline, vxcSpline, mesh, oldDenSpline, admix) #calculate new density
    print "Didn't converge after: ", maxSteps, "steps!"
    return 0


#Code works for Z=2, 4, 10, 12, 18. These are spherically symmetrical
#First iteration is usually much slower then other iterations...
#For Z=2  energy is -2.83484     should be -2.834836 (all from NIST)
#For Z=4  energy is -14.4470     should be -14.447209
#For Z=10 energy is -128.2338    should be -128.233481
#For Z=12 energy is -199.1393    should be -199.139406
#For Z=18 energy is -525.948     should be -525.946195
Z=2.
numElectrons=Z
maxSteps=100
relPrecision=1e-6
mesh=logspace(1., -6.-log(Z-1.), 200) #if you change Z have to change mesh also. Probably can guess this better....

print "Z = ", Z
print "Relative precision in total energy = ", relPrecision
print
runSelfConsistentLoop(Z, numElectrons, mesh, maxSteps, relPrecision)