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

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))-2*Z/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,Z,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', 'Z'],type_converters=weave.converters.blitz, compiler = 'gcc')
    
    Solution = zeros(len(R),dtype=float)
    Solution[1]=1e-7
    weave.inline(code_Numerov, ['F', 'Nmax', 'dh', 'Solution'],type_converters=weave.converters.blitz, compiler = 'gcc')

    if (Normalize):
        Norm = abs(integrate.romb(Solution**2,R[1]-R[0]))
        Solution *= 1./sqrt(Norm)
    return Solution

def Shoot(E,R,l,Z):
    u = ComputeSchrod(E,R,l,Z)
    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(Z,R,n_max,lmax,Nmax=1000):
    E0=-1.2*Z**2
    Eb=[]
    for l in range(lmax+1):
        dE = abs(E0)/12.
        u0 = Shoot(E0,R,l,Z)
        Ebl=[]
        for i in range(Nmax):
            E0 += dE
            u1 = Shoot(E0,R,l,Z)
            if u0*u1<0:
                Ebound = optimize.brentq(Shoot, E0-dE, E0, args=(R,l,Z), xtol=1e-16)
                Ebl.append([l,Ebound])
                print 'Found bound state at ', Ebound
                if len(Ebl)>=(n_max-l): break
            u0=u1
            dE = dE/1.091
        Eb += Ebl
        E0 = Ebl[0][1]
    
    Eb.sort(comp)
    return Eb

def ChargeDensity(Eb,R,Z,Z_screened):
    rho = zeros(len(R),dtype=float)
    Nelec=0
    for i,(l,Ene) in enumerate(Eb):
        u = ComputeSchrod(Ene,R,l,Z_screened,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
    return rho


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

Z_screened=1.
Z=2
n_max=2 #5
lmax=1  #3
#R0 = linspace(1e-7,90,2**12+1)   # For high precision
R0 = linspace(1e-7,20,2**10+1)
R=R0[::-1]

Eb = FindBoundStates(Z_screened,R,n_max,lmax)
rho = ChargeDensity(Eb,R,Z,Z_screened)


# Computing U by odeint
rhoSpline = interpolate.UnivariateSpline(R[::-1],rho[::-1],s=0)
U = integrate.odeint(rhsPotential, [0.0, 5.], R0, args=(rhoSpline,))[:,0]


# Computing U using Numerov
U2 = HartreeU(R0,rho[::-1])

plot(R0, U)
plot(R0, U2)
show()