import copy
import sys
from itertools import cycle
from math import pi

import numpy as np
from numpy.linalg import eigh
from scipy.special import gamma
import ase.units as units
from ase.data import atomic_numbers, atomic_names, chemical_symbols
from ase.utils import seterr

import gpaw.cgpaw as cgpaw
from gpaw.xc import XC
from gpaw.gaunt import gaunt
from gpaw.atom.configurations import configurations
from gpaw.atom.radialgd import (AERadialGridDescriptor,
                                AbinitRadialGridDescriptor)


# Velocity of light in atomic units:
c = 2 * units._hplanck / (units._mu0 * units._c * units._e**2)

# Colors for s, p, d, f, g:


class _Colors:
    def __init__(self, items):
        self.items = items

    def __getitem__(self, i):
        return self.items[i % len(self.items)]

    def __iter__(self):
        yield from cycle(self.items)


colors = _Colors('krgbycm')


class GaussianBasis:
    def __init__(self, l, alpha_B, rgd, eps=1.0e-7):
        """Guassian basis set for spherically symmetric atom.

        l: int
            Angular momentum quantum number.
        alpha_B: ndarray
            Exponents.
        rgd: GridDescriptor
            Grid descriptor.
        eps: float
            Cutoff for eigenvalues of overlap matrix."""

        self.l = l
        self.alpha_B = alpha_B
        self.rgd = rgd
        self.eps = eps

        A_BB = np.add.outer(alpha_B, alpha_B)
        M_BB = np.multiply.outer(alpha_B, alpha_B)

        # Overlap matrix:
        S_BB = (2 * M_BB**0.5 / A_BB)**(l + 1.5)

        # Kinetic energy matrix:
        T_BB = 2**(l + 2.5) * M_BB**(0.5 * l + 0.75) / gamma(l + 1.5) * (
            gamma(l + 2.5) * M_BB / A_BB**(l + 2.5) -
            0.5 * (l + 1) * gamma(l + 1.5) / A_BB**(l + 0.5) +
            0.25 * (l + 1) * (2 * l + 1) * gamma(l + 0.5) / A_BB**(l + 0.5))

        # Derivative matrix:
        D_BB = 2**(l + 2.5) * M_BB**(0.5 * l + 0.75) / gamma(l + 1.5) * (
            0.5 * (l + 1) * gamma(l + 1) / A_BB**(l + 1) -
            gamma(l + 2) * alpha_B / A_BB**(l + 2))

        # 1/r matrix:
        K_BB = 2**(l + 2.5) * M_BB**(0.5 * l + 0.75) / gamma(l + 1.5) * (
            0.5 * gamma(l + 1) / A_BB**(l + 1))

        # Find set of linearly independent functions.
        # We have len(alpha_B) gaussians (index B) and self.nbasis
        # linearly independent functions (index b).
        s_B, U_BB = eigh(S_BB)
        self.nbasis = int((s_B > eps).sum())

        Q_Bb = np.dot(U_BB[:, -self.nbasis:],
                      np.diag(s_B[-self.nbasis:]**-0.5))

        self.T_bb = np.dot(np.dot(Q_Bb.T, T_BB), Q_Bb)
        self.D_bb = np.dot(np.dot(Q_Bb.T, D_BB), Q_Bb)
        self.K_bb = np.dot(np.dot(Q_Bb.T, K_BB), Q_Bb)

        r_g = rgd.r_g
        prefactors_B = (2 * (2 * alpha_B)**(l + 1.5) / gamma(l + 1.5))**0.5
        with seterr(divide='ignore', invalid='ignore'):
            # Avoid errors division by zero when l<0.0:
            gaussians_Bg = np.exp(-np.outer(alpha_B, r_g**2)) * r_g**l
            self.basis_bg = Q_Bb.T @ (prefactors_B[:, None] * gaussians_Bg)

    def __len__(self):
        return self.nbasis

    def expand(self, C_xb):
        return np.dot(C_xb, self.basis_bg)

    def calculate_potential_matrix(self, vr_g):
        vr2dr_g = vr_g * self.rgd.r_g * self.rgd.dr_g
        V_bb = np.inner(self.basis_bg[:, 1:],
                        self.basis_bg[:, 1:] * vr2dr_g[1:])
        return V_bb


