import io
import pytest

from gpaw import GPAW
import numpy as np
from ase.dft.bandgap import bandgap
from ase.units import Ha

import numpy.testing as npt
from gpaw.io.logger import GPAWLogger
from gpaw.wavefunctions.base import eigenvalue_string
from gpaw.test.sic._utils import (mk_arr_from_str,
                                  extract_lagrange_section,
                                  MockWorld)
from gpaw.mpi import rank


@pytest.mark.slow
@pytest.mark.old_gpaw_only
@pytest.mark.sic
def test_fdsic(in_tmp_dir, gpw_files):
    """
    Test Perdew-Zunger Self-Interaction
    Correction in PW mode using DirectMin
    :param in_tmp_dir:
    :return:
    """
    calc = GPAW(gpw_files['h2o_fdsic'])
    H2O = calc.atoms
    H2O.calc = calc
    e = H2O.get_potential_energy()
    f = H2O.get_forces()
    efermi = calc.wfs.fermi_levels[0] * Ha
    gap = bandgap(calc, efermi=efermi)[0]
    e_old = -18.144428

    assert e == pytest.approx(e_old, abs=1e-3)
    # Numeric forces, generated by disabled code below
    f_num = np.array([[2.11270273, 4.89616573e-01, -6.00482442e-04],
                      [-2.15829241, 3.54950512e-01, 1.04418211e-04],
                      [6.67703026e-01, -8.89596180e-01, 8.83126024e-05]])

    assert f == pytest.approx(f_num, abs=5e-2)
    assert gap == pytest.approx(10.215, abs=1e-2)

    numeric = False
    if numeric:
        from gpaw.test import calculate_numerical_forces
        f_num = calculate_numerical_forces(H2O, 0.001)
        print('Numerical forces')
        print(f_num)
        print(f - f_num, np.abs(f - f_num).max())

    calc.write('h2o.gpw', mode='all')
    from gpaw import restart
    H2O, calc = restart('h2o.gpw', txt='-')
    H2O.positions += 1.0e-6
    e = H2O.get_potential_energy()
    f = H2O.get_forces()
    niter = calc.get_number_of_iterations()
    assert niter == pytest.approx(4, abs=3)
    assert e == pytest.approx(e_old, abs=1e-3)
    assert f == pytest.approx(f_num, abs=5e-2)

    if rank == 0:
        logger = GPAWLogger(MockWorld(rank=0))
        string_io = io.StringIO()
        logger.fd = string_io
        calc.wfs.summary_func(logger)
        lstr = extract_lagrange_section(string_io.getvalue())

        expect_lagrange_str = """\
        Band         L_ii   Occupancy   Band      L_ii   Occupancy
           0    -21.19082    1.00000    0    -21.19081    1.00000
           1    -20.53526    1.00000    1    -20.53521    1.00000
           2    -13.88424    1.00000    2    -13.88428    1.00000
           3    -13.83555    1.00000    3    -13.83554    1.00000
           4     -0.58073    0.00000    4     -0.58073    0.00000
           5      1.21087    0.00000    5      1.21087    0.00000
        """
        expect_eigen_str = """\
        Band  Eigenvalues  Occupancy  Eigenvalues  Occupancy
           0    -28.89620    1.00000    -28.89620    1.00000
           1    -16.92142    1.00000    -16.92142    1.00000
           2    -12.82836    1.00000    -12.82835    1.00000
           3    -10.79987    1.00000    -10.79987    1.00000
           4     -0.58369    0.00000     -0.58369    0.00000
           5      1.21382    0.00000      1.21382    0.00000
        """

        npt.assert_allclose(
            mk_arr_from_str(expect_lagrange_str, 6),
            mk_arr_from_str(lstr, 6),
            atol=0.3,
        )

        npt.assert_allclose(
            mk_arr_from_str(expect_eigen_str, 5),
            mk_arr_from_str(eigenvalue_string(calc.wfs), 5, skip_rows=1),
            atol=0.3,
        )
