import pprint

import pytest
from qcelemental import constants
from qcelemental.testing import compare, compare_values

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

try:
    import qcmanybody as qcmb
except ModuleNotFoundError:
    schema_versions = schema_versions0
except RuntimeError:
    # py314 + qcmanybody 0.5.1
    schema_versions = schema_versions0
else:
    if qcmb.__version__.startswith("0.5."):
        schema_versions = schema_versions0
    else:
        schema_versions = schema_versionsALL5


@pytest.fixture
def he_tetramer(schema_versions, request):
    models, _, _ = schema_versions

    a2 = 2 / constants.bohr2angstroms
    return models.Molecule(
        symbols=["He", "He", "He", "He"],
        fragments=[[0], [1], [2], [3]],
        geometry=[0, 0, 0, 0, 0, a2, 0, a2, 0, 0, a2, a2],
    )


@uusing("qcmanybody")
@pytest.mark.parametrize(
    "program,basis,keywords",
    [
        pytest.param("cfour", "aug-pvdz", {"frozen_core": False}, id="cfour_conv", marks=using("cfour")),
        pytest.param(
            "gamess", "accd", {"contrl__ispher": 1, "mp2__nacore": 0}, id="gamess_conv", marks=using("gamess")
        ),
        pytest.param(
            "nwchem",
            "aug-cc-pvdz",
            {"basis__spherical": True, "scf__thresh": 1.0e-8, "mp2__freeze": False},
            id="nwchem_conv",
            marks=using("nwchem"),
        ),
        pytest.param(
            "psi4",
            "aug-cc-pvdz",
            {"e_convergence": 1.0e-10, "d_convergence": 1.0e-10, "scf_type": "pk", "mp2_type": "conv"},
            id="psi4_conv",
            marks=using("psi4"),
        ),
    ],
)
@pytest.mark.parametrize(
    "mbe_keywords,anskey,calcinfo_nmbe",
    [
        pytest.param(
            {"bsse_type": "nocp", "return_total_data": False, "max_nbody": 3},
            "nocp_corrected_interaction_energy_through_3_body",
            14,
            id="3b_nocp",
        ),
    ],
)
def test_nbody_he4_single(
    schema_versions, program, basis, keywords, mbe_keywords, anskey, calcinfo_nmbe, he_tetramer, request
):
    models, retver, _ = schema_versions

    if from_v2(request.node.name):
        from qcmanybody.models.v2 import ManyBodyInput

        AtomicSpecification = models.AtomicSpecification
        subptcl = "cluster_results"
    else:
        from qcmanybody.models.v1 import AtomicSpecification, ManyBodyInput

        subptcl = "component_results"

    # "to_v2"
    subres = "cluster_results" if ("v2" in request.node.name) else "component_results"
    import qcmanybody as qcmb

    if qcmb.__version__.startswith("0.5."):
        subres = "component_results"

    atomic_spec = AtomicSpecification(
        model={"method": "mp2", "basis": basis},
        program=program,
        driver="energy",
        keywords=keywords,
        protocols={"stdout": False},
    )
    mbe_model = ManyBodyInput(
        specification={
            "specification": atomic_spec,
            "keywords": mbe_keywords,
            "driver": "energy",
            "protocols": {subptcl: "all"},
        },
        molecule=he_tetramer,
    )

    mbe_model = checkver_and_convert(mbe_model, request.node.name, "pre")
    ret = qcng.compute(mbe_model, "qcmanybody", raise_error=True, return_version=retver)
    ret = checkver_and_convert(ret, request.node.name, "post")
    pprint.pprint(ret.model_dump(), width=200)

    assert ret.extras == {}, f"[w] extras wrongly present: {ret.extras.keys()}"

    refs = {
        "nocp_corrected_total_energy_through_1_body": -11.530668717083888,
        "nocp_corrected_total_energy_through_2_body": -11.522851206178828,
        "nocp_corrected_total_energy_through_3_body": -11.523095269671348,
        "nocp_corrected_total_energy_through_4_body": -11.523038093664368,
        "nocp_corrected_interaction_energy_through_1_body": 0.0,
        "nocp_corrected_interaction_energy_through_2_body": 0.007817510905059777,
        "nocp_corrected_interaction_energy_through_3_body": 0.0075734474125397355,
        "nocp_corrected_interaction_energy_through_4_body": 0.007630623419519367,
        "nocp_corrected_2_body_contribution_to_energy": 0.007817510905059777,
        "nocp_corrected_3_body_contribution_to_energy": -0.00024406349252004134,
        "nocp_corrected_4_body_contribution_to_energy": 5.717600697963121e-05,
    }

    ans = refs[anskey]
    ref_nmbe = calcinfo_nmbe
    atol = 1.0e-8

    for skp, ref in refs.items():
        if not "4_body" in skp:
            assert compare_values(ref, getattr(ret.properties, skp), atol=atol, label=f"[b] skprop {skp}")
        else:
            assert getattr(ret.properties, skp) is None

    assert compare_values(
        refs["nocp_corrected_total_energy_through_3_body"],
        ret.properties.nocp_corrected_total_energy,
        atol=atol,
        label=f"[d] skprop tot",
    )
    assert compare_values(
        refs["nocp_corrected_interaction_energy_through_3_body"],
        ret.properties.nocp_corrected_interaction_energy,
        atol=atol,
        label=f"[d] skprop IE",
    )

    assert compare_values(ans, ret.properties.return_energy, atol=atol, label=f"[f] skprop {skp}")
    assert compare_values(ans, ret.return_result, atol=atol, label=f"[g] ret")

    assert ret.properties.calcinfo_nmbe == ref_nmbe, f"[i] {ret.properties.calcinfo_nmbe} != {ref_nmbe}"
    ret_subres = getattr(ret, subres)
    assert len(ret_subres) == ref_nmbe, f"[k] {len(ret_subres)} != {ref_nmbe}; mbe protocol did not take"
    if ref_nmbe > 0:
        an_atres = next(iter(ret_subres.values()))
        assert an_atres.stdout is None, f"[l] atomic protocol did not take"


