from scipy import *
from scipy import integrate
from scipy import optimize
from scipy import weave
from scipy import interpolate
from pylab import *
from excor import ExchangeCorrelation
import sys
import copy


code_Numerov="""
// void Numerov(const container& F, int Nmax, double dh, container& Solution)
  #line 14 "inc_atom1.py"
  double dx = dh;
  double h2 = dx*dx;
  double h12 = h2/12;
  
  double w0 = (1-h12*F(0))*Solution(0);
  double Fx = F(1);
  double w1 = (1-h12*Fx)*Solution(1);
  double Phi = Solution(1);
  
  double w2;
  for (int i=2; i<Nmax; i++){
    w2 = 2*w1 - w0 + h2*Phi*Fx;
    w0 = w1;
    w1 = w2;
    Fx = F(i);
    Phi = w2/(1-h12*Fx);
    Solution(i) = Phi;
  }

"""

code_NumerovU="""
//void NumerovInhom(const container& U, int Nmax, double dh, container& Solution)
  #line 38 "inc_atom1.py"
  double dx = dh;
  double h2 = dx*dx;
  double h12 = h2/12;
  
  double w0 = Solution(0)-h12*U(0);
  double Ux = U(1);
  double w1 = Solution(1)-h12*Ux;
  double Phi = Solution(1);
  
  double w2;
  for (int i=2; i<Nmax; i++){
    w2 = 2*w1 - w0 + h2*Ux;
    w0 = w1;
    w1 = w2;
    Ux = U(i);
    Phi = w2+h12*Ux;
    Solution(i) = Phi;
  }

"""


code_Sch="""
  #line 62 "inc_atom1.py"
   double En=E;
   for (int i=0; i<Nmax; i++){
      F(i) = l*(l+1)/(Rm(i)*Rm(i))+Uks(i)/Rm(i)-En;
   }
"""

code_Hart="""
  #line 69 "inc_atom1.py"
   for (int i=0; i<Nmax; i++){
      U(i) = -8*M_PI*R0(i)*rho(i);
   }
"""



def ComputeSchrod(E,Rm,l,Uks):

    dh=(Rm[0]-Rm[-1])/(len(Rm)-1.)
    Nmax = len(Rm)
    
    F = zeros(len(Rm),dtype=float)
    weave.inline(code_Sch, ['F', 'Nmax', 'Rm', 'E', 'l', 'Uks'],type_converters=weave.converters.blitz, compiler = 'gcc')
    
    Solution = zeros(len(Rm),dtype=float)
    Solution[1]=1e-7
    weave.inline(code_Numerov, ['F', 'Nmax', 'dh', 'Solution'],type_converters=weave.converters.blitz, compiler = 'gcc')

    norm = integrate.romb(Solution**2,dh)
    Solution *= 1./sqrt(abs(norm))
    
    return Solution

def Shoot(E,R,l,Uks):
    u = ComputeSchrod(E,R,l,Uks)
    return u[-1]+(u[-2]-u[-1])*(0.0-R[-1])/(R[-2]-R[-1])

def comp(x,y):
    if abs(x[1]-y[1])<1e-5:
        # sort according to l quantum number when energy is degenerate
        return cmp(x[0],y[0])
    # If energy different, sort according to energy
    return cmp(x[1],y[1])

def FindBoundStates(Uks,Z,R,n_max,lmax):
    E0=-1.2*Z**2
    Eb=[]
    for l in range(lmax+1):

        Eshift=0.5
        # starting at -E0,....Eshift-1e-4
        Esearch = -logspace(-4,log10(-E0+Eshift),200)[::-1] + Eshift
        
        u0 = Shoot(Esearch[0],R,l,Uks)
        Ebl=[]
        for i in range(1,len(Esearch)):
            u1 = Shoot(Esearch[i],R,l,Uks)
            if u0*u1<0:
                Ebound = optimize.brentq(Shoot, Esearch[i-1], Esearch[i], args=(R,l,Uks), xtol=1e-16)
                Ebl.append([l,Ebound])
                print 'Found bound state at ', Ebound/2., 'Hartree'
                if len(Ebl)>=(n_max-l): break
            u0=u1
        Eb += Ebl
        E0 = Ebl[0][1]
    
    Eb.sort(comp)
    return Eb

