"""
Tests the DQM compute dispatch module
"""
import pprint
import re
import sys
import tempfile
from pathlib import Path

import pytest

import qcengine as qcng
from qcengine.testing import checkver_and_convert, from_v2, schema_versions, using

_canonical_methods = [
    # needs attn ("adcc", {"method": "adc2", "basis": "6-31G"}, {"n_triplets": 3}),
    pytest.param("cfour", {"method": "hf", "basis": "6-31G"}, {}, marks=using("cfour")),
    pytest.param("dftd3", {"method": "b3lyp-d3"}, {}, marks=using("classic-dftd3")),
    pytest.param("gamess", {"method": "hf", "basis": "n31"}, {"basis__NGAUSS": 6}, marks=using("gamess")),
    pytest.param("jaguar", {"method": "hf", "basis": "6-31G"}, {"nops": 1}, marks=using("jaguar")),
    pytest.param("gcp", {"method": "hf3c"}, {}, marks=using("classic-gcp")),
    pytest.param("mctc-gcp", {"method": "dft/sv"}, {}, marks=using("mctc-gcp")),
    # needs attn ("molpro", {"method": "hf", "basis": "6-31G"}, {}),
    # needs attn ("mopac", {"method": "PM6"}, {}),
    pytest.param("mp2d", {"method": "MP2-DMP2"}, {}, marks=using("mp2d")),
    # needs attn ("mrchem", {"method": "blyp"}, {"world_prec": 1.0e-3}),
    pytest.param("nwchem", {"method": "hf", "basis": "6-31G"}, {}, marks=using("nwchem")),
    # needs attn ("openmm", {"method": "openff-1.0.0", "basis": "smirnoff"}, {}),
    pytest.param("psi4", {"method": "hf", "basis": "6-31G"}, {"gradient_write": True}, marks=using("psi4")),
    # needs attn ("qchem", {"method": "hf", "basis": "6-31G"}, {}),
    # needs attn ("qcore", {"method": "pbe", "basis": "6-31G"}, {}),
    # needs attn ("rdkit", {"method": "UFF"}, {}),
    # needs attn ("terachem", {"method": "bad"}, {}),
    # needs attn ("terachem_pbs", {"method": "b3lyp", "basis": "6-31G"}, {}),
    # needs attn ("torchani", {"method": "ANI1x"}, {}),
    # needs attn ("turbomole", {"method": "pbe", "basis": "6-31G"}, {}),
    # needs attn ("xtb", {"method": "GFN2-xTB"}, {}),
]


def _get_molecule(program, method, molcls):
    if program in ["openmm", "terachem_pbs"]:
        dmol = qcng.get_molecule("water", return_dict=True)
    elif program == "gamess" and method == "ccsd(t)":
        dmol = qcng.get_molecule("water", return_dict=True)
    else:
        dmol = qcng.get_molecule("hydrogen", return_dict=True)

    return molcls(**dmol)


