import pytest
from gpaw.mpi import world
from gpaw.utilities import compiled_with_sl
import numpy as np
from ase import Atoms
from ase.build import bulk
from ase.dft.kpoints import monkhorst_pack

from gpaw import GPAW, FermiDirac
from gpaw.mpi import serial_comm
from gpaw.test import equal
from gpaw.xc.rpa import RPACorrelation
from gpaw.xc.fxc import FXCCorrelation

pytestmark = pytest.mark.skipif(
    world.size != 1 and not compiled_with_sl(),
    reason='world.size != 1 and not compiled_with_sl()')


@pytest.mark.slow
def test_ralda_ralda_energy_Si(in_tmp_dir):
    a0 = 5.43
    cell = bulk('Si', 'fcc', a=a0).get_cell()
    Si = Atoms('Si2', cell=cell, pbc=True,
               scaled_positions=((0, 0, 0), (0.25, 0.25, 0.25)))

    kpts = monkhorst_pack((2, 2, 2))
    kpts += np.array([1 / 4, 1 / 4, 1 / 4])

    calc = GPAW(mode='pw',
                kpts=kpts,
                occupations=FermiDirac(0.001),
                communicator=serial_comm)
    Si.calc = calc
    Si.get_potential_energy()
    calc.diagonalize_full_hamiltonian(nbands=50)

    rpa = RPACorrelation(calc)
    E_rpa1 = rpa.calculate(ecut=[25, 50])

    fxc = FXCCorrelation(calc, xc='RPA', nlambda=16)
    E_rpa2 = fxc.calculate(ecut=[25, 50])

    fxc = FXCCorrelation(calc, xc='rALDA', unit_cells=[1, 1, 2])
    E_ralda = fxc.calculate(ecut=[25, 50])

    fxc = FXCCorrelation(calc, xc='rAPBE', unit_cells=[1, 1, 2])
    E_rapbe = fxc.calculate(ecut=[25, 50])

    fxc = FXCCorrelation(calc, xc='rALDA', av_scheme='wavevector')
    E_raldawave = fxc.calculate(ecut=[25, 50])

    fxc = FXCCorrelation(calc, xc='rAPBE', av_scheme='wavevector')
    E_rapbewave = fxc.calculate(ecut=[25, 50])

    fxc = FXCCorrelation(calc, xc='JGMs', av_scheme='wavevector',
                         Eg=3.1, nlambda=2)
    E_JGMs = fxc.calculate(ecut=[25, 50])

    fxc = FXCCorrelation(calc, xc='CP_dyn', av_scheme='wavevector',
                         nfrequencies=2, nlambda=2)
    E_CPdyn = fxc.calculate(ecut=[25, 50])

    equal(E_rpa1[-1], E_rpa2[-1], 0.01)
    equal(E_rpa2[-1], -12.6495, 0.002)
    equal(E_ralda[-1], -11.368, 0.002)
    equal(E_rapbe[-1], -11.149, 0.002)
    equal(E_raldawave[-1], -11.0910, 0.002)
    equal(E_rapbewave[-1], -10.831, 0.002)
    equal(E_JGMs[-1], -11.2561, 0.001)
    equal(E_CPdyn[-1], -7.8640, 0.001)