def coefs(rgd, l, vr_g, e, scalar_relativistic=True, Z=None):
    r_g = rgd.r_g

    x0_g = 2 * (e * r_g - vr_g) * r_g
    x1_g = 2 * (l + 1) * r_g / rgd.dr_g + r_g**2 * rgd.d2gdr2()
    x2_g = r_g**2 / rgd.dr_g**2
    x = 1.0

    if scalar_relativistic:
        x = (1 + l + l**2 - (Z / c)**2)**0.5 - l
        r_g = r_g.copy()
        r_g[0] = 1.0
        v_g = vr_g / r_g
        M_g = 1 + (e - v_g) / (2 * c**2)
        kappa_g = (rgd.derivative(vr_g) - v_g) / r_g / (2 * c**2 * M_g)
        x0_g = (2 * M_g * (e * r_g - vr_g) * r_g +
                (l + x - 1) * kappa_g * r_g +
                (l + x) * (l + x - 1) - l * (l + 1))
        x1_g = (2 * (l + x) * r_g / rgd.dr_g +
                r_g**2 * rgd.d2gdr2() +
                r_g**2 * kappa_g / rgd.dr_g)

    cm1_g = x2_g - x1_g / 2
    c0_g = x0_g - 2 * x2_g
    cp1_g = x2_g + x1_g / 2

    return cm1_g, c0_g, cp1_g, x