@pytest.mark.parametrize(
    "memory_trickery",
    [
        pytest.param({}, id="none"),
        pytest.param(
            # native keywords that CONTRADICT config.memory below
            {
                "cfour": {"memory_size": "5000"},
                "gamess": {"system__mwords": "5000"},
                "nwchem": {"memory": "5 gb"},
                "psi4": {},  # no contradictory memory keyword in psi
            },
            id="dsl",
        ),
    ],
)
@pytest.mark.parametrize("program, model, keywords", _canonical_methods)
def test_local_options_memory_gib(program, model, keywords, memory_trickery, schema_versions, request):
    """Ensure memory handling implemented in harness (if applicable).

    For available harnesses, run minimal calc at specific total node memory, both through runtime
      config alone and with clashing (and non-QCEngine-like) keyword spec. Check memory quantity
      shows up in ``TaskConfig``.
    For ``managed-memory``-active harnesses, check that memory registers in output.

    New Harness Instructions
    ------------------------
    * Make sure minimal calc is in _canonical_methods above.
    * If ``managed_memory=True`` in harness, add regex to ``stdout_ref`` below to check that memory
      is specifiable.
    * If this test doesn't work, implement or adjust ``config.memory`` in your harness.

    """
    models, retver, _ = schema_versions

    harness = qcng.get_program(program)
    molecule = _get_molecule(program, model["method"], models.Molecule)

    addl_keywords = memory_trickery.get(program, memory_trickery)
    use_keywords = {**keywords, **addl_keywords}

    #  <<  Config

    config = qcng.config.get_config(
        hostname="something",
        task_config={
            "ncores": 1,
            "nnodes": 1,
            "memory": 1.555,
        },
    )

    #  <<  Run

    if from_v2(request.node.name):
        inp = models.AtomicInput(
            molecule=molecule,
            specification=models.AtomicSpecification(driver="energy", model=model, keywords=use_keywords),
        )
    else:
        inp = models.AtomicInput(molecule=molecule, driver="energy", model=model, keywords=use_keywords)

    inp = checkver_and_convert(inp, request.node.name, "pre")
    ret = qcng.compute(inp, program, raise_error=True, task_config=config.model_dump(), return_version=retver)
    ret = checkver_and_convert(ret, request.node.name, "post")

    pprint.pprint(ret.model_dump(), width=200)
    assert ret.success is True

    #  <<  Reference

    stdout_ref = {  # 1.555 GiB = 208708567 quad-words
        "cfour": "Allocated    1592 MB of main memory",
        "gamess": "208000000 WORDS OF MEMORY AVAILABLE",
        "nwchem": r"total\s+=\s+2087085\d\d doubles =\s+1592.3 Mbytes",  # doubles is quad-words. Mbytes is MiB
        "psi4": "1592 MiB Core",
    }

    #  <<  Test

    assert config.ncores == 1
    assert pytest.approx(config.memory, 0.1) == 1.555

    if harness._defaults["managed_memory"] is True:
        assert re.search(stdout_ref[program], ret.stdout), f"Memory pattern not found: {stdout_ref[program]}"


@pytest.mark.parametrize("program, model, keywords", _canonical_methods)
def test_local_options_scratch(program, model, keywords, schema_versions, request):
    """Ensure scratch handling implemented in harness (if applicable).

    For available harnesses, run minimal calc at specific scratch directory name (randomly generated
      during test) and skip scratch clean-up. Check scratch settings show up in ``TaskConfig``.
    For ``scratch``-active harnesses, check that an expected file is written to and left behind in
      scratch directory. Check any scratch-related printing in output.

    New Harness Instructions
    ------------------------
    * Make sure minimal calc is in _canonical_methods above.
    * If ``scratch=True`` in harness, add single file (preferrably output) glob to ``scratch_sample``
      below to check that program scratch is directable.
    * If ``scratch=True`` in harness, if scratch directory mentioned in output, add regex to
      ``stdout_ref`` below to check that program scratch is directable. Otherwise, add an
      always-passing regex.
    * If this test doesn't work, implement or adjust ``config.scratch_directory`` and
      ``config.scratch_messy`` in your harness.

    """
    models, retver, _ = schema_versions

    harness = qcng.get_program(program)
    molecule = _get_molecule(program, model["method"], models.Molecule)

    #  <<  Config

    scratch_directory = tempfile.mkdtemp(suffix="_" + program)

    config = qcng.config.get_config(
        hostname="something",
        task_config={
            "scratch_directory": scratch_directory,
            "scratch_messy": True,
        },
    )

    #  <<  Run

    if from_v2(request.node.name):
        inp = models.AtomicInput(
            molecule=molecule, specification=models.AtomicSpecification(driver="energy", model=model, keywords=keywords)
        )
    else:
        inp = models.AtomicInput(molecule=molecule, driver="energy", model=model, keywords=keywords)

    inp = checkver_and_convert(inp, request.node.name, "pre")
    ret = qcng.compute(inp, program, raise_error=True, task_config=config.model_dump(), return_version=retver)
    ret = checkver_and_convert(ret, request.node.name, "post")

    pprint.pprint(ret.model_dump(), width=200)
    assert ret.success is True

    #  <<  Reference

    stdout_ref = {
        "cfour": "University of Florida",  # freebie
        "dftd3": "Grimme",  # freebie
        "gamess": "IOWA STATE UNIVERSITY",  # freebie
        "jaguar": "Jaguar version",  # freebie
        "gcp": "Grimme",  # freebie
        "mctc-gcp": "Grimme",  # freebie
        "mp2d": "Beran",  # freebie
        "nwchem": "E. Apra",  # freebie
        "psi4": rf"Scratch directory: {scratch_directory}/tmp\w+_psi_scratch/",
    }
    if sys.platform.startswith("win"):
        # too hard to regex Windows paths that need escape chars
        stdout_ref["psi4"] = f"Scratch directory: "

    # a scratch file (preferrably output) expected after job if scratch not cleaned up
    scratch_sample = {
        "cfour": "*/NEWFOCK",
        "dftd3": "*/dftd3_geometry.xyz",  # no outfiles
        "gamess": "*/gamess.dat",
        "jaguar": "*/dispatch.out",
        "gcp": "*/gcp_geometry.xyz",  # no outfiles
        "mctc-gcp": "*/gcp_geometry.xyz",  # no outfiles
        "mp2d": "*/mp2d_geometry",  # no outfiles
        "nwchem": "*/nwchem.db",
        "psi4": "*/psi.*.35",
    }

    #  <<  Test

    assert config.scratch_directory.endswith(program)

    if harness._defaults["scratch"] is True:
        sample_file = list(Path(scratch_directory).glob(scratch_sample[program]))
        assert len(sample_file) == 1, f"Scratch sample not found: {scratch_sample[program]} in {scratch_directory}"

        assert re.search(stdout_ref[program], ret.stdout), f"Scratch pattern not found: {stdout_ref[program]}"