@uusing("qcmanybody")
@pytest.mark.parametrize(
    "qcprog",
    [
        pytest.param("cfour", marks=using("cfour")),
        pytest.param("gamess", marks=using("gamess")),
        pytest.param("nwchem", marks=using("nwchem")),
        pytest.param("psi4", marks=using("psi4")),
    ],
)
def test_bsse_ene_tu6_cp_ne2(schema_versions, request, qcprog):
    """
    from https://github.com/psi4/psi4/blob/master/tests/tu6-cp-ne2/input.dat
    Example potential energy surface scan and CP-correction for Ne2
    """
    models, retver, _ = schema_versions

    if from_v2(request.node.name):
        from qcmanybody.models.v2 import ManyBodyInput
    else:
        from qcmanybody.models.v1 import ManyBodyInput

    tu6_ie_scan = {2.5: 0.757717, 3.0: 0.015685, 4.0: -0.016266}  # Ang: kcal/mol IE

    keywords = {
        "cfour": {"frozen_core": True},
        "gamess": {"contrl__ispher": 1},
        "nwchem": {"ccsd__freeze__atomic": True, "basis__spherical": True},
        "psi4": {"freeze_core": True},
    }
    basis = {
        "cfour": "aug-pvdz",
        "gamess": "accd",
        "nwchem": "aug-cc-pvdz",
        "psi4": "aug-cc-pvdz",
    }

    mbe_data = {
        "specification": {
            "specification": {
                "model": {
                    "method": "ccsd(t)",
                    "basis": basis[qcprog],
                },
                "driver": "energy",
                "program": qcprog,
                "keywords": keywords[qcprog],
            },
            "keywords": {
                "bsse_type": "cp",
            },
            "driver": "energy",
        },
        "molecule": None,
    }

    for R in tu6_ie_scan:
        nene = models.Molecule(
            symbols=["Ne", "Ne"], fragments=[[0], [1]], geometry=[0, 0, 0, 0, 0, R / constants.bohr2angstroms]
        )
        mbe_data["molecule"] = nene

        mbe_model = ManyBodyInput(**mbe_data)
        if qcprog == "gamess":
            with pytest.raises(RuntimeError) as exe:
                qcng.compute(mbe_model, "qcmanybody", raise_error=True, return_version=retver)
            assert "GAMESS+QCEngine can't handle ghost atoms yet" in str(exe.value)
            pytest.xfail("GAMESS can't do ghosts")

        mbe_model = checkver_and_convert(mbe_model, request.node.name, "pre")
        ret = qcng.compute(mbe_model, "qcmanybody", raise_error=True, return_version=retver)
        ret = checkver_and_convert(ret, request.node.name, "post")
        pprint.pprint(ret.model_dump(), width=200)

        assert compare_values(
            tu6_ie_scan[R], ret.return_result * constants.hartree2kcalmol, atol=1.0e-4, label=f"CP-CCSD(T) [{R:3.1f}]"
        )
        assert compare(3, ret.properties.calcinfo_nmbe, label="nmbe")
        assert compare(1, ret.properties.calcinfo_nmc, label="nmc")
        assert compare(2, ret.properties.calcinfo_nfr, label="nfr")
        assert compare(2, ret.properties.calcinfo_natom, label="nat")


