import numpy as np
from ase.build import molecule
from ase.data.vdw import vdw_radii
from gpaw import GPAW
from gpaw.cluster import Cluster
from gpaw.test import equal
from gpaw.solvation import (SolvationGPAW, EffectivePotentialCavity,
                            Power12Potential, LinearDielectric)

SKIP_REF_CALC = True

energy_eps = 0.0005 / 8.
forces_eps = 3e-2

h = 0.3
vac = 3.0
u0 = .180
T = 298.15
vdw_radii = vdw_radii.copy()
vdw_radii[1] = 1.09


def atomic_radii(atoms):
    return [vdw_radii[n] for n in atoms.numbers]

atoms = Cluster(molecule('H2O'))
atoms.minimal_box(vac, h)

convergence = {
    'energy': energy_eps,
    'forces': forces_eps ** 2,  # Force error is squared
    'density': 10.,
    'eigenstates': 10.,
}

if not SKIP_REF_CALC:
    atoms.calc = GPAW(xc='LDA', h=h, convergence=convergence)
    Eref = atoms.get_potential_energy()
    print(Eref)
    Fref = atoms.get_forces()
    print(Fref)
else:
    # setups: 0.9.11271, same settings as above
    Eref = -11.9932

    Fref = np.array(
        [[1.95122040e-12, -1.17770462e-12, -6.04993798e+00],
         [6.61270337e-14, 1.58227909e+00, 6.06605145e-02],
         [1.35947527e-13, -1.58227909e+00, 6.06605145e-02]])

atoms.calc = SolvationGPAW(
    xc='LDA', h=h, convergence=convergence,
    cavity=EffectivePotentialCavity(
        effective_potential=Power12Potential(atomic_radii=atomic_radii, u0=u0),
        temperature=T
    ),
    dielectric=LinearDielectric(epsinf=1.0),
)
Etest = atoms.get_potential_energy()
Eeltest = atoms.calc.get_electrostatic_energy()
Ftest = atoms.get_forces()
equal(Etest, Eref, energy_eps * atoms.calc.get_number_of_electrons())
equal(Ftest, Fref, forces_eps)
equal(Eeltest, Etest)
