"""Check TDDFT ionizations with Yukawa potential."""
from ase import Atoms
from ase.units import Hartree
from gpaw import GPAW
from gpaw.mpi import world
from gpaw.cluster import Cluster
from gpaw.occupations import FermiDirac
from gpaw.test import equal
from gpaw.eigensolvers import RMMDIIS
from gpaw.lrtddft import LrTDDFT
import _gpaw


def test_rsf_yukawa_lrtddft_short(in_tmp_dir):
    libxc_version = getattr(_gpaw, 'libxc_version', '2.x.y')
    if int(libxc_version.split('.')[0]) < 3:
        from unittest import SkipTest
        raise SkipTest

    o_plus = Cluster(Atoms('Be', positions=[[0, 0, 0]]))
    o_plus.set_initial_magnetic_moments([1.0])
    o_plus.minimal_box(2.5, h=0.35)

    def get_paw():
        """Return calculator object."""
        c = {'energy': 0.05, 'eigenstates': 0.05, 'density': 0.05}
        return GPAW(convergence=c, eigensolver=RMMDIIS(),
                    nbands=3,
                    xc='PBE',
                    #                experimental={'niter_fixdensity': 2},
                    parallel={'domain': world.size}, h=0.35,
                    occupations=FermiDirac(width=0.0, fixmagmom=True))

    calc_plus = get_paw()
    calc_plus.set(txt='Be_plus_LCY_PBE_083.log', charge=1)
    o_plus.calc = calc_plus
    o_plus.get_potential_energy()
    calc_plus.set(xc='LCY-PBE:omega=0.83:unocc=True',
                  experimental={'niter_fixdensity': 2})
    o_plus.get_potential_energy()
    lr = LrTDDFT(calc_plus, txt='LCY_TDDFT_Be.log',
                 restrict={'istart': 0, 'jend': 1})
    equal(lr.xc.omega, 0.83)
    lr.write('LCY_TDDFT_Be.ex.gz')
    e_ion = 9.3
    ip_i = 13.36
    # reading is problematic with EXX on more than one core
    if world.rank == 0:
        lr2 = LrTDDFT.read('LCY_TDDFT_Be.ex.gz')
        lr2.diagonalize()
        equal(lr2.xc.omega, 0.83)
        ion_i = lr2[0].get_energy() * Hartree + e_ion
        equal(ion_i, ip_i, 0.3)
