"""This module defines a Hamiltonian."""

import functools

import numpy as np
from ase.units import Ha

from gpaw.arraydict import ArrayDict
from gpaw.external import create_external_potential
from gpaw.lfc import LFC
from gpaw.poisson import PoissonSolver
from gpaw.spinorbit import soc
from gpaw.transformers import Transformer
from gpaw.utilities import (pack2, pack_atomic_matrices, unpack,
                            unpack_atomic_matrices)
from gpaw.utilities.partition import AtomPartition

ENERGY_NAMES = ['e_kinetic', 'e_coulomb', 'e_zero', 'e_external', 'e_xc',
                'e_entropy', 'e_total_free', 'e_total_extrapolated']


def apply_non_local_hamilton(dH_asp, collinear, P, out=None):
    if out is None:
        out = P.new()
    for a, I1, I2 in P.indices:
        if collinear:
            dH_ii = unpack(dH_asp[a][P.spin])
            out.array[:, I1:I2] = np.dot(P.array[:, I1:I2], dH_ii)
        else:
            dH_xp = dH_asp[a]
            # We need the transpose because
            # we are dotting from the left
            dH_ii = unpack(dH_xp[0]).T
            dH_vii = [unpack(dH_p).T for dH_p in dH_xp[1:]]
            out.array[:, 0, I1:I2] = (np.dot(P.array[:, 0, I1:I2],
                                             dH_ii + dH_vii[2]) +
                                      np.dot(P.array[:, 1, I1:I2],
                                             dH_vii[0] - 1j * dH_vii[1]))
            out.array[:, 1, I1:I2] = (np.dot(P.array[:, 1, I1:I2],
                                             dH_ii - dH_vii[2]) +
                                      np.dot(P.array[:, 0, I1:I2],
                                             dH_vii[0] + 1j * dH_vii[1]))
    return out