class Channel:
    def __init__(self, l, s=0, f_n=(), basis=None):
        self.l = l
        self.s = s
        self.basis = basis

        self.C_nb = None                       # eigenvectors
        self.e_n = None                        # eigenvalues
        self.f_n = np.array(f_n, dtype=float)  # occupation numbers
        self.phi_ng = None                     # wave functions

        self.name = 'spdfg'[l]
        self.solve2ok = False

    def solve(self, vr_g):
        """Diagonalize Schrödinger equation in basis set."""
        H_bb = self.basis.calculate_potential_matrix(vr_g)
        H_bb += self.basis.T_bb
        self.e_n, C_bn = eigh(H_bb)
        self.C_nb = C_bn.T
        self.phi_ng = self.basis.expand(self.C_nb[:len(self.f_n)])

    def solve2(self, vr_g, scalar_relativistic=True, Z=None, rgd=None):
        rgd = rgd or self.basis.rgd
        r_g = rgd.r_g
        l = self.l
        u_g = rgd.empty()
        self.solve2ok = True
        for n in range(len(self.f_n)):
            e = self.e_n[n]

            if e > 0.0:
                # Skip orbitals with positive energies
                self.e_n[n] = 42
                self.phi_ng[n] = 0.0
                continue

            # Find classical turning point:
            x = vr_g * r_g + 0.5 * l * (l + 1) - e * r_g**2
            g0 = rgd.round(4.0)
            while x[g0] > 0:
                g0 -= 1

            iter = 0
            ok = False
            while True:
                du1dr, a = self.integrate_outwards(u_g, rgd, vr_g, g0, e,
                                                   scalar_relativistic, Z)
                u1 = u_g[g0]
                du2dr = self.integrate_inwards(u_g, rgd, vr_g, g0, e,
                                               scalar_relativistic, Z)
                u2 = u_g[g0]
                A = du1dr / u1 - du2dr / u2
                u_g[g0:] *= u1 / u2
                norm = rgd.integrate(u_g**2, -2) / (4 * pi)
                u_g /= norm**0.5
                a /= norm**0.5

                nodes = (u_g[:-1] * u_g[1:] < 0).sum()

                if abs(A) < 1e-7 and nodes == n:
                    ok = True
                    break

                if nodes > n:
                    e *= 1.2
                elif nodes < n:
                    e *= 0.8
                else:
                    e += 0.5 * A * u_g[g0]**2
                    if e > 0:
                        break

                iter += 1
                assert iter < 400, (n, l, e)

            if ok:
                self.e_n[n] = e
                self.phi_ng[n, 1:] = u_g[1:] / r_g[1:]
                self.phi_ng[n, 0] = a
            else:
                self.solve2ok = False

    def calculate_density(self, n=None):
        """Calculate density."""
        if n is None:
            n_g = 0.0
            for n, f in enumerate(self.f_n):
                n_g += f * self.calculate_density(n)
        else:
            n_g = self.phi_ng[n]**2 / (4 * pi)
        return n_g

    def calculate_kinetic_energy_density(self, n):
        """Calculate kinetic energy density."""
        phi_g = self.phi_ng[n]
        rgd = self.basis.rgd
        tau_g = rgd.derivative(phi_g)**2 / (8 * pi)
        if self.l > 0:
            tau_g[1:] += (self.l * (self.l + 1) *
                          (phi_g[1:] / rgd.r_g[1:])**2 / (8 * pi))
        return tau_g

    def get_eigenvalue_sum(self):
        f_n = self.f_n
        return np.dot(f_n, self.e_n[:len(f_n)])

    def integrate_outwards(self, u_g, rgd, vr_g, g0, e,
                           scalar_relativistic=True, Z=None, pt_g=None):
        l = self.l
        r_g = rgd.r_g

        cm1_g, c0_g, cp1_g, x = coefs(rgd, l, vr_g, e, scalar_relativistic, Z)

        b_g = rgd.zeros()

        if Z is not None:
            a0 = 1.0
            if scalar_relativistic:
                dadr = 2 * ((l + x - 1) * c**2 / Z - Z) / (1 + 2 * (l + x))
            else:
                dadr = -Z / (l + 1)
            a1 = a0 + dadr * r_g[1]
        else:
            assert not scalar_relativistic
            if pt_g is None:
                a0 = 1.0
            else:
                b_g[1:] = 2 * pt_g[1:] * r_g[1:]**(2 - l)
                a0 = pt_g[1] / r_g[1]**l / (vr_g[1] / r_g[1] - e)
            a1 = a0

        if 0:
            u_g[0] = 0.0
            g = 1
            agm1 = a0
            ag = a1
            while True:
                u_g[g] = ag * r_g[g]**(l + x)
                agp1 = -(agm1 * cm1_g[g] + ag * c0_g[g] + b_g[g]) / cp1_g[g]
                if g == g0:
                    break
                g += 1
                agm1 = ag
                ag = agp1
        else:
            a_g = np.zeros_like(u_g)
            a_g[:2] = (a0, a1)
            cgpaw.integrate_outwards(g0, cm1_g, c0_g, cp1_g, b_g, a_g)
            u_g[:g0 + 1] = a_g[:g0 + 1] * r_g[:g0 + 1]**(l + x)
            g = g0
            agm1, ag, agp1 = a_g[g - 1:g + 2]

        r = r_g[g0]
        dr = rgd.dr_g[g0]
        da = 0.5 * (agp1 - agm1)
        dudr = (l + x) * r**(l + x - 1) * ag + r**(l + x) * da / dr

        if l - 1 + x < 0:
            phi0 = a0 * (r_g[1] * 0.1)**(l - 1 + x)
        else:
            phi0 = a0 * 0.0**(l - 1 + x)

        return dudr, phi0

    def integrate_inwards(self, u_g, rgd, vr_g, g0, e,
                          scalar_relativistic=True, Z=None, gmax=None):
        l = self.l
        r_g = rgd.r_g

        cm1_g, c0_g, cp1_g, x = coefs(rgd, l, vr_g, e, scalar_relativistic, Z)

        cm1_g[:g0] = 1.0  # prevent division by zero
        c0_g /= -cm1_g
        cp1_g /= -cm1_g

        if gmax is None:
            gmax = len(u_g)

        g = gmax - 2
        agp1 = 1.0
        u_g[gmax - 1] = agp1 * r_g[gmax - 1]**(l + x)
        with np.errstate(over='raise'):
            try:
                ag = np.exp(-(-2 * e)**0.5 * (r_g[gmax - 2] - r_g[gmax - 1]))
            except FloatingPointError:
                ag = 2e50

        if 0:
            while True:
                u_g[g] = ag * r_g[g]**(l + x)
                if ag > 1e50:
                    u_g[g:] /= 1e50
                    ag = ag / 1e50
                    agp1 = agp1 / 1e50
                agm1 = agp1 * cp1_g[g] + ag * c0_g[g]
                if g == g0:
                    break
                g -= 1
                agp1 = ag
                ag = agm1
        else:
            a_g = np.zeros_like(u_g)
            a_g[g:g + 2] = (ag, agp1)
            cgpaw.integrate_inwards(g, g0, c0_g, cp1_g, a_g)
            u_g[g0:g + 2] = a_g[g0:g + 2] * r_g[g0:g + 2]**(l + x)
            g = g0
            agm1, ag, agp1 = a_g[g - 1:g + 2]

        # print(agm1, ag, agp1, u_g[g - 1:]);asdfg
        r = r_g[g]
        dr = rgd.dr_g[g]
        da = 0.5 * (agp1 - agm1)
        dudr = (l + x) * r**(l + x - 1) * ag + r**(l + x) * da / dr

        return dudr