def ChargeDensity(Eb,R,Z,Uks):
    rho = zeros(len(R),dtype=float)
    Eband=0.0
    Nelec=0
    for i,(l,Ene) in enumerate(Eb):
        #print 'state=', l, Ene
        u = ComputeSchrod(Ene,R,l,Uks)
    
        dN = 2*(2*l+1)
        if Z >= Nelec+dN:
            ferm=1.
        else:
            ferm= (Z-Nelec)/(dN+0.0)
            
        Nelec += dN*ferm
        rho += u*u*dN*ferm / (4*pi*R**2)
        Eband += Ene*dN*ferm
        
        if Nelec>=Z: break
    
        #print 'Adding ', dN*ferm, 'electrons of l=', l, 'Ene=', Ene
    return (rho, Eband)

def rhsPotential(u,r,rhoSpline):
    """Right hand side (derivative) for the Poisson equation to computer V_H
       gives an array d/dr(U,U') == (U',U'')
    """
    return [u[1],-8*pi*r*rhoSpline(r)]
 

def HartreeU(R0,rho):
    U = zeros(len(R0),dtype=float)
    Nmax = len(R0)
    weave.inline(code_Hart,['U', 'R0', 'rho', 'Nmax'],type_converters=weave.converters.blitz, compiler = 'gcc')
    dh = (R0[-1]-R0[0])/(len(R0)-1.)
    Solution = zeros(len(R0),dtype=float)
    Solution[0]=0.0
    Solution[1]=5*dh
    weave.inline(code_NumerovU,['U', 'Nmax', 'dh', 'Solution'],type_converters=weave.converters.blitz, compiler = 'gcc')
    return Solution

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


D_nmax={1:1, 2:1, 3:2, 4:2, 5:2, 6:2, 7:2, 8:2, 9:2, 10:2, 11:3, 12:3, 13:3,14:3, 15:3, 16:3, 17:3, 18:3, 19:4}


Z=18
Etol=1e-6
Nitt=200
n_max=D_nmax[Z]
lmax=D_nmax[Z]-1
R0 = linspace(1e-7,10,2**14+1)
R=R0[::-1]

mixr=0.5  # Need to mix to get converged for larger Z

Uks = -2*ones(len(R))
Eold=0.0
for itt in range(Nitt):
    Eb = FindBoundStates(Uks,Z,R,n_max,lmax)
    (rho, Eband) = ChargeDensity(Eb,R,Z,Uks)
    rho0 = rho[::-1]

    if itt>0: rho0 = mixr*rho0 + (1-mixr)*rho_old
    rho_old = copy.deepcopy(rho0)
    
    U = HartreeU(R0,rho0)
    ualpha = (2*Z-U[-1])/R0[-1]
    U += ualpha*R0
    
    
    exc = ExchangeCorrelation()
    Vxc = [2*exc.Vx(rs(rh))+2*exc.Vc(rs(rh)) for rh in rho0]
    
    Uks0 = (U-2*Z) + Vxc*R0  # Kohn-Sham potential*r: Uks/r == V_ks = V_H/r - 2*Z/r + Vxc
    Uks = Uks0[::-1]

    # Total energy
    ExcVxc = [2*exc.EcVc(rs(rh))+2*exc.ExVx(rs(rh)) for rh in rho0]

    pot=(ExcVxc*R0**2-0.5*U*R0)*rho0*4*pi
    Eint = integrate.romb(pot,R0[1]-R0[0])
    Etot = Eint + Eband
    
    print 'Iteration', itt, ' Enery[Ry]=', Etot, 'Energy[Hartree]=', Etot/2., 'Ediff=', abs(Etot-Eold)
    if abs(Etot-Eold)<Etol: break
    
    Eold=Etot
    
plot(R0, U, label='U')
plot(R0, Vxc, label='Vxc')
plot(R0, Uks0, label='Uks0')
legend(loc='best')
plot()
show()

plot(R0, rho0*4*pi*R0**2, label='rho')
legend(loc='best')
show()