from ase.build import molecule
from ase.parallel import parprint

from gpaw import GPAW
from gpaw.cluster import Cluster
from gpaw.analyse.hirshfeld import HirshfeldDensity, HirshfeldPartitioning
from gpaw.analyse.wignerseitz import WignerSeitz
from gpaw.test import equal
import pytest


@pytest.mark.later
def test_utilities_partitioning(in_tmp_dir):
    h = 0.4
    gpwname = 'H2O' + str(h) + '.gpw'

    def run(lastres=[]):
        results = []

        # Hirshfeld ----------------------------------------

        if 1:

            hd = HirshfeldDensity(calc)

            # check for the number of electrons
            expected = [[None, 10],
                        [[0, 1, 2], 10],
                        [[1, 2], 2],
                        [[0], 8],
                        ]
            for gridrefinement in [1, 2, 4]:
                # Test for all gridrefinements for get_all_electron_density
                parprint('grid refinement', gridrefinement)
                for result in expected:
                    indicees, result = result
                    full, gd = hd.get_density(indicees, gridrefinement)
                    parprint('indicees', indicees, end=': ')
                    parprint('result, expected:', gd.integrate(full), result)
                    if gridrefinement < 4:
                        # The highest level of gridrefinement gets wrong
                        # electron numbers
                        equal(gd.integrate(full), result, 1.e-8)
                    else:
                        equal(gd.integrate(full), result, 1.e-4)

            hp = HirshfeldPartitioning(calc)
            vr = hp.get_effective_volume_ratios()
            parprint('Hirshfeld:', vr)
            if len(lastres):
                equal(vr, lastres.pop(0), 1.e-10)
            results.append(vr)

        # Wigner-Seitz ----------------------------------------

        if 1:
            ws = WignerSeitz(calc.density.finegd, mol, calc)

            vr = ws.get_effective_volume_ratios()
            parprint('Wigner-Seitz:', vr)
            if len(lastres):
                equal(vr, lastres.pop(0), 1.e-10)
            results.append(vr)

        return results

    mol = Cluster(molecule('H2O'))
    mol.minimal_box(2.5, h=h)

    # calculate
    if 1:
        parprint('### fresh:')
        calc = GPAW(mode='fd',
                    nbands=6,
                    h=h,
                    txt=None)
    if 1:
        calc.calculate(mol)
        calc.write(gpwname)
        lastres = run()

    # load previous calculation
    if 1:
        parprint('### reloaded:')
        calc = GPAW(gpwname, txt=None)
        mol = calc.get_atoms()
        run(lastres)

    # periodic modulo test
    parprint('### periodic:')
    mol.set_pbc(True)
    mol.translate(-mol[0].position)
    mol.translate([-1.e-24, 0, 0])
    calc.calculate(mol)
    run()