@uusing("qcmanybody")
def test_mbe_error(schema_versions, request):
    models, retver, _ = schema_versions

    if from_v2(request.node.name):
        from qcmanybody.models.v2 import ManyBodyInput

        AtomicSpecification = models.AtomicSpecification
    else:
        from qcmanybody.models.v1 import AtomicSpecification, ManyBodyInput

    mbe_data = {
        "specification": {
            "specification": AtomicSpecification(
                **{
                    "model": {
                        "method": "nonsense",
                        "basis": "nonsense",
                    },
                    "driver": "energy",
                    "program": "cms",
                }
            ),
            "keywords": {},
            "driver": "energy",
        },
        "molecule": None,
    }

    nene = models.Molecule(
        symbols=["Ne", "Ne"], fragments=[[0], [1]], geometry=[0, 0, 0, 0, 0, 3.0 / constants.bohr2angstroms]
    )
    mbe_data["molecule"] = nene

    mbe_model = ManyBodyInput(**mbe_data)

    # test 1
    with pytest.raises(RuntimeError) as exc:
        qcng.compute(mbe_model, "qcmanybody", raise_error=True, return_version=retver)

    assert "Program cms is not registered to QCEngine" in str(exc.value)

    # test 2
    ret = qcng.compute(mbe_model, "qcmanybody", return_version=retver)
    assert ret.success is False
    assert "Program cms is not registered to QCEngine" in ret.error.error_message