class DiracChannel(Channel):
    def __init__(self, k, f_n, basis):
        l = (abs(2 * k + 1) - 1) // 2
        super().__init__(l, 0, f_n, basis)
        self.k = k
        self.j = abs(k) - 0.5
        self.c_nb = None  # eigenvectors (small component)

        self.name += '(%d/2)' % (2 * self.j)

    def solve(self, vr_g):
        """Solve Dirac equation in basis set."""
        nb = len(self.basis)
        V_bb = self.basis.calculate_potential_matrix(vr_g)
        H_bb = np.zeros((2 * nb, 2 * nb))
        H_bb[:nb, :nb] = V_bb
        H_bb[nb:, nb:] = V_bb - 2 * c**2 * np.eye(nb)
        H_bb[nb:, :nb] = -c * (-self.basis.D_bb.T + self.k * self.basis.K_bb)
        e_n, C_bn = eigh(H_bb)
        if self.k < 0:
            n0 = nb
        else:
            n0 = nb + 1
        self.e_n = e_n[n0:].copy()
        self.C_nb = C_bn[:nb, n0:].T.copy()  # large component
        self.c_nb = C_bn[nb:, n0:].T.copy()  # small component

    def calculate_density(self, n=None):
        """Calculate density."""
        if n is None:
            n_g = Channel.calculate_density(self)
        else:
            n_g = (self.basis.expand(self.C_nb[n])**2 +
                   self.basis.expand(self.c_nb[n])**2) / (4 * pi)
            if self.basis.l < 0:
                n_g[0] = n_g[1]
        return n_g