# from gpaw.utilities.debug import frozen
# @frozen
class Hamiltonian:

    def __init__(self, gd, finegd, nspins, collinear, setups, timer, xc, world,
                 redistributor, vext=None):
        self.gd = gd
        self.finegd = finegd
        self.nspins = nspins
        self.collinear = collinear
        self.setups = setups
        self.timer = timer
        self.xc = xc
        self.world = world
        self.redistributor = redistributor

        self.ncomponents = self.nspins if self.collinear else 1 + 3

        self.atomdist = None
        self.dH_asp = None
        self.vt_xG = None
        self.vt_sG = None
        self.vt_vG = None
        self.vHt_g = None
        self.vt_xg = None
        self.vt_sg = None
        self.vt_vg = None
        self.atom_partition = None

        # Energy contributioons that sum up to e_total_free:
        self.e_kinetic = None
        self.e_coulomb = None
        self.e_zero = None
        self.e_external = None
        self.e_xc = None
        self.e_entropy = None
        self.e_band = None

        self.e_total_free = None
        self.e_total_extrapolated = None
        self.e_kinetic0 = None

        self.ref_vt_sG = None
        self.ref_dH_asp = None

        if isinstance(vext, dict):
            vext = create_external_potential(**vext)
        self.vext = vext  # external potential

        self.positions_set = False
        self.spos_ac = None
        self.soc = False

    @property
    def dH(self):
        return functools.partial(apply_non_local_hamilton,
                                 self.dH_asp, self.collinear)

    def update_atomic_hamiltonians(self, value):
        if isinstance(value, dict):
            dtype = complex if self.soc else float
            tmp = self.setups.empty_atomic_matrix(self.ncomponents,
                                                  self.atom_partition, dtype)
            tmp.update(value)
            value = tmp
        assert isinstance(value, ArrayDict) or value is None, type(value)
        if value is not None:
            value.check_consistency()
        self.dH_asp = value

    def __str__(self):
        s = 'Hamiltonian:\n'
        s += ('  XC and Coulomb potentials evaluated on a {0}*{1}*{2} grid\n'
              .format(*self.finegd.N_c))
        s += '  Using the %s Exchange-Correlation functional\n' % self.xc.name
        # We would get the description of the XC functional here,
        # except the thing has probably not been fully initialized yet.
        if self.vext is not None:
            s += '  External potential:\n    {0}\n'.format(self.vext)
        return s

    def summary(self, wfs, log):
        log('Energy contributions relative to reference atoms:',
            '(reference = {0:.6f})\n'.format(self.setups.Eref * Ha))

        energies = [('Kinetic:      ', self.e_kinetic),
                    ('Potential:    ', self.e_coulomb),
                    ('External:     ', self.e_external),
                    ('XC:           ', self.e_xc),
                    ('Entropy (-ST):', self.e_entropy),
                    ('Local:        ', self.e_zero)]

        for name, e in energies:
            log('%-14s %+11.6f' % (name, Ha * e))

        log('--------------------------')
        log('Free energy:   %+11.6f' % (Ha * self.e_total_free))
        log('Extrapolated:  %+11.6f' % (Ha * self.e_total_extrapolated))
        log()
        self.xc.summary(log)

        try:
            workfunctions = self.get_workfunctions(wfs)
        except ValueError:
            pass
        else:
            log('Dipole-layer corrected work functions: {:.6f}, {:.6f} eV'
                .format(*np.array(workfunctions) * Ha))
            log()

    def get_workfunctions(self, wfs):
        """
        Returns the work functions, in Hartree, for a dipole-corrected
        simulation. Returns None if no dipole correction is present.
        (wfs can be obtained from calc.wfs)
        """
        try:
            dipole_correction = self.poisson.correction
        except AttributeError:
            raise ValueError(
                'Work function not defined if no field-free region. Consider '
                'using a dipole correction if you are looking for a '
                'work function.')
        c = self.poisson.c  # index of axis perpendicular to dipole-layer
        if not self.gd.pbc_c[c]:
            # zero boundary conditions
            vacuum = 0.0
        else:
            v_q = self.pd3.gather(self.vHt_q)
            if self.pd3.comm.rank == 0:
                axes = (c, (c + 1) % 3, (c + 2) % 3)
                v_g = self.pd3.ifft(v_q, local=True).transpose(axes)
                vacuum = v_g[0].mean()
            else:
                vacuum = np.nan

        fermilevel = wfs.fermi_level
        wf1 = vacuum - fermilevel + dipole_correction
        wf2 = vacuum - fermilevel - dipole_correction
        return np.array([wf1, wf2])

    def set_positions_without_ruining_everything(self, spos_ac,
                                                 atom_partition):
        self.spos_ac = spos_ac
        rank_a = atom_partition.rank_a

        # If both old and new atomic ranks are present, start a blank dict if
        # it previously didn't exist but it will needed for the new atoms.
        # XXX what purpose does this serve?  In what case does it happen?
        # How would one even go about figuring it out?  Why does it all have
        # to be so unreadable? -Ask
        #
        if (self.atom_partition is not None and
            self.dH_asp is None and (rank_a == self.gd.comm.rank).any()):
            self.update_atomic_hamiltonians({})

        if self.atom_partition is not None and self.dH_asp is not None:
            self.timer.start('Redistribute')
            self.dH_asp.redistribute(atom_partition)
            self.timer.stop('Redistribute')

        self.atom_partition = atom_partition
        self.atomdist = self.redistributor.get_atom_distributions(spos_ac)

    def set_positions(self, spos_ac, atom_partition):
        self.vbar.set_positions(spos_ac, atom_partition)
        self.xc.set_positions(spos_ac)
        self.set_positions_without_ruining_everything(spos_ac, atom_partition)
        self.positions_set = True

    def initialize(self):
        self.vt_xg = self.finegd.empty(self.ncomponents)
        self.vt_sg = self.vt_xg[:self.nspins]
        self.vt_vg = self.vt_xg[self.nspins:]
        self.vHt_g = self.finegd.zeros()
        self.vt_xG = self.gd.empty(self.ncomponents)
        self.vt_sG = self.vt_xG[:self.nspins]
        self.vt_vG = self.vt_xG[self.nspins:]

    def update(self, density, wfs=None, kin_en_using_band=True):
        """Calculate effective potential.

        The XC-potential and the Hartree potential are evaluated on
        the fine grid, and the sum is then restricted to the coarse
        grid."""

        self.timer.start('Hamiltonian')
        if self.vt_sg is None:
            with self.timer('Initialize Hamiltonian'):
                self.initialize()

        finegrid_energies = self.update_pseudo_potential(density)
        coarsegrid_e_kinetic = self.calculate_kinetic_energy(density)

        with self.timer('Calculate atomic Hamiltonians'):
            W_aL = self.calculate_atomic_hamiltonians(density)

        atomic_energies = self.update_corrections(density, W_aL)

        # Make energy contributions summable over world:
        finegrid_energies *= self.finegd.comm.size / self.world.size
        coarsegrid_e_kinetic *= self.gd.comm.size / self.world.size
        # (careful with array orderings/contents)

        if 0:
            print('kinetic', atomic_energies[0], coarsegrid_e_kinetic)
            print('coulomb', atomic_energies[1], finegrid_energies[0])
            print('zero', atomic_energies[2], finegrid_energies[1])
            print('xc', atomic_energies[4], finegrid_energies[3])
            print('external', atomic_energies[3], finegrid_energies[2])

        energies = atomic_energies  # kinetic, coulomb, zero, external, xc
        energies[1:] += finegrid_energies  # coulomb, zero, external, xc
        energies[0] += coarsegrid_e_kinetic  # kinetic

        with self.timer('Communicate'):  # time possible load imbalance
            self.world.sum(energies)
        if not kin_en_using_band:
            assert wfs is not None
            with self.timer('New Kinetic Energy'):
                energies[0] = \
                    self.calculate_kinetic_energy_directly(density,
                                                           wfs)

        (self.e_kinetic0, self.e_coulomb, self.e_zero,
         self.e_external, self.e_xc) = energies

        self.timer.stop('Hamiltonian')

    def update_corrections(self, dens, W_aL):
        self.timer.start('Atomic')
        self.update_atomic_hamiltonians(None)  # XXXX

        e_kinetic = 0.0
        e_coulomb = 0.0
        e_zero = 0.0
        e_external = 0.0
        e_xc = 0.0

        D_asp = self.atomdist.to_work(dens.D_asp)
        dtype = complex if self.soc else float
        dH_asp = self.setups.empty_atomic_matrix(self.ncomponents,
                                                 D_asp.partition, dtype)

        for a, D_sp in D_asp.items():
            W_L = W_aL[a]
            setup = self.setups[a]

            if self.nspins == 2:
                D_p = D_sp.sum(0)
            else:
                D_p = D_sp[0]
            dH_p = (setup.K_p + setup.M_p +
                    setup.MB_p + 2.0 * np.dot(setup.M_pp, D_p) +
                    np.dot(setup.Delta_pL, W_L))
            e_kinetic += np.dot(setup.K_p, D_p) + setup.Kc
            e_zero += setup.MB + np.dot(setup.MB_p, D_p)
            e_coulomb += setup.M + np.dot(D_p, (setup.M_p +
                                                np.dot(setup.M_pp, D_p)))

            if self.soc:
                dH_vii = soc(setup, self.xc, D_sp)
                dH_sp = np.zeros_like(D_sp, dtype=complex)
                dH_sp[1:] = pack2(dH_vii)
            else:
                dH_sp = np.zeros_like(D_sp)

            if setup.hubbard_u is not None:
                eU, dHU_sp = setup.hubbard_u.calculate(setup, D_sp)
                e_xc += eU
                dH_sp += dHU_sp

            dH_sp[:self.nspins] += dH_p

            if self.vext is not None:
                self.vext.paw_correction(setup.Delta_pL[:, 0], dH_sp)

            if self.vext and self.vext.get_name() == 'CDFTPotential':
                # cDFT atomic hamiltonian, eq. 25
                # energy correction added in cDFT main
                h_cdft_a, h_cdft_b = self.vext.get_atomic_hamiltonians(
                    setups=setup.Delta_pL[:, 0], atom=a)

                dH_sp[0] += h_cdft_a
                dH_sp[1] += h_cdft_b

            if self.ref_dH_asp:
                assert self.collinear
                dH_sp += self.ref_dH_asp[a]

            dH_asp[a] = dH_sp

        self.timer.start('XC Correction')
        for a, D_sp in D_asp.items():
            e_xc += self.xc.calculate_paw_correction(self.setups[a], D_sp,
                                                     dH_asp[a], a=a)
        self.timer.stop('XC Correction')
        for a, D_sp in D_asp.items():
            e_kinetic -= (D_sp * dH_asp[a]).sum().real

        self.update_atomic_hamiltonians(self.atomdist.from_work(dH_asp))
        self.timer.stop('Atomic')

        # Make corrections due to non-local xc:
        # self.Enlxc = 0.0  # XXXxcfunc.get_non_local_energy()
        e_kinetic += self.xc.get_kinetic_energy_correction() / self.world.size

        return np.array([e_kinetic, e_coulomb, e_zero, e_external, e_xc])

    def get_energy(self, e_entropy, wfs, kin_en_using_band=True):
        """Sum up all eigenvalues weighted with occupation numbers"""
        self.e_band = wfs.calculate_band_energy()
        if kin_en_using_band:
            self.e_kinetic = self.e_kinetic0 + self.e_band
        else:
            self.e_kinetic = self.e_kinetic0
        self.e_entropy = e_entropy
        if 0:
            print(self.e_kinetic0,
                  self.e_band,
                  self.e_coulomb,
                  self.e_external,
                  self.e_zero,
                  self.e_xc,
                  self.e_entropy)

        self.e_total_free = (self.e_kinetic + self.e_coulomb +
                             self.e_external + self.e_zero + self.e_xc +
                             self.e_entropy)
        self.e_total_extrapolated = (
            self.e_total_free +
            wfs.occupations.extrapolate_factor * e_entropy)

        return self.e_total_free

    def linearize_to_xc(self, new_xc, density):
        # Store old hamiltonian
        ref_vt_sG = self.vt_sG.copy()
        ref_dH_asp = self.dH_asp.copy()
        self.xc = new_xc
        self.xc.set_positions(self.spos_ac)
        self.update(density)

        ref_vt_sG -= self.vt_sG
        for a, dH_sp in self.dH_asp.items():
            ref_dH_asp[a] -= dH_sp
        self.ref_vt_sG = ref_vt_sG
        self.ref_dH_asp = self.atomdist.to_work(ref_dH_asp)

    def calculate_forces(self, dens, F_av):
        ghat_aLv = dens.ghat.dict(derivative=True)
        nct_av = dens.nct.dict(derivative=True)
        vbar_av = self.vbar.dict(derivative=True)

        self.calculate_forces2(dens, ghat_aLv, nct_av, vbar_av)
        F_coarsegrid_av = np.zeros_like(F_av)

        # Force from compensation charges:
        _Q, Q_aL = dens.calculate_multipole_moments()
        for a, dF_Lv in ghat_aLv.items():
            F_av[a] += np.dot(Q_aL[a], dF_Lv)

        # Force from smooth core charge:
        for a, dF_v in nct_av.items():
            F_coarsegrid_av[a] += dF_v[0]

        # Force from zero potential:
        for a, dF_v in vbar_av.items():
            F_av[a] += dF_v[0]

        self.xc.add_forces(F_av)
        self.gd.comm.sum(F_coarsegrid_av, 0)
        self.finegd.comm.sum(F_av, 0)
        if self.vext:
            if self.vext.get_name() == 'CDFTPotential':
                F_av += self.vext.get_cdft_forces()
        F_av += F_coarsegrid_av

    def apply_local_potential(self, psit_nG, Htpsit_nG, s):
        vt_G = self.vt_sG[s]
        if psit_nG.ndim == 3:
            Htpsit_nG += psit_nG * vt_G
        else:
            for psit_G, Htpsit_G in zip(psit_nG, Htpsit_nG):
                Htpsit_G += psit_G * vt_G

    def apply(self, a_xG, b_xG, wfs, kpt, calculate_P_ani=True):
        """Apply the Hamiltonian operator to a set of vectors.

        Parameters:

        a_nG: ndarray
            Set of vectors to which the overlap operator is applied.
        b_nG: ndarray, output
            Resulting S times a_nG vectors.
        wfs: WaveFunctions
            Wave-function object defined in wavefunctions.py
        kpt: KPoint object
            k-point object defined in kpoint.py.
        calculate_P_ani: bool
            When True, the integrals of projector times vectors
            P_ni = <p_i | a_nG> are calculated.
            When False, existing P_ani are used

        """

        wfs.kin.apply(a_xG, b_xG, kpt.phase_cd)
        self.apply_local_potential(a_xG, b_xG, kpt.s)
        shape = a_xG.shape[:-3]
        P_axi = wfs.pt.dict(shape)

        if calculate_P_ani:  # TODO calculate_P_ani=False is experimental
            wfs.pt.integrate(a_xG, P_axi, kpt.q)
        else:
            for a, P_ni in kpt.P_ani.items():
                P_axi[a][:] = P_ni

        for a, P_xi in P_axi.items():
            dH_ii = unpack(self.dH_asp[a][kpt.s])
            P_axi[a] = np.dot(P_xi, dH_ii)
        wfs.pt.add(b_xG, P_axi, kpt.q)

    def get_xc_difference(self, xc, density):
        """Calculate non-selfconsistent XC-energy difference."""
        if density.nt_sg is None:
            density.interpolate_pseudo_density()
        nt_sg = density.nt_sg
        if hasattr(xc, 'hybrid'):
            xc.calculate_exx()
        finegd_e_xc = xc.calculate(density.finegd, nt_sg)
        D_asp = self.atomdist.to_work(density.D_asp)
        atomic_e_xc = 0.0
        for a, D_sp in D_asp.items():
            setup = self.setups[a]
            atomic_e_xc += xc.calculate_paw_correction(setup, D_sp, a=a)
        e_xc = finegd_e_xc + self.world.sum(atomic_e_xc)
        return e_xc - self.e_xc

    def estimate_memory(self, mem):
        nbytes = self.gd.bytecount()
        nfinebytes = self.finegd.bytecount()
        arrays = mem.subnode('Arrays', 0)
        arrays.subnode('vHt_g', nfinebytes)
        arrays.subnode('vt_sG', self.nspins * nbytes)
        arrays.subnode('vt_sg', self.nspins * nfinebytes)
        self.xc.estimate_memory(mem.subnode('XC'))
        self.poisson.estimate_memory(mem.subnode('Poisson'))
        self.vbar.estimate_memory(mem.subnode('vbar'))

    def write(self, writer):
        # Write all eneriges:
        for name in ENERGY_NAMES:
            energy = getattr(self, name)
            if energy is not None:
                energy *= Ha
            writer.write(name, energy)

        writer.write(
            potential=self.gd.collect(self.vt_xG) * Ha,
            atomic_hamiltonian_matrices=pack_atomic_matrices(self.dH_asp) * Ha)

        self.xc.write(writer.child('xc'))

        if hasattr(self.poisson, 'write'):
            self.poisson.write(writer.child('poisson'))

    def read(self, reader):
        h = reader.hamiltonian

        # Read all energies:
        for name in ENERGY_NAMES:
            energy = h.get(name)
            if energy is not None:
                energy /= reader.ha
            setattr(self, name, energy)

        # Read pseudo potential on the coarse grid
        # and broadcast on kpt/band comm:
        self.initialize()
        self.gd.distribute(h.potential / reader.ha, self.vt_xG)

        self.atom_partition = AtomPartition(self.gd.comm,
                                            np.zeros(len(self.setups), int),
                                            name='hamiltonian-init-serial')

        # Read non-local part of hamiltonian
        self.update_atomic_hamiltonians({})
        dH_sP = h.atomic_hamiltonian_matrices / reader.ha

        if self.gd.comm.rank == 0:
            self.update_atomic_hamiltonians(
                unpack_atomic_matrices(dH_sP, self.setups))

        if hasattr(self.poisson, 'read'):
            self.poisson.read(reader)
            self.poisson.set_grid_descriptor(self.finegd)

    def calculate_kinetic_energy_directly(self, density, wfs):

        """
        Calculate kinetic energy as 1/2 (nable psi)^2
        it gives better estimate of kinetic energy during the SCF.
        Important for direct min.

        'calculate_kinetic_energy' method gives a correct
        value of kinetic energy only at self-consistent solution.

        :param density:
        :param wfs:
        :return: total kinetic energy
        """
        # pseudo-part
        if wfs.mode == 'lcao':
            return self.calculate_kinetic_energy_using_kin_en_matrix(
                density, wfs)
        elif wfs.mode == 'pw':
            e_kin = 0.0
            for kpt in wfs.kpt_u:
                for f, psit_G in zip(kpt.f_n, kpt.psit_nG):
                    if f > 1.0e-10:
                        G2_G = wfs.pd.G2_qG[kpt.q]
                        e_kin += f * wfs.pd.integrate(
                            0.5 * G2_G * psit_G, psit_G).real
        else:
            e_kin = 0.0

            def Lapl(psit_G, kpt):
                Lpsit_G = np.zeros_like(psit_G)
                wfs.kin.apply(psit_G, Lpsit_G, kpt.phase_cd)
                return Lpsit_G

            for kpt in wfs.kpt_u:
                for f, psit_G in zip(kpt.f_n, kpt.psit_nG):
                    if f > 1.0e-10:
                        e_kin += f * wfs.integrate(
                            Lapl(psit_G, kpt), psit_G, False)
            e_kin = e_kin.real
            e_kin = wfs.gd.comm.sum(e_kin)

        e_kin = wfs.kd.comm.sum(e_kin)  # ?
        # paw corrections
        e_kin_paw = 0.0
        for a, D_sp in density.D_asp.items():
            setup = wfs.setups[a]
            D_p = D_sp.sum(0)
            e_kin_paw += np.dot(setup.K_p, D_p) + setup.Kc
        e_kin_paw = density.gd.comm.sum(e_kin_paw)
        return e_kin + e_kin_paw

    def calculate_kinetic_energy_using_kin_en_matrix(self, density,
                                                     wfs):
        """
        E_k = sum_{M'M} rho_MM' T_M'M
        better agreement between gradients of energy and
        the total energy during the direct minimisation.
        This is important when the line search is used.
        Also avoids using the eigenvalues which are
        not calculated during the direct minimisation.

        'calculate_kinetic_energy' method gives a correct
        value of kinetic energy only at self-consistent solution.

        :param density:
        :param wfs:
        :return: total kinetic energy
        """
        # pseudo-part
        e_kinetic = 0.0
        e_kin_paw = 0.0

        for kpt in wfs.kpt_u:
            # calculation of the density matrix directly
            # can be expansive for a large scale
            # as there are lot of empty states
            # when the exponential transformation is used
            # (n_bands=n_basis_functions.)
            #
            # rho_MM = \
            #     wfs.calculate_density_matrix(kpt.f_n, kpt.C_nM)
            # e_kinetic += np.einsum('ij,ji->', kpt.T_MM, rho_MM)
            #
            # the code below is faster
            self.timer.start('Pseudo part')
            occ = kpt.f_n > 1e-10
            x_nn = np.dot(kpt.C_nM[occ],
                          np.dot(kpt.T_MM,
                                 kpt.C_nM[occ].T.conj())).real
            e_kinetic += np.einsum('i,ii->', kpt.f_n[occ], x_nn)
            self.timer.stop('Pseudo part')
        # del rho_MM

        e_kinetic = wfs.kd.comm.sum(e_kinetic)
        # paw corrections
        for a, D_sp in density.D_asp.items():
            setup = wfs.setups[a]
            D_p = D_sp.sum(0)
            e_kin_paw += np.dot(setup.K_p, D_p) + setup.Kc

        e_kin_paw = self.gd.comm.sum(e_kin_paw)

        return e_kinetic.real + e_kin_paw