@uusing("psi4")
@uusing("qcmanybody")
@pytest.mark.parametrize(
    "optimizer,bsse_type,sio",
    [
        pytest.param("optking", "nil", "", marks=using("optking")),
        pytest.param("optking", "cp", False, marks=using("optking_v2")),
        pytest.param("optking", "nocp", False, marks=using("optking_v2")),
        pytest.param("geometric", "nil", "", marks=using("geometric")),
        # needs qcsk v2-compat geometric harness
        # pytest.param("geometric", "nocp", False, marks=using("geometric_genopt")),
        # pytest.param("geometric", "cp", True, marks=using("geometric_genopt")),
    ],
)
def test_bsse_opt_hf_trimer(schema_versions, request, optimizer, bsse_type, sio):
    models, retver, _ = schema_versions

    if bsse_type != "nil" and "as_v2" not in request.node.name:
        pytest.skip(
            "ManyBody Optimization is only available for QCSchema v2. The experimental v1 GeneralizedOptimization is retired."
        )

    Molecule = models.Molecule
    if from_v2(request.node.name):
        from qcmanybody.models.v2 import ManyBodyInput

        subptcl = "cluster_results"
    else:
        from qcmanybody.models.v1 import ManyBodyInput

        subptcl = "component_results"
    # to_v2
    subres = "cluster_results" if ("v2" in request.node.name) else "component_results"

    initial_molecule = Molecule.from_data(
        """
F         -0.04288        2.78905        0.00000
H          0.59079        2.03435        0.00000
--
F         -1.94320       -0.70822        0.00000
H         -1.60642        0.21789       -0.00000
--
F          2.03569       -0.60531       -0.00000
H          1.06527       -0.77673        0.00000
units ang
"""
    )

    at_spec = {
        # schema_name needed for differentiation in genopt
        "schema_name": "qcschema_atomic_specification" if from_v2(request.node.name) else "qcschema_input",
        "program": "psi4",
        "driver": "gradient",
        "model": {
            "method": "hf",
            "basis": "6-31g",
        },
        "keywords": {
            "scf_type": "df",
        },
    }
    if not from_v2(request.node.name):
        at_spec.pop("program")
        at_spec.pop("driver")

    mbe_spec = {
        # schema_name needed for differentiation in genopt
        "schema_name": "qcschema_many_body_specification"
        if from_v2(request.node.name)
        else "qcschema_manybodyspecification",
        "specification": {
            "model": {
                "method": "hf",
                "basis": "6-31g",
            },
            "driver": "energy",
            "program": "psi4",
            "keywords": {},
            "protocols": {
                "stdout": False,
            },
            "extras": {
                "psiapi": True,
            },
        },
        "keywords": {
            "bsse_type": bsse_type,
            "supersystem_ie_only": sio,
        },
        "driver": "energy",
        "protocols": {
            subptcl: "all",
        },
    }

    if from_v2(request.node.name):
        opt_data = {
            "initial_molecule": initial_molecule,
            "specification": {
                "specification": at_spec if (bsse_type == "nil") else mbe_spec,
                "keywords": {
                    "g_convergence": "nwchem_loose",
                },
                "protocols": {
                    "trajectory_results": "initial_and_final",
                },
            },
        }
    else:
        opt_data = {
            "initial_molecule": initial_molecule,
            "input_specification": at_spec if (bsse_type == "nil") else mbe_spec,
            "keywords": {
                "program": "psi4",
                "g_convergence": "nwchem_loose",
            },
            "protocols": {
                "trajectory": "initial_and_final",
            },
        }

    opt_data = checkver_and_convert(opt_data, request.node.name, "pre", cast_dict_as="OptimizationInput")
    ret = qcng.compute(opt_data, optimizer, raise_error=True, return_version=retver)
    ret = checkver_and_convert(ret, request.node.name, "post")

    pprint.pprint(ret.model_dump(), width=200)

    r_fh_hb_xptd = {
        "nil": 2.18 / constants.bohr2angstroms,
        "nocp": 2.18 / constants.bohr2angstroms,
        "cp": 2.27 / constants.bohr2angstroms,
    }[bsse_type]
    r_fh_computed = ret.final_molecule.measure([1, 3])
    assert (
        pytest.approx(r_fh_computed, 1.0e-2) == r_fh_hb_xptd
    ), f"hydrogen bond length computed ({r_fh_computed}) != expected ({r_fh_hb_xptd})"
    ret_traj = ret.trajectory_results if "v2" in request.node.name else ret.trajectory
    assert (
        len(ret_traj) == 2
    ), f"trajectory protocol did not take. len(ret.trajectory)={len(ret_traj)} != 2 (initial_and_final)"
    if bsse_type != "nil":
        xptd_nmbe = {
            ("nocp", False): 7,
            ("nocp", True): 4,
            ("cp", False): 10,
            ("cp", True): 7,
        }[(bsse_type, sio)]
        ret_last_subres = getattr(ret_traj[-1], subres)
        assert xptd_nmbe == len(ret_last_subres), f"mbe protocol did not take"
        assert ret_last_subres['["(auto)", [1, 2, 3], [1, 2, 3]]'].stdout is None, f"atomic protocol did not take"


