"""Integration tests for GW flows."""

from __future__ import annotations

import socket

import pytest

import abipy.data as abidata
from abipy import abilab, flowtk

hostname = socket.gethostname()

skip_hosts = [
    "scope.pcpm.ucl.ac.be",
]


def make_g0w0_inputs(ngkpt, tvars):
    """
    Input files for the calculation of the GW corrections.

    Returns: gs_input, nscf_input, scr_input, sigma_input
    """
    multi = abilab.MultiDataset(structure=abidata.cif_file("si.cif"), pseudos=abidata.pseudos("14si.pspnc"), ndtset=4)

    # This grid is the most economical, but does not contain the Gamma point.
    scf_kmesh = dict(ngkpt=ngkpt, shiftk=[0.5, 0.5, 0.5, 0.5, 0.0, 0.0, 0.0, 0.5, 0.0, 0.0, 0.0, 0.5])

    # This grid contains the Gamma point, which is the point at which
    # we will compute the (direct) band gap.
    gw_kmesh = dict(ngkpt=ngkpt, shiftk=[0.0, 0.0, 0.0, 0.0, 0.5, 0.5, 0.5, 0.0, 0.5, 0.5, 0.5, 0.0])

    # Global variables. gw_kmesh is used in all datasets except DATASET 1.
    ecut = 4

    multi.set_vars(
        ecut=ecut,
        pawecutdg=ecut * 2 if multi.ispaw else None,
        istwfk="*1",
        paral_kgb=tvars.paral_kgb,
        gwpara=2,
    )
    multi.set_kmesh(**gw_kmesh)

    # Dataset 1 (GS run)
    multi[0].set_kmesh(**scf_kmesh)
    multi[0].set_vars(tolvrs=1e-6, nband=4)

    # Dataset 2 (NSCF run)
    multi[1].set_vars(iscf=-2, tolwfr=1e-10, nband=10, nbdbuf=2)

    # Dataset3: Calculation of the screening.
    multi[2].set_vars(
        optdriver=3,
        nband=8,
        ecutwfn=ecut,
        symchi=1,
        inclvkb=0,
        ecuteps=2.0,
    )

    # Dataset4: Calculation of the Self-Energy matrix elements (GW corrections)
    kptgw = [
        -2.50000000e-01,
        -2.50000000e-01,
        0.00000000e00,
        -2.50000000e-01,
        2.50000000e-01,
        0.00000000e00,
        5.00000000e-01,
        5.00000000e-01,
        0.00000000e00,
        -2.50000000e-01,
        5.00000000e-01,
        2.50000000e-01,
        5.00000000e-01,
        0.00000000e00,
        0.00000000e00,
        0.00000000e00,
        0.00000000e00,
        0.00000000e00,
    ]

    multi[3].set_vars(
        optdriver=4,
        nband=10,
        ecutwfn=ecut,
        ecuteps=2.0,
        ecutsigx=2.0,
        symsigma=1,
        # gw_qprange=0,
    )

    bdgw = [4, 5]
    multi[3].set_kptgw(kptgw, bdgw)

    return multi.split_datasets()


@pytest.mark.skipif(hostname in skip_hosts, reason=f"Skipped on {hostname}")
def itest_g0w0_flow(fwp, tvars):
    """Test flow for G0W0 calculations."""
    scf, nscf, scr, sig = make_g0w0_inputs(ngkpt=[2, 2, 2], tvars=tvars)

    flow = flowtk.g0w0_flow(fwp.workdir, scf, nscf, scr, sig, manager=fwp.manager)
    # Will remove output files at run-time.
    flow.set_garbage_collector()
    flow.build_and_pickle_dump(abivalidate=True)

    for task in flow[0]:
        task.start_and_wait()

    flow.check_status(show=True)
    assert all(work.finalized for work in flow)
    if not flow.all_ok:
        flow.debug()
        raise RuntimeError

    scf_task = flow[0][0]
    nscf_task = flow[0][1]
    scr_task = flow[0][2]
    sig_task = flow[0][3]

    # Test garbage_collector
    # The WFK|SCR file should have been removed because we call set_garbage_collector
    assert not scf_task.outdir.has_abiext("WFK")
    assert not nscf_task.outdir.has_abiext("WFK")
    assert not scr_task.outdir.has_abiext("SCR")
    assert not scr_task.outdir.has_abiext("SUS")

    # The sigma task should produce a SIGRES file.
    sigfile = sig_task.outdir.list_filepaths(wildcard="*SIGRES.nc")[0]
    assert sigfile
    with abilab.abiopen(sigfile) as sigres:
        sigres.to_string(verbose=2)
        assert sigres.nsppol == 1

    # Test SigmaTask inspect method
    # if has_matplotlib():
    # sig_task.inspect(show=False)

    # Test get_results for Sigma and Scr
    scr_task.get_results()
    sig_task.get_results()

    # Test SCR.nc file (this is optional)
    if scr_task.scr_path:
        with scr_task.open_scr() as scr:
            scr.to_string(verbose=2)
            assert len(scr.wpts) == 2
            assert scr.nwre == 1 and scr.nwim == 1
            for iq, qpoint in enumerate(scr.qpoints[:2]):
                # print(qpoint)
                qpt, iqcheck = scr.reader.find_qpoint_fileindex(qpoint)
                assert iqcheck == iq
                em1 = scr.get_em1(qpoint)
                # print(em1)

    # TODO Add more tests
    # assert flow.validate_json_schema()


@pytest.mark.skipif(hostname in skip_hosts, reason=f"Skipped on {hostname}")
def itest_htc_g0w0(fwp, tvars):
    """Testing G0W0Work."""
    structure = abilab.Structure.from_file(abidata.cif_file("si.cif"))
    pseudos = abidata.pseudos("14si.pspnc")

    flow = flowtk.Flow(fwp.workdir, manager=fwp.manager)

    scf_kppa = 10
    nscf_nband = 10
    # nscf_ngkpt = [4,4,4]
    # nscf_shiftk = [0.0, 0.0, 0.0]
    ecut, ecuteps, ecutsigx = 4, 2, 3
    # scr_nband = 50
    # sigma_nband = 50

    extra_abivars = dict(
        ecut=ecut,
        istwfk="*1",
        paral_kgb=tvars.paral_kgb,
    )

    multi = abilab.g0w0_with_ppmodel_inputs(
        structure,
        pseudos,
        scf_kppa,
        nscf_nband,
        ecuteps,
        ecutsigx,
        ecut=ecut,
        pawecutdg=None,
        accuracy="normal",
        spin_mode="unpolarized",
        smearing=None,
        # ppmodel="godby", charge=0.0, scf_algorithm=None, inclvkb=2, scr_nband=None,
        # sigma_nband=None, gw_qprange=1):
    )
    multi.set_vars(paral_kgb=tvars.paral_kgb)

    scf_input, nscf_input, scr_input, sigma_input = multi.split_datasets()
    work = flowtk.G0W0Work(scf_input, nscf_input, scr_input, sigma_input)

    flow.register_work(work)
    flow.allocate()
    flow.connect_signals()

    fwp.scheduler.add_flow(flow)
    assert fwp.scheduler.start() == 0
    assert not fwp.scheduler.exceptions
    assert fwp.scheduler.nlaunch == 4

    # The sigma task should produce a SCR file.
    assert len(work[2].outdir.list_filepaths(wildcard="*SCR")) == 1

    flow.show_status()
    if not flow.all_ok:
        flow.debug()
        raise RuntimeError

    assert all(work.finalized for work in flow)

    # assert flow.validate_json_schema()