class RealSpaceHamiltonian(Hamiltonian):
    def __init__(self, gd, finegd, nspins, collinear, setups, timer, xc, world,
                 vext=None,
                 psolver=None, stencil=3, redistributor=None,
                 charge: float = 0.0):
        Hamiltonian.__init__(self, gd, finegd, nspins, collinear,
                             setups, timer, xc,
                             world, vext=vext,
                             redistributor=redistributor)

        # Solver for the Poisson equation:
        if psolver is None:
            psolver = {}
        if isinstance(psolver, dict):
            psolver = PoissonSolver(**psolver)
        self.poisson = psolver
        self.poisson.set_grid_descriptor(self.finegd)

        # Restrictor function for the potential:
        self.restrictor = Transformer(self.finegd, self.redistributor.aux_gd,
                                      stencil)
        self.restrict = self.restrictor.apply

        self.vbar = LFC(self.finegd, [[setup.vbar] for setup in setups],
                        forces=True)

        self.vbar_g = None
        self.npoisson = None

    def restrict_and_collect(self, a_xg, b_xg=None, phases=None):
        if self.redistributor.enabled:
            tmp_xg = self.restrictor.apply(a_xg, output=None, phases=phases)
            b_xg = self.redistributor.collect(tmp_xg, b_xg)
        else:
            b_xg = self.restrictor.apply(a_xg, output=b_xg, phases=phases)
        return b_xg

    def __str__(self):
        s = Hamiltonian.__str__(self)

        degree = self.restrictor.nn * 2 - 1
        name = ['linear', 'cubic', 'quintic', 'heptic'][degree // 2]
        s += ('  Interpolation: tri-%s ' % name +
              '(%d. degree polynomial)\n' % degree)
        s += '  Poisson solver: %s' % self.poisson.get_description()
        return s

    def set_positions(self, spos_ac, rank_a):
        Hamiltonian.set_positions(self, spos_ac, rank_a)
        if self.vbar_g is None:
            self.vbar_g = self.finegd.empty()
        self.vbar_g[:] = 0.0
        self.vbar.add(self.vbar_g)

    def update_pseudo_potential(self, dens):
        self.timer.start('vbar')
        e_zero = self.finegd.integrate(self.vbar_g, dens.nt_g,
                                       global_integral=False)

        vt_g = self.vt_sg[0]
        vt_g[:] = self.vbar_g
        self.timer.stop('vbar')

        e_external = 0.0
        if self.vext is not None:
            if self.vext.get_name() == 'CDFTPotential':
                vext_g = self.vext.get_potential(self.finegd).copy()
                e_external += self.vext.get_cdft_external_energy(
                    dens,
                    self.nspins, vext_g, vt_g, self.vbar_g, self.vt_sg)

            else:
                vext_g = self.vext.get_potential(self.finegd)
                vt_g += vext_g
                e_external = self.finegd.integrate(vext_g, dens.rhot_g,
                                                   global_integral=False)

        if self.nspins == 2:
            self.vt_sg[1] = vt_g

        self.timer.start('XC 3D grid')
        e_xc = self.xc.calculate(self.finegd, dens.nt_sg, self.vt_sg)
        e_xc /= self.finegd.comm.size
        self.timer.stop('XC 3D grid')

        self.timer.start('Poisson')
        # npoisson is the number of iterations:
        self.npoisson = self.poisson.solve(self.vHt_g, dens.rhot_g,
                                           charge=-dens.charge,
                                           timer=self.timer)
        self.timer.stop('Poisson')

        self.timer.start('Hartree integrate/restrict')

        e_coulomb = 0.5 * self.finegd.integrate(self.vHt_g, dens.rhot_g,
                                                global_integral=False)

        for vt_g in self.vt_sg:
            vt_g += self.vHt_g

        self.timer.stop('Hartree integrate/restrict')
        energies = np.array([e_coulomb, e_zero, e_external, e_xc])
        return energies

    def calculate_kinetic_energy(self, density):
        # XXX new timer item for kinetic energy?
        self.timer.start('Hartree integrate/restrict')
        self.restrict_and_collect(self.vt_sg, self.vt_sG)

        e_kinetic = 0.0
        s = 0
        for vt_G, nt_G in zip(self.vt_sG, density.nt_sG):
            if self.ref_vt_sG is not None:
                vt_G += self.ref_vt_sG[s]

            if s < self.nspins:
                e_kinetic -= self.gd.integrate(vt_G, nt_G - density.nct_G,
                                               global_integral=False)
            else:
                e_kinetic -= self.gd.integrate(vt_G, nt_G,
                                               global_integral=False)
            s += 1
        self.timer.stop('Hartree integrate/restrict')
        return e_kinetic

    def calculate_atomic_hamiltonians(self, dens):
        def getshape(a):
            return sum(2 * l + 1 for l, _ in enumerate(self.setups[a].ghat_l)),
        W_aL = ArrayDict(self.atomdist.aux_partition, getshape, float)
        if self.vext:
            if self.vext.get_name() != 'CDFTPotential':
                vext_g = self.vext.get_potential(self.finegd)
                dens.ghat.integrate(self.vHt_g + vext_g, W_aL)
            else:
                dens.ghat.integrate(self.vHt_g, W_aL)
        else:
            dens.ghat.integrate(self.vHt_g, W_aL)

        return self.atomdist.to_work(self.atomdist.from_aux(W_aL))

    def calculate_forces2(self, dens, ghat_aLv, nct_av, vbar_av):
        if self.nspins == 2:
            vt_G = self.vt_sG.mean(0)
        else:
            vt_G = self.vt_sG[0]

        self.vbar.derivative(dens.nt_g, vbar_av)
        if self.vext:
            if self.vext.get_name() == 'CDFTPotential':
                # CDFT force added in calculate_forces
                dens.ghat.derivative(self.vHt_g, ghat_aLv)
            else:
                vext_g = self.vext.get_potential(self.finegd)
                dens.ghat.derivative(self.vHt_g + vext_g, ghat_aLv)
        else:
            dens.ghat.derivative(self.vHt_g, ghat_aLv)
        dens.nct.derivative(vt_G, nct_av)

    def get_electrostatic_potential(self, dens):
        self.update(dens)

        v_g = self.finegd.collect(self.vHt_g, broadcast=True)
        v_g = self.finegd.zero_pad(v_g)
        if hasattr(self.poisson, 'correction'):
            assert self.poisson.c == 2
            v_g[:, :, 0] = self.poisson.correction
        return v_g