@uusing("qcmanybody")
@pytest.mark.parametrize("schver", [2])  # GeneralizedOptimization retired (this was the v1 precursor)
@pytest.mark.parametrize("bsse_type", ["mbe", "ssfc"])  # aka nocp, cp
@pytest.mark.parametrize(
    "qcprog,qc_keywords",
    [
        pytest.param("psi4", {}, marks=using("psi4")),
        pytest.param("cfour", {}, marks=using("cfour")),
        pytest.param("nwchem", {}, marks=using("nwchem")),
    ],
)
@pytest.mark.parametrize(
    "optimizer,opt_keywords",
    [
        pytest.param("optking", {}, marks=using("optking_v2")),
        # pytest.param(
        #     "geometric", {"convergence_set": "interfrag_tight", "maxiter": 15}, marks=using("geometric_genopt")
        # ),
    ],
)
def test_bsse_opt_lif_dimer(optimizer, opt_keywords, bsse_type, qcprog, qc_keywords, schver):
    # in geomeTRIC: test_lif_bsse

    lif = {
        "symbols": ["Li", "F"],
        "geometry": [0, 0, 0, 0, 0, 3],
        "fragments": [[0], [1]],
        "fragment_charges": [+1, -1],
    }

    mbe_subprop = "component_properties" if schver == 1 else "cluster_properties"
    mbe_subptcl = "component_results" if schver == 1 else "cluster_results"
    opt_subptcl = "trajectory" if schver == 1 else "trajectory_results"

    mbe_spec = {
        "schema_name": "qcschema_many_body_specification" if schver == 2 else "qcschema_manybodyspecification",
        "specification": {
            "model": {
                "method": "hf",
                "basis": "6-31G",
            },
            "driver": "energy",
            "program": qcprog,
            "keywords": qc_keywords,
            "protocols": {
                "stdout": False,
            },
        },
        "keywords": {
            "bsse_type": bsse_type,
            "supersystem_ie_only": True,
        },
        "protocols": {
            mbe_subptcl: "all",
        },
        "driver": "energy",
        "program": "qcmanybody",
    }

    if schver == 2:
        opt_data = {
            "initial_molecule": lif,
            "specification": {
                "specification": at_spec if (bsse_type == "nil") else mbe_spec,
                "keywords": {
                    **opt_keywords,
                },
                "protocols": {
                    opt_subptcl: "final",
                },
            },
        }
    else:
        opt_data = {
            "initial_molecule": lif,
            "input_specification": mbe_spec,
            "keywords": {
                "program": "nonsense",
                **opt_keywords,
            },
            "protocols": {
                opt_subptcl: "final",
            },
        }

    ret = qcng.compute_procedure(opt_data, optimizer, raise_error=True, return_version=schver)

    # printing will show up if job fails
    pprint.pprint(ret.model_dump(), width=200)

    assert ret.success

    ret_traj = ret.trajectory_results if schver == 2 else ret.trajectory
    assert len(getattr(ret_traj[0], mbe_subprop)) == (5 if bsse_type == "ssfc" else 3)

    assert ret.provenance.creator.lower() == optimizer
    assert ret_traj[0].provenance.creator == "QCManyBody"
    atres = list(getattr(ret_traj[0], mbe_subptcl).values())[0]
    assert atres.provenance.creator.lower() == qcprog

    Rlif = ret.final_molecule.measure([0, 1])
    Rref = 3.016 if bsse_type == "ssfc" else 2.969
    assert compare_values(Rlif, Rref, "bond length", atol=1.0e-3)