class AllElectronAtom:
    def __init__(self, symbol, xc='LDA', spinpol=False, dirac=False,
                 configuration=None,
                 ee_interaction=True,
                 Z=None,
                 log=None,
                 scalar_relativistic=True):
        """All-electron calculation for spherically symmetric atom.

        symbol: str (or int)
            Chemical symbol (or atomic number).
        xc: str
            Name of XC-functional.
        spinpol: bool
            If true, do spin-polarized calculation.  Default is spin-paired.
        dirac: bool
            Solve Dirac equation instead of Schrödinger equation.
        configuration: list
            Electronic configuration for symbol, format as in
            gpaw.atom.configurations
        ee_interaction: bool
            Use ee_interaction=False to turn off electron-electron interaction.
            Default is True.
        log: stream
            Text output."""

        if isinstance(symbol, int):
            symbol = chemical_symbols[symbol]
        self.symbol = symbol
        self.Z = Z or atomic_numbers[symbol]

        self.nspins = 1 + int(bool(spinpol))

        self.dirac = bool(dirac)

        self.ee_interaction = ee_interaction

        if configuration is not None:
            self.configuration = copy.deepcopy(configuration)
        else:
            self.configuration = None

        self.scalar_relativistic = bool(scalar_relativistic)

        if isinstance(xc, str):
            self.xc = XC(xc)
        else:
            self.xc = xc

        self.fd = log or sys.stdout

        self.vr_sg = None  # potential * r
        self.n_sg = 0.0  # density
        self.rgd = None  # radial grid descriptor

        # Energies:
        self.ekin = None
        self.eeig = None
        self.eH = None
        self.eZ = None

        self.channels = None

        self.initialize_configuration(self.configuration)

        self.log('Z:              ', self.Z)
        self.log('Name:           ', atomic_names[atomic_numbers[symbol]])
        self.log('Symbol:         ', symbol)
        self.log('XC-functional:  ', self.xc.name)
        self.log('Equation:       ', ['Schrödinger', 'Dirac'][self.dirac])

        self.method = 'Gaussian basis-set'

    def log(self, *args, **kwargs):
        print(file=self.fd, *args, **kwargs)

    def initialize_configuration(self, configuration=None):
        self.f_lsn = {}

        if configuration is None:
            configs = configurations[self.symbol][1]
        elif isinstance(configuration, str):
            configs = []
            if configuration[0] == '[':
                symbol, configuration = configuration[1:].split(']')
                configs = configurations[symbol][1]
            for nlf in configuration.split(','):
                configs.append((int(nlf[0]),
                                'spdfg'.index(nlf[1]),
                                int(nlf[2:]),
                                -1.0))
        else:
            configs = configuration

        for n, l, f, e in configs:

            if l not in self.f_lsn:
                self.f_lsn[l] = [[] for s in range(self.nspins)]
            if self.nspins == 1:
                self.f_lsn[l][0].append(f)
            else:
                # Use Hund's rule:
                f0 = min(f, 2 * l + 1)
                self.f_lsn[l][0].append(f0)
                self.f_lsn[l][1].append(f - f0)

        if 0:
            n = 2 + len(self.f_lsn[2][0])
            if self.f_lsn[0][0][n] == 2:
                self.f_lsn[0][0][n] = 1
                self.f_lsn[2][0][n - 3] += 1

    def add(self, n, l, df=+1, s=None):
        """Add (remove) electrons."""
        if s is None:
            if self.nspins == 1:
                s = 0
            else:
                self.add(n, l, 0.5 * df, 0)
                self.add(n, l, 0.5 * df, 1)
                return

        if l not in self.f_lsn:
            self.f_lsn[l] = [[] for x in range(self.nspins)]

        f_n = self.f_lsn[l][s]
        if len(f_n) < n - l:
            f_n.extend([0] * (n - l - len(f_n)))
        f_n[n - l - 1] += df

    def initialize(self, ngpts=2000, rcut=50.0,
                   alpha1=0.01, alpha2=None, ngauss=50,
                   eps=1.0e-7):
        """Initialize basis sets and radial grid.

        ngpts: int
            Number of grid points for radial grid.
        rcut: float
            Cutoff for radial grid.
        alpha1: float
            Smallest exponent for gaussian.
        alpha2: float
            Largest exponent for gaussian.
        ngauss: int
            Number of gaussians.
        eps: float
            Cutoff for eigenvalues of overlap matrix."""

        if alpha2 is None:
            alpha2 = 50.0 * self.Z**2

        if 1:
            # Use grid with r(0)=0, r(1)=a and r(ngpts)=rcut:
            a = 1 / alpha2**0.5 / 20
            b = (rcut - a * ngpts) / (rcut * ngpts)
            b = 1 / round(1 / b)
            self.rgd = AERadialGridDescriptor(a, b, ngpts)
        else:
            from scipy.optimize import root
            rT = self.Z / 137**2
            r1 = rT / 10 / 5
            sol = root(lambda d: r1 / d * (np.exp(d * (ngpts - 1)) - 1) - rcut,
                       0.1)
            d = sol.x[0]
            a = r1 / d
            self.rgd = AbinitRadialGridDescriptor(a, d, ngpts)

        self.log('Grid points:     %d (%.5f, %.5f, %.5f, ..., %.3f, %.3f)' %
                 ((self.rgd.N,) + tuple(self.rgd.r_g[[0, 1, 2, -2, -1]])))

        # Distribute exponents between alpha1 and alpha2:
        alpha_B = alpha1 * (alpha2 / alpha1)**np.linspace(0, 1, ngauss)
        self.log('Exponents:       %d (%.3f, %.3f, ..., %.3f, %.3f)' %
                 ((ngauss,) + tuple(alpha_B[[0, 1, -2, -1]])))

        # Maximum l value:
        lmax = max(self.f_lsn.keys())

        self.channels = []
        nb_l = []
        if not self.dirac:
            for l in range(lmax + 1):
                basis = GaussianBasis(l, alpha_B, self.rgd, eps)
                nb_l.append(len(basis))
                for s in range(self.nspins):
                    self.channels.append(Channel(l, s, self.f_lsn[l][s],
                                                 basis))
        else:
            for K in range(1, lmax + 2):
                leff = (K**2 - (self.Z / c)**2)**0.5 - 1
                basis = GaussianBasis(leff, alpha_B, self.rgd, eps)
                nb_l.append(len(basis))
                for k, l in [(-K, K - 1), (K, K)]:
                    if l > lmax:
                        continue
                    f_n = self.f_lsn[l][0]
                    j = abs(k) - 0.5
                    f_n = (2 * j + 1) / (4 * l + 2) * np.array(f_n)
                    self.channels.append(DiracChannel(k, f_n, basis))

        self.log('Basis functions: %s (%s)' %
                 (', '.join([str(nb) for nb in nb_l]),
                  ', '.join('spdf'[:lmax + 1])))

        self.vr_sg = self.rgd.zeros(self.nspins)
        self.vr_sg[:] = -self.Z

    def solve(self):
        """Diagonalize Schrödinger equation."""
        self.eeig = 0.0
        for channel in self.channels:
            if self.method == 'Gaussian basis-set':
                channel.solve(self.vr_sg[channel.s])
            else:
                channel.solve2(self.vr_sg[channel.s], self.scalar_relativistic,
                               self.Z)

            self.eeig += channel.get_eigenvalue_sum()

    def calculate_density(self):
        """Calculate elctron density and kinetic energy."""
        self.n_sg = self.rgd.zeros(self.nspins)
        for channel in self.channels:
            self.n_sg[channel.s] += channel.calculate_density()

    def calculate_electrostatic_potential(self):
        """Calculate electrostatic potential and energy."""
        n_g = self.n_sg.sum(0)
        if self.ee_interaction:
            self.vHr_g = self.rgd.poisson(n_g)
        else:
            self.vHr_g = self.rgd.zeros()
        self.eH = 0.5 * self.rgd.integrate(n_g * self.vHr_g, -1)
        self.eZ = -self.Z * self.rgd.integrate(n_g, -1)

    def calculate_xc_potential(self):
        self.vxc_sg = self.rgd.zeros(self.nspins)
        self.exc = self.xc.calculate_spherical(self.rgd, self.n_sg,
                                               self.vxc_sg)

    def step(self):
        self.solve()
        self.calculate_density()
        self.calculate_electrostatic_potential()
        self.calculate_xc_potential()
        self.vr_sg = self.vxc_sg * self.rgd.r_g
        self.vr_sg += self.vHr_g
        self.vr_sg -= self.Z
        self.ekin = (self.eeig -
                     self.rgd.integrate((self.vr_sg * self.n_sg).sum(0), -1))

    def run(self, mix=0.4, maxiter=117, dnmax=1e-9):
        if self.channels is None:
            self.initialize()

        if self.dirac:
            equation = 'Dirac'
        elif self.scalar_relativistic:
            equation = 'scalar-relativistic Schrödinger'
        else:
            equation = 'non-relativistic Schrödinger'
        self.log(f'\nSolving {equation} equation using {self.method}:')

        dn = self.Z

        vr_old_sg = None
        n_old_sg = None
        for iter in range(maxiter):
            self.log('.', end='')
            self.fd.flush()
            if iter > 0:
                self.vr_sg *= mix
                self.vr_sg += (1 - mix) * vr_old_sg
                dn = self.rgd.integrate(abs(self.n_sg - n_old_sg).sum(0))
                if dn <= dnmax:
                    self.log('\nConverged in', iter, 'steps')
                    break

            vr_old_sg = self.vr_sg
            n_old_sg = self.n_sg
            self.step()

        self.summary()

        if self.method != 'Gaussian basis-set':
            for channel in self.channels:
                assert channel.solve2ok

        if dn > dnmax:
            raise RuntimeError('Did not converge!')

    def refine(self):
        self.method = 'finite difference'
        self.run(dnmax=1e-6, mix=0.14, maxiter=200)

    def summary(self):
        self.write_states()
        self.write_energies()

    def write_states(self):
        self.log('\n state  occupation         eigenvalue          <r>')
        if self.dirac:
            self.log(' nl(j)               [Hartree]        [eV]    [Bohr]')
        else:
            self.log(' nl                  [Hartree]        [eV]    [Bohr]')
        self.log('-----------------------------------------------------')
        states = []
        for ch in self.channels:
            for n, f in enumerate(ch.f_n):
                states.append((ch.e_n[n], n, ch.s, ch))
        states.sort()
        for e, n, s, ch in states:
            name = str(n + ch.l + 1) + ch.name
            if self.nspins == 2:
                name += '(%s)' % '+-'[s]
            n_g = ch.calculate_density(n)
            rave = self.rgd.integrate(n_g, 1)
            self.log(' %-7s  %6.3f %13.6f  %13.5f %6.3f' %
                     (name, ch.f_n[n], e, e * units.Hartree, rave))

    def write_energies(self):
        self.log('\nEnergies:          [Hartree]           [eV]')
        self.log('--------------------------------------------')
        for text, e in [('kinetic      ', self.ekin),
                        ('coulomb (e-e)', self.eH),
                        ('coulomb (e-n)', self.eZ),
                        ('xc           ', self.exc),
                        ('total        ',
                         self.ekin + self.eH + self.eZ + self.exc)]:
            self.log(f' {text} {e:+13.6f}  {units.Hartree * e:+13.5f}')

        if not self.dirac:
            self.calculate_exx()
            self.log(f'\nExact exchange energy: {self.exx:.6f} Hartree, '
                     f'{self.exx * units.Hartree:.5f} eV')

    def get_channel(self, l=None, s=0, k=None):
        if self.dirac:
            for channel in self.channels:
                if channel.k == k:
                    return channel
        else:
            for channel in self.channels:
                if channel.l == l and channel.s == s:
                    return channel
        raise ValueError

    def get_orbital(self, n, l=None, s=0, k=None):
        channel = self.get_channel(l, s, k)
        return channel.basis.expand(channel.C_nb[n])

    def plot_wave_functions(self, rc=4.0, show=True):
        import matplotlib.pyplot as plt
        for ch in self.channels:
            for n in range(len(ch.f_n)):
                fr_g = ch.basis.expand(ch.C_nb[n]) * self.rgd.r_g
                # fr_g = ch.phi_ng[n]
                name = str(n + ch.l + 1) + ch.name
                lw = 2
                if self.nspins == 2:
                    name += '(%s)' % '+-'[ch.s]
                    if ch.s == 1:
                        lw = 1
                if self.dirac and ch.k > 0:
                    lw = 1
                ls = ['-', '--', '-.', ':'][ch.l]
                n_g = ch.calculate_density(n)
                rave = self.rgd.integrate(n_g, 1)
                gave = self.rgd.round(rave)
                fr_g *= np.sign(fr_g[gave])
                plt.plot(self.rgd.r_g, fr_g,
                         ls=ls, lw=lw, color=colors[n + ch.l], label=name)

        plt.legend(loc='best')
        plt.xlabel('r [Bohr]')
        plt.ylabel('$r\\phi(r)$')
        plt.axis(xmin=0, xmax=rc)
        if show:
            plt.show()

    def logarithmic_derivative(self, l, energies, rcut):
        ch = Channel(l)
        gcut = self.rgd.round(rcut)
        u_g = self.rgd.empty()
        logderivs = []
        d0 = 42.0
        offset = 0
        for e in energies:
            dudr = ch.integrate_outwards(u_g, self.rgd, self.vr_sg[0],
                                         gcut, e, self.scalar_relativistic,
                                         self.Z)[0]
            d1 = np.arctan(dudr / u_g[gcut]) / pi + offset
            if d1 > d0:
                offset -= 1
                d1 -= 1
            logderivs.append(d1)
            d0 = d1

        return np.array(logderivs)

    def calculate_exx(self, s=None):
        if s is None:
            self.exx = sum(self.calculate_exx(s)
                           for s in range(self.nspins)) / self.nspins
            return self.exx

        states = []
        lmax = 0
        for ch in self.channels:
            l = ch.l
            for n, phi_g in enumerate(ch.phi_ng):
                f = ch.f_n[n]
                if f > 0 and ch.s == s:
                    states.append((l, f * self.nspins / 2.0 / (2 * l + 1),
                                   phi_g))
                    if l > lmax:
                        lmax = l

        G_LLL = gaunt(lmax)

        exx = 0.0
        j1 = 0
        for l1, f1, phi1_g in states:
            f = 1.0
            for l2, f2, phi2_g in states[j1:]:
                n_g = phi1_g * phi2_g
                for l in range((l1 + l2) % 2, l1 + l2 + 1, 2):
                    G = (G_LLL[l1**2:(l1 + 1)**2,
                               l2**2:(l2 + 1)**2,
                               l**2:(l + 1)**2]**2).sum()
                    vr_g = self.rgd.poisson(n_g, l)
                    e = f * self.rgd.integrate(vr_g * n_g, -1) / 4 / pi
                    exx -= e * G * f1 * f2
                f = 2.0
            j1 += 1

        return exx


