import io
import pytest

from gpaw import GPAW, restart
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.old_gpaw_only
@pytest.mark.sic
def test_pwsic(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_pwsic"])
    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]

    # Numeric forces, generated by disabled code below
    f2 = np.array(
        [
            [0.22161312, -0.98564396, -0.00204214],
            [-0.34986867, 0.17494903, 0.00029861],
            [0.01085528, 0.56112341, 0.00129632],
        ]
    )
    assert f2 == pytest.approx(f, abs=3e-2)
    assert e == pytest.approx(-9.98919, abs=1e-3)
    assert gap == pytest.approx(9.555, 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")

    H2O, calc = restart("h2o.gpw", txt="-")
    H2O.positions += 1.0e-6
    f3 = H2O.get_forces()
    niter = calc.get_number_of_iterations()
    assert niter == pytest.approx(4, abs=3)
    assert f2 == pytest.approx(f3, abs=3e-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.28876    1.00000    0    -21.29052    1.00000
           1    -21.02682    1.00000    1    -21.03380    1.00000
           2    -13.96149    1.00000    2    -13.94529    1.00000
           3    -13.91558    1.00000    3    -13.92438    1.00000
           4     -0.94710    0.00000    4     -0.94712    0.00000
           5      0.76805    0.00000    5      0.76806    0.00000
        """
        expect_eigen_str = """\
        Band  Eigenvalues  Occupancy  Eigenvalues  Occupancy
           0    -30.18943    1.00000    -30.19005    1.00000
           1    -16.73073    1.00000    -16.73110    1.00000
           2    -12.77049    1.00000    -12.77136    1.00000
           3    -10.50200    1.00000    -10.50147    1.00000
           4     -0.94720    0.00000     -0.94720    0.00000
           5      0.76815    0.00000      0.76815    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,
        )
