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

def normalize(u):
    " normalizes solution "
    u2 = u*u    
    nrm = integrate.simps(u2, r_mesh)
    return -u/sqrt(abs(nrm))

def Schoedinger_deriv1(u, r, l, eps, Z, Veff):
    """ This function returns derivative to solve Schroedinger
    equation for hydrogen atom.
    """
    Veff = interpolate.splev(r, Veff)
    return array([u[1], 2*(l*(l+1)/(2*r**2)+Veff-eps)*u[0]])


def Shoot1(eps, r_mesh, l, Z, Veff):
    " Starting from r=infty, finds the solution of Schroedinger equation at r=0"
    # initial condition
    u0 = array([0.,1e-10])
    # actual integration of Sch. equation
    ub = integrate.odeint(Schoedinger_deriv1, u0, r_mesh, args=(l, eps, Z, Veff))
    u_at_zero = ub[-1,0] + (ub[-2,0]-ub[-1,0])*(0.0-r_mesh[-1])/(r_mesh[-2]-r_mesh[-1])
    return u_at_zero

def FindBoundStates1(l, Z, r_mesh, nmax, eps_start, Veff):
    " Finds nmax-l eigenstates of the radial Schroedinger equation for given l"
    eps = eps_start
    states = []
    for n in range(l,nmax):
        deps = 0.3/(n+1)**3
        val0 = Shoot1(eps, r_mesh, l, Z, Veff)
        val1 = val0
        while (val0*val1>0):
            eps += deps
            val0 = val1
            val1 = Shoot1(eps, r_mesh, l, Z, Veff)
            
        zero = optimize.brentq(Shoot1, eps-deps, eps, args=(r_mesh, l, Z, Veff) )
        states.append([n,l,zero])
        
    return states
        

def FindAllBoundStates1(lmax, Z, r_mesh, eguess, Veff):
    # All bound states will be saved
    bstates=[]
    for l in range(lmax):
        eguess = -0.5*Z*Z/(l+1.)**2-3. 
        bstates.append(FindBoundStates1(l, Z, r_mesh, lmax, eguess, Veff)) # Finds new bound state at certain l
        eguess = bstates[l][0][2]         # guess for the first bound state (uses the fact that E(1s)>E(1p)>E(1d)...)
    return bstates

def PlotBoundStates(lmax, r_mesh, bstates, Veff, aplot=False):
    # Prints the bound states and plots them
    for l in range(lmax):
        for ni in range(0,lmax-l):
            print "  ", l, ni, bstates[l][ni]
            eps = bstates[l][ni][2]
            ub = integrate.odeint(Schoedinger_deriv1, [0,1], r_mesh, args=(l, eps, Z, Veff)) # recomputes psi
            u = normalize(ub[:,0]) # and normalizes it
            if (aplot): plot(r_mesh, u)        # plotting
    if (aplot):
        axis([0., 50, -0.3, 0.8])
        show()

def addl(a, b):
    return a + b

def cmp_energy(a, b):
    return cmp(a[2], b[2])
    
def ComputeDensity(Z, r_mesh, bstates, Veff):
    states = reduce(addl, bstates)
    states.sort(cmp_energy)
    rho = zeros(len(r_mesh), dtype=float)
    N = 0
    Eband=0
    for state in states:
        l = state[1]
        eps = state[2]
        ub = integrate.odeint(Schoedinger_deriv1, [0,1], r_mesh, args=(l, eps, Z, Veff)) # recomputes psi
        u = normalize(ub[:,0]) # and normalizes it
        u2 = u**2/(4*pi*r_mesh**2)
        
        if (N+2*(2*l+1) < Z):
            ferm = 1
        else:
            ferm = (Z-N)/(2.*(2.*l+1.))

            
        rho = rho + array(u2)*2*(2*l+1)*ferm
        N = N + 2*(2*l+1)*ferm
        Eband += 2*(2*l+1)*ferm*eps
        if (N>Z): break

    return (rho,Eband)

