from math import log

import pytest
from ase import Atoms
from ase.db import connect
from ase.io import iread, read, write
from ase.units import Bohr

from gpaw import GPAW, FermiDirac


@pytest.mark.later
def test_generic_hydrogen(in_tmp_dir, needs_ase_master):
    a = 4.0
    h = 0.2
    hydrogen = Atoms('H',
                     [(a / 2, a / 2, a / 2)],
                     cell=(a, a, a))

    params = dict(mode='fd',
                  h=h,
                  nbands=1,
                  convergence={'energy': 1e-7})
    hydrogen.calc = GPAW(**params,
                         txt='h.txt')
    e1 = hydrogen.get_potential_energy()
    assert e1 == pytest.approx(0.526939, abs=0.001)

    dens = hydrogen.calc.density
    c = dens.gd.find_center(dens.nt_sG[0]) * Bohr
    assert abs(c - a / 2).max() == pytest.approx(0, abs=1e-13)

    kT = 0.001
    hydrogen.calc = GPAW(**params,
                         occupations=FermiDirac(width=kT))
    e2 = hydrogen.get_potential_energy()
    assert e1 == pytest.approx(e2 + log(2) * kT, abs=3.0e-7)

    # Test ase.db a bit:
    for name in ['h.json', 'h.db']:
        con = connect(name, append=False)
        con.write(hydrogen)
        id = con.write(hydrogen, foo='bar', data={'abc': [1, 2, 3]})
        assert id == 2
        assert con.reserve(foo='bar') is None
        row = con.get(foo='bar')
        assert row.energy == e2
        assert sum(row.data.abc) == 6
        del con[1]
        assert con.reserve(x=42) == 3

        write('x' + name, hydrogen)
        write('xx' + name, [hydrogen, hydrogen])

        assert read(name + '@foo=bar')[0].get_potential_energy() == e2
        for n, h in zip([1, 0], iread(name + '@:')):
            assert n == len(h)

    # Test parsing of GPAW's text output:
    h = read('h.txt')
    error = abs(h.calc.get_eigenvalues() -
                hydrogen.calc.get_eigenvalues()).max()
    assert error < 1e-5, error

    # Test get_electrostatic_potential() method
    v = hydrogen.calc.get_electrostatic_potential()
    print(v.shape, v.ptp())