class CLICommand:
    """Solve radial equation for an atom.

    Example:

        gpaw atom Li -f PBE -p  # plot wave functions for a lithium atom
    """

    @staticmethod
    def add_arguments(parser):
        add = parser.add_argument
        add('symbol')
        add('-f', '--xc-functional', type=str, default='LDA',
            help='Exchange-Correlation functional ' +
            '(default value LDA)',
            metavar='<XC>')
        add('-a', '--add', metavar='states',
            help='Add electron(s). Use "1s0.5a" to add 0.5 1s ' +
            'electrons to the alpha-spin channel (use "b" for ' +
            'beta-spin).  The number of electrons defaults to ' +
            'one. Examples: "1s", "2p2b", "4f0.1b,3d-0.1a".')
        add('--spin-polarized', action='store_true')
        add('-d', '--dirac', action='store_true',
            help='Solve Dirac equation.')
        add('-p', '--plot', action='store_true')
        add('-e', '--exponents',
            help='Exponents a: exp(-a*r^2).  Use "-e 0.1:20.0:30" ' +
            'to get 30 exponents from 0.1 to 20.0.')
        add('-l', '--logarithmic-derivatives',
            metavar='spdfg,e1:e2:de,radius',
            help='Plot logarithmic derivatives. ' +
            'Example: -l spdf,-1:1:0.05,1.3. ' +
            'Energy range and/or radius can be left out.')
        add('-n', '--ngrid', help='Specify number of grid points.')
        add('-R', '--rcut', help='Radial cutoff.')
        add('-r', '--refine', action='store_true')
        add('-s', '--scalar-relativistic',
            action='store_const', const=True,
            help='Use scalar-relativistic setups (default).')
        add('-S', '--non-relativistic',
            action='store_const', const=False, dest='scalar_relativistic',
            help='Don\'t use non-relativistic setups.')
        add('--no-ee-interaction', action='store_true',
            help='Turn off electron-electron interaction.')

    @staticmethod
    def run(args):
        main(args)


