from scipy import *
from scipy import weave
from scipy import integrate
from scipy import optimize
from scipy import interpolate
from pylab import *
from scipy.misc import derivative
from excor import ExchangeCorrelation

code_Numerov="""
// void Numerov(const container& F, int Nmax, double dh, container& Solution)
  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)
  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="""
   for (int i=0; i<Nmax; i++){
      F(i) = l*(l+1)/(R(i)*R(i))+Uks(i)/R(i)-E;
   }
"""

code_Hart="""
   for (int i=0; i<Nmax; i++){
      U(i) = -8*M_PI*R(i)*rho(i);
   }
"""

def ComputeSchrod(E,R,l,Uks,Normaliza=False):
    dh=(R[0]-R[-1])/(len(R)-1.)
    Nmax = len(R)
    
    F = zeros(len(R),dtype=float)
    weave.inline(code_Sch, ['F', 'Nmax', 'R', 'E', 'l', 'Uks'],type_converters=weave.converters.blitz, compiler = 'gcc')
    
    Solution = zeros(len(R),dtype=float)
    Solution[1]=1e-10
    weave.inline(code_Numerov, ['F', 'Nmax', 'dh', 'Solution'],type_converters=weave.converters.blitz, compiler = 'gcc')

    
    ## Normalize
    typical_value = Solution[int(0.9*len(Solution))]
    ## To avoid owerfloow
    Solution *= 1./typical_value
    Norm = abs(integrate.romb(Solution**2,R[1]-R[0]))
    Solution *= 1./sqrt(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,Nmax=1000):
    E0 = -1.2*Z**2
    Eb=[]
    for l in range(lmax+1):
        
        #E0 = -1.2*Z**2
        Eshift=0.5 # sometimes energies can be positive!!!
        Esearch = -logspace(-4,log10(-E0+Eshift),200)[::-1] + Eshift 

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

    print 'Found the following bound states:', Eb
    return Eb

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

def rs(rho):
    "Given density, returns rs."
    if rho<1e-100: return 1e100
    return pow(3/(4*pi*rho),1/3.)

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


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}
D_lmax={1:0, 2:0, 3:1, 4:1, 5:1, 6:1, 7:1, 8:1, 9:1, 10:1, 11:2, 12:2}

Z=4
n_max=D_nmax[Z]
lmax=D_lmax[Z]
R0 = linspace(1e-7,20,2**12+1)
R=R0[::-1]
Uks = -2.*ones(len(R0),dtype=float)

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

for itt in range(30):
    
    Eb = FindBoundStates(Uks,Z,R,n_max,lmax)
    rho = ChargeDensity(Eb,R,Z,Uks)
    rho0 = rho[::-1]

    if itt>0: rho0 = rho0*mixr + (1-mixr)*rho_old
    rho_old = copy(rho0)
    
    # Computing U using Numerov
    U = HartreeU(R0,rho0)
    alpha = (2*Z-U[-1])/R0[-1]
    U += alpha*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]
    
    print 'Itteration', itt, 'finished'
    
    
plot(R0, rho0*4*pi*R0**2, label='rho')
legend(loc='best')
show()