#!/usr/bin/env python3
"""LDA calculation of hydrogen in spherical coordinates.

Reference values (LDA):
https://www.nist.gov/pml/atomic-reference-data-electronic-structure-calculations-hydrogen

We actually solve for r times the wavefunction, rpsi, which simplifies
the equation:

     1  __2           /                    Z  \
  - --- \/  (r psi) + | vHa[n] + vxc[n] - --- | (r psi) = eps (r psi)
     2                \                    r  /

The GPAW all-electron module can also be used for this task.
"""

import pylab as plt
import numpy as np
from gpaw.utilities import hartree as _hartree
from gpaw.xc.lda import lda_x, lda_c

# Set up grid:
ng = 400
rmax = 12.0
# We don't want 0 on the grid, displace by one grid point:
r0 = np.linspace(0.0, rmax, ng + 1)
r = r0[1:]  # <-- This is our grid
dr = r[1] - r[0]

# Set up Laplacian:
lap = np.zeros((ng, ng))
lap.flat[::ng + 1] = -2.0 / dr**2
lap.flat[1::ng + 1] = lap.flat[ng::ng + 1] = 1.0 / dr**2
T = -0.5 * lap

# External potential for hydrogen
Z = 1.0
vHa_nuc = -Z / r


# H 1s wavefunction:
rpsi_init = 2.0 / np.sqrt(4.0 * np.pi) * r * np.exp(-r)

def test_potential(func, r, n, itest, dn=1e-6):
    # Test whether potential is the functional derivative of the energy
    # in a given point for a given function.
    e, v = func(r, n)
    n0 = n[itest]
    n[itest] += 0.5 * dn
    eplus, _ = func(r, n)
    n[itest] -= dn
    eminus, _ = func(r, n)
    n[itest] = n0

    v_fd = (eplus[itest] - eminus[itest]) / dn
    err = abs(v_fd - v[itest])
    print('{} test: err={} fd={} ref={}'.format(func.__name__, err,
                                                v_fd, v[itest]))


def get_ex(r, rho):
    # Get energy density and potential for LDA exchange
    prefactor = -(3.0 / np.pi)**(1. / 3.)
    ex = 3. / 4. * prefactor * rho**(4. / 3.)
    vx = prefactor * rho**(1. / 3.)
    return ex, vx


def _get_e_gpaw(func, r, rho, *args):
    # Shortcut for calling GPAW
    e = np.zeros(len(rho))
    v = np.zeros(len(rho))
    func(0, e, rho, v, *args)
    return e, v

def get_ex_gpaw(r, rho):
    # Same as get_ex, but uses GPAW (for testing)
    return _get_e_gpaw(lda_x, r, rho)

def get_ec_gpaw(r, rho):
    # Same as get_ex, but correlation
    return _get_e_gpaw(lda_c, r, rho, None)


def hartree(r, rrho):
    # Call spherical Hartree solver from GPAW
    dr = r[1] - r[0]
    rvHa = np.zeros(len(r))
    # _hartree(l, rho_r_dr, r, vHa_r)
    _hartree(0, rrho * dr, r, rvHa)
    return rvHa

def hartree1(r, rho):
    # Same as hartree() except we take rho -- for testing the same way
    # as the XC functions
    rrho = r * rho
    rvHa = hartree(r, rrho)
    vHa = rvHa / r
    EHa = vHa * rho
    return EHa, vHa


rpsi = rpsi_init
r2 = r**2
r2rho = 0.0

for i in range(10):
    print('=============')
    print('iter {}'.format(i))
    print('=============')
    r2rho_new = rpsi**2
    r2rho_change = 4.0 * np.pi * abs(r2rho_new - r2rho).sum() * dr
    print('drho change', r2rho_change)
    r2rho = r2rho_new
    rrho = r2rho / r
    rho = rrho / r
    charge = 4.0 * np.pi * r2rho.sum() * dr
    print('charge', charge)
    assert abs(charge - 1) < 1e-6, charge

    # ex is an array of the energy density at each point
    ex, vx = get_ex(r, rho)
    #ex1, vx1 = get_ex_gpaw(r, rho)
    #ex_err = np.abs(ex - ex1).max()
    # assert ex_err < 1e-14
    ec, vc = get_ec_gpaw(r, rho)
    #test_potential(get_ec_gpaw, r, rho, 10)
    #test_potential(get_ex, r, rho, 10)

    exc = ex + ec
    vxc = vx + vc
    Exc = 4.0 * np.pi * (exc * r2).sum() * dr
    print('Exc', Exc)

    # The GPAW Hartree solver wants a grid which starts at 0:
    rrho0 = np.empty(ng + 1)
    rrho0[0] = 0.0
    rrho0[1:] = rrho
    rvHa_el = hartree(r0, rrho0)[1:]

    # test_potential(hartree1, r, rho, 10)
    EHa_el = 4.0 * np.pi * 0.5 * (rrho * rvHa_el).sum() * dr
    print('EHa_el', EHa_el)

    EHa_nuc = -4.0 * np.pi * Z * rrho.sum() * dr
    print('EHa_nuc', EHa_nuc)
    EHa = EHa_el + EHa_nuc
    print('EHa_total', EHa)

    veff = (rvHa_el - Z) / r + vxc
    H = T + np.diag(veff)

    eps_n, rpsi_gn = np.linalg.eigh(H)
    print('1s eigenval', eps_n[0])
    Etot = eps_n[0] - 4.0 * np.pi * (r2rho * veff).sum() * dr + EHa + Exc
    rpsi = rpsi_gn[:, 0].copy()
    rpsi *= np.sign(rpsi[1])

    norm = 4.0 * np.pi * (rpsi**2).sum() * dr
    rpsi /= np.sqrt(norm)


fig, ax = plt.subplots(1, 2)

psiax = ax[0]
psiax.plot(r, rpsi, label='r psi')
psiax.plot(r, r2rho, label='r^2 rho')
psiax.legend(loc='best')

vax = ax[1]
vHa_el = rvHa_el / r
veff_test = vx + vc + vHa_el + vHa_nuc
vax.plot(r, vx, label='vx')
vax.plot(r, vc, label='vc')
vax.plot(r, vHa_el, label='vHa_el')
vax.plot(r, vHa_nuc, label='vHa_nuc')
vax.plot(r, veff, label='veff')
vax.plot(r, veff_test, ls='--', label='veff [sum]')
vax.legend(loc='best')
vax.axis(ymax=1, ymin=-2)

plt.show()