import ase.io
import numpy as np
import pytest
from ase.build import molecule
from ase.calculators.vdwcorrection import vdWTkatchenko09prl
from ase.parallel import barrier, parprint

from gpaw import GPAW
from gpaw.analyse.hirshfeld import HirshfeldPartitioning
from gpaw.analyse.vdwradii import vdWradii
from gpaw.cluster import Cluster
from gpaw.test import equal


@pytest.mark.later
def test_vdw_ts09(in_tmp_dir):
    h = 0.4
    s = Cluster(molecule('LiH'))
    s.minimal_box(3., h=h)

    def print_charge_and_check(hp, q=0, label='unpolarized'):
        q_a = np.array(hp.get_charges())
        parprint('Charges ({0})='.format(label), q_a, ', sum=', q_a.sum())
        equal(q_a.sum(), q, 0.03)
        return q_a

    # spin unpolarized

    if 1:
        out_traj = 'LiH.traj'
        out_txt = 'LiH.txt'

        cc = GPAW(mode='fd', h=h, xc='PBE', txt=out_txt)

        # this is needed to initialize txt output
        cc.initialize(s)

        hp = HirshfeldPartitioning(cc)
        c = vdWTkatchenko09prl(hp,
                               vdWradii(s.get_chemical_symbols(), 'PBE'))
        s.calc = c
        E = s.get_potential_energy()
        F_ac = s.get_forces()
        s.write(out_traj)
        q_a = print_charge_and_check(hp)

        barrier()

        # test I/O, accuracy due to text output
        accuracy = 1.e-5
        for fname in [out_traj, out_txt]:
            s_out = ase.io.read(fname)
            equal(s_out.get_potential_energy(), E, accuracy)
            for fi, fo in zip(F_ac, s_out.get_forces()):
                equal(fi, fo, accuracy)

    # spin polarized

    if 0:
        ccs = GPAW(mode='fd', h=h, xc='PBE', spinpol=True,
                   txt=None)
        hps = HirshfeldPartitioning(ccs)
        cs = vdWTkatchenko09prl(hps, vdWradii(s.get_chemical_symbols(), 'PBE'))
        s.calc = cs
        Es = s.get_potential_energy()
        Fs_ac = s.get_forces()

        qs_a = print_charge_and_check(hps, label='spin')

        equal(q_a, qs_a, 1.e-6)
        equal(E, Es, 1.e-4)
        equal(F_ac, Fs_ac, 1.e-4)

    # charged

    if 0:
        cc = cc.new(charge=1)
        hpp = HirshfeldPartitioning(cc)
        cp = vdWTkatchenko09prl(hpp,
                                vdWradii(s.get_chemical_symbols(), 'PBE'))
        s.calc = cp
        E = s.get_potential_energy()

        print_charge_and_check(hpp, 1, label='+1')
