import psi4
import optking
from .utils import utils

_schver = 2 if utils.psi4_runs_v2_qcschema(psi4.__version__) else 1

# HF SCF CC-PVDZ geometry optimization of HOOH with Z-matrix input
def test_B_dB_matrices(check_iter):
    hooh = psi4.geometry(
        """
      H
      O 1 0.9
      O 2 1.4 1 100.0
      H 3 0.9 2 100.0 1 114.0
    """
    )

    psi4.core.clean_options()
    psi4_options = {
        "basis": "cc-pvdz",
        "g_convergence": "gau_tight",
        "scf_type": "pk",
        "TEST_B": True,
        "TEST_DERIVATIVE_B": True,
    }
    psi4.set_options(psi4_options)
    json_output = optking.optimize_psi4("hf")  # Uses default program (psi4)

    if _schver == 1:
        E = json_output["energies"][-1]  # TEST
        nucenergy = json_output["trajectory"][-1]["properties"]["nuclear_repulsion_energy"]  # TEST
    elif _schver == 2:
        E = json_output["trajectory_properties"][-1]["return_energy"]  # TEST
        nucenergy = json_output["trajectory_results"][-1]["properties"]["nuclear_repulsion_energy"]  # TEST
    refnucenergy = 38.06177  # TEST
    refenergy = -150.786766850  # TEST
    if _schver == 1:
        assert "test_b" in json_output["keywords"]  # TEST
        assert "test_derivative_b" in json_output["keywords"]  # TEST
        assert "g_convergence" in json_output["keywords"]  # TEST
    elif _schver == 2:
        assert "test_b" in json_output["input_data"]["specification"]["keywords"]  # TEST
        assert "test_derivative_b" in json_output["input_data"]["specification"]["keywords"]  # TEST
        assert "g_convergence" in json_output["input_data"]["specification"]["keywords"]  # TEST
    assert psi4.compare_values(refnucenergy, nucenergy, 3, "Nuclear repulsion energy")  # TEST
    assert psi4.compare_values(refenergy, E, 8, "Reference energy")  # TEST
    utils.compare_iterations(json_output, 7, check_iter)


def test_maxiter(check_iter):
    h2o = psi4.geometry(
        """
     O
     H 1 1.0
     H 1 1.0 2 104.5
    """
    )

    psi4.core.clean_options()
    psi4options = {
        # Throw a bunch of options at psi4
        "diis": 0,
        "basis": "STO-3G",
        "e_convergence": 1e-10,
        "d_convergence": 1e-10,
        "scf_type": "PK",
        "geom_maxiter": 2,
    }
    psi4.set_options(psi4options)

    json_output = optking.optimize_psi4("hf")

    if _schver == 1:
        assert "geom_maxiter" in json_output["keywords"]  # TEST
    elif _schver == 2:
        assert "geom_maxiter" in json_output["input_data"]["specification"]["keywords"]  # TEST
    assert "Maximum number of steps exceeded" in json_output["error"]["error_message"]  # TEST
    assert "OptError" in str(json_output["error"]["error_type"])  # TEST


# Test the energy of geometry output, when maxiter is reached.
def test_maxiter_geom():
    h2o = psi4.geometry(
        """
        O
        H 1 1.0
        H 1 1.0 2 104.5
    """
    )

    psi4.core.clean_options()
    psi4options = {
        "basis": "cc-pvdz",
        "e_convergence": 10,
        "d_convergence": 10,
        "scf_type": "pk",
        "geom_maxiter": 2
    }
    psi4.set_options(psi4options)

    result = optking.optimize_psi4("hf")
    print(result)

    nextStepSchema = result["final_molecule"]  # TEST
    nextStepMolecule = psi4.core.Molecule.from_schema(nextStepSchema)  # TEST
    psi4.core.set_active_molecule(nextStepMolecule)  # TEST
    psi4.set_options(psi4options)
    nextStepEnergy = psi4.driver.energy("scf/cc-pvdz")  # TEST
    REF_energy = -76.0270381300  # TEST
    REF_energy = -76.0270521  # TEST changed. Energy different now that update occurs
    assert psi4.compare_values(REF_energy, nextStepEnergy, 5, "Energy of next-step molecule")
    # TEST