def parse_ld_str(s, energies=None, r=2.0):
    parts = s.split(',')
    lvalues = ['spdfg'.find(x) for x in parts.pop(0)]
    if parts:
        e1, e2, de = (float(x) for x in parts.pop(0).split(':'))
    else:
        e1, e2, de = energies
    if parts:
        r = float(parts.pop())
    energies = np.linspace(e1, e2, int((e2 - e1) / de) + 1)
    return lvalues, energies, r


def main(args):
    symbol = args.symbol

    nlfs = []
    if args.add:
        for x in args.add.split(','):
            n = int(x[0])
            l = 'spdfg'.find(x[1])
            x = x[2:]
            if x and x[-1] in 'ab':
                s = int(x[-1] == 'b')
                args.spin_polarized = True
                x = x[:-1]
            else:
                s = None
            if x:
                f = float(x)
            else:
                f = 1
            nlfs.append((n, l, f, s))

    aea_kwargs = dict(xc=args.xc_functional,
                      spinpol=args.spin_polarized,
                      dirac=args.dirac,
                      ee_interaction=not args.no_ee_interaction)
    if args.scalar_relativistic is not None:
        aea_kwargs['scalar_relativistic'] = args.scalar_relativistic
    aea = AllElectronAtom(symbol, **aea_kwargs)

    kwargs = {}
    if args.exponents:
        parts = args.exponents.split(':')
        kwargs['alpha1'] = float(parts[0])
        if len(parts) > 1:
            kwargs['alpha2'] = float(parts[1])
            if len(parts) > 2:
                kwargs['ngauss'] = int(parts[2])

    for n, l, f, s in nlfs:
        aea.add(n, l, f, s)

    if args.ngrid:
        kwargs['ngpts'] = int(args.ngrid)
    if args.rcut:
        kwargs['rcut'] = float(args.rcut)

    aea.initialize(**kwargs)
    aea.run()

    if args.refine or args.scalar_relativistic:
        aea.refine()

    if args.logarithmic_derivatives:
        lvalues, energies, r = parse_ld_str(args.logarithmic_derivatives,
                                            (-1, 1, 0.05))
        import matplotlib.pyplot as plt
        for l in lvalues:
            ld = aea.logarithmic_derivative(l, energies, r)
            plt.plot(energies, ld, colors[l])
        plt.show()

    if args.plot:
        aea.plot_wave_functions()