@pytest.mark.parametrize("ncores", [1, 3])
@pytest.mark.parametrize("program, model, keywords", _canonical_methods)
def test_local_options_ncores(program, model, keywords, ncores, schema_versions, request):
    """Ensure multithreading implemented in harness (if applicable) or multithreaded runs don't
    break harness (if inapplicable).

    For available harnesses, run minimal calc with single and multiple cores; check ncores count
      shows up in ``TaskConfig``.
    For ``thread_parallel``-active harnesses, check ncores count registers in output.

    New Harness Instructions
    ------------------------
    * Make sure minimal calc is in _canonical_methods above.
    * If ``thread_parallel=True`` in harness, add regex to ``stdout_ref`` below to check ncores the
      program sees.
    * If this test doesn't work, implement or adjust ``config.ncores`` in your harness.

    """
    models, retver, _ = schema_versions

    harness = qcng.get_program(program)
    molecule = _get_molecule(program, model["method"], models.Molecule)

    #  <<  Config

    config = qcng.config.get_config(
        hostname="something",
        task_config={
            "ncores": ncores,
            "nnodes": 1,
        },
    )

    #  <<  Run

    if from_v2(request.node.name):
        inp = models.AtomicInput(
            molecule=molecule, specification=models.AtomicSpecification(driver="energy", model=model, keywords=keywords)
        )
    else:
        inp = models.AtomicInput(molecule=molecule, driver="energy", model=model, keywords=keywords)

    inp = checkver_and_convert(inp, request.node.name, "pre")
    ret = qcng.compute(inp, program, raise_error=True, task_config=config.model_dump(), return_version=retver)
    ret = checkver_and_convert(ret, request.node.name, "post")

    pprint.pprint(ret.model_dump(), width=200)
    assert ret.success is True

    #  <<  Reference

    stdout_ref = {
        "cfour": rf"Running with {ncores} threads/proc",
        "gamess": rf"MEMDDI DISTRIBUTED OVER\s+{ncores} PROCESSORS",
        "jaguar": rf"Using up to\s+{ncores} threads per process",
        # "gamess": rf"PARALLEL VERSION RUNNING ON\s+{ncores} PROCESSORS IN\s+1 NODES",  # no line for serial
        # nwchem is node_parallel only
        "psi4": rf"Threads:\s+{ncores}",
    }

    #  <<  Test

    assert config.ncores == ncores
    assert config.nnodes == 1

    if harness._defaults["thread_parallel"] is True:
        assert re.search(stdout_ref[program], ret.stdout), f"Thread pattern not found: {stdout_ref[program]}"