def Poisson_derivative(y, r, rho_spline):
    """Right hand side of the differential equation.
    Here y = [u, v].
    """
    rho = interpolate.splev(r, rho_spline)
    return array([y[1], -4*pi*r*rho]) # (\dot{U}, \dot{\dot{U}})

def SolvePoisson(Z, Rmesh, rho_spline):
    """Solve the Poisson equation: given spline of charge density rho_spline
    it computes Hartree potential, which satisfies Poisson equation
    U''(r) = -4*pi*r*rho(r)
    """
    y_t = integrate.odeint(Poisson_derivative, [0, 1.0], Rmesh, args = (rho_spline,)) # integration of the equation
    # adding homogeneous solution to satisfay boundary conditions: U(0)=0, U(infinity)=Z
    U = y_t[:,0]
    alpha = (Z - U[-1])/Rmesh[-1]
    U += alpha*Rmesh
    return U

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

def HartreeEnergy(rho, Rmesh, Uhartree):
    iarr = rho*Uhartree*Rmesh*(2*pi)

    return integrate.simps(iarr, Rmesh)

def ExcVxc(rho, Rmesh, exc):
    
    EV = [exc.ExVx(rs(rh)) + exc.EcVc(rs(rh)) for rh in rho]
    iarr = EV*rho*Rmesh**2*(4*pi)

    return integrate.simps(iarr, Rmesh)
    
if  __name__ == '__main__':

    # mesh of r-points
    r_mesh = logspace(2,-6,400)
    # increasing mesh
    Rmesh = r_mesh[::-1]
    exc = ExchangeCorrelation()
    
    # Z and maximal number of l needed
    Z = 2
    lmax = 1
    admix = 0.8
    Maxitt = 5
    
    rho = array(map(lambda r: exp(-Z/2.*r)*Z**4/(64.*pi), Rmesh)) # starting guess for density

    all_rho=[]
    pEtotal=0
    
    for it in range(Maxitt):
        rho_spline = interpolate.splrep(Rmesh, rho, s=0)   # spline density with cubic spline

        # Hartree potential
        Uhartree = SolvePoisson(Z, Rmesh, rho_spline)
        # Exchange correlation potential
        Vxc = [exc.Vx(rs(rh)) + exc.Vc(rs(rh)) for rh in rho]

        Veff0 = (-Z*ones(len(Rmesh)) + Uhartree)/Rmesh + Vxc
        Veff = interpolate.splrep(Rmesh, Veff0, s=0)

        # starting guess for the lowest energy
        eguess = -0.6*Z**2

        bstates = FindAllBoundStates1(lmax, Z, r_mesh, eguess, Veff)

        # Prints the bound states
        PlotBoundStates(lmax, r_mesh, bstates, Veff)

        # New density
        (rnrho,Eband) = ComputeDensity(Z, r_mesh, bstates, Veff)
        nrho = rnrho[::-1]

        # We store the new density for plotting
        nrh = [nrho[i]*Rmesh[i]**2 for i in range(len(Rmesh))]
        all_rho.append(nrh)

        # Total energy and various contribution to the total energy are computed
        Ehartree = HartreeEnergy(nrho, Rmesh, Uhartree)
        Exc = ExcVxc(nrho, Rmesh, exc)
        Etotal = Eband - Ehartree + Exc
        print 'Eband=', Eband, '  Ehartree = ', Ehartree, '  Exc = ', Exc
        print 'Etotal = ', Etotal, '  Ediff = ', Etotal-pEtotal
        
        rho = (1.-admix)*rho + admix*nrho
        pEtotal = Etotal

    styles = ['^-', 'v-', '<-', '>-', 'o-', 's-']
    colors = ['y', 'm', 'c', 'b', 'g', 'r']

    for i in range(len(all_rho)):
        plot(Rmesh, all_rho[i], styles[i%len(styles)]+colors[i%len(colors)])
        
#    axis([0,10,0,0.05])
    show()
