import pytest
import numpy as np
from ase import Atoms

import gpaw.dscf as dscf
from gpaw import GPAW, MixerSum
from gpaw.mpi import rank


@pytest.mark.skip(reason='TODO')
def test_dscf_forces():
    L = 10.0

    eigensolver = 'cg'
    xc = 'PBE'

    E0 = []
    E1 = []
    d0 = 1.4
    delta = 0.05
    for i in range(3):
        d = d0 + (i - 1) * delta
        atoms = Atoms('N2',
                      [[0.5 * L, 0.5 * L, 0.5 * L - 0.5 * d],
                       [0.5 * L, 0.5 * L, 0.5 * L + 0.5 * d]])
        atoms.set_cell([L, L, L])

        calc_gs = GPAW(h=0.2,
                       nbands=-5,
                       xc=xc,
                       occupations={'name': 'fermi-dirac', 'width': 0.05},
                       eigensolver=eigensolver,
                       spinpol=True,
                       txt=None,
                       mixer=MixerSum(beta=0.1, nmaxold=5, weight=50.0),
                       convergence={'energy': 100,
                                    'density': 100,
                                    'eigenstates': 1.0e-9,
                                    'bands': -1})
        atoms.calc = calc_gs

        E0.append(atoms.get_potential_energy())
        if i == 1:
            F0 = atoms.get_forces()

        calc_es = GPAW(h=0.2,
                       nbands=-5,
                       xc=xc,
                       width=0.05,
                       eigensolver=eigensolver,
                       spinpol=True,
                       txt=None,
                       mixer=MixerSum(beta=0.1, nmaxold=5, weight=50.0),
                       convergence={'energy': 100,
                                    'density': 100,
                                    'eigenstates': 1.0e-9,
                                    'bands': -1})

        n = 5  # LUMO
        molecule = [0, 1]
        wf_u = [kpt.psit_nG[n] for kpt in calc_gs.wfs.kpt_u]
        p_uai = [dict([(molecule[a], P_ni[n])
                       for a, P_ni in kpt.P_ani.items()])
                 for kpt in calc_gs.wfs.kpt_u]

        atoms.calc = calc_es
        lumo = dscf.AEOrbital(calc_es, wf_u, p_uai)
        dscf.dscf_calculation(calc_es, [[1.0, lumo, 1]], atoms)

        E1.append(atoms.get_potential_energy())
        if i == 1:
            F1 = atoms.get_forces()

    f0 = np.sqrt(((F0[1, :] - F0[0, :])**2).sum()) * 0.5
    f0b = (E0[1] - E0[0]) / delta     # backward gradient
    f0f = (E0[2] - E0[1]) / delta     # forward gradient
    if rank == 0:
        print('Ground state')
        print(E0)
        print(f0b, '<', f0, '<', f0f)
    assert f0 > f0b
    assert f0 < f0f

    f1 = np.sqrt(((F1[1, :] - F1[0, :])**2).sum()) * 0.5
    f1b = (E1[1] - E1[0]) / delta     # backward gradient
    f1f = (E1[2] - E1[1]) / delta     # forward gradient
    if rank == 0:
        print('Excited state')
        print(E1)
        print(f1b, '<', f1, '<', f1f)
    assert f1 > f1b
    assert f1 < f1f
