#! SCF STO-3G geometry optimzation, with Z-matrix input
import os
import pathlib
import json
import qcelemental as qcel
import psi4
import optking
from .utils import utils

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


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

    psi4.core.clean_options()
    psi4_options = {
        "diis": False,
        "basis": "sto-3g",
        "e_convergence": 10,
        "d_convergence": 10,
        "scf_type": "pk",
        "print_trajectory_xyz_file": True,
    }
    psi4.set_options(psi4_options)

    json_output = optking.optimize_psi4("hf")

    if _schver == 1:
        E = json_output["energies"][-1]
        nucenergy = json_output["trajectory"][-1]["properties"]["nuclear_repulsion_energy"]  # TEST
    elif _schver == 2:
        E = json_output["trajectory_properties"][-1]["return_energy"]
        nucenergy = json_output["trajectory_results"][-1]["properties"]["nuclear_repulsion_energy"]  # TEST
    refnucenergy = 8.9064983474  # TEST
    refenergy = -74.9659011923  # TEST
    assert psi4.compare_values(refnucenergy, nucenergy, 3, "Nuclear repulsion energy")  # TEST
    assert psi4.compare_values(refenergy, E, 6, "Reference energy")  # TEST
    utils.compare_iterations(json_output, 4, check_iter)

    ##
    # Test writing a trajectory file
    ##
    traj_file = pathlib.Path(f'opt_traj.{os.getpid()}.xyz')
    assert traj_file.exists()

    with traj_file.open() as f:
        lines = f.readlines()
    # Remove before assertions
    os.system(f'rm {str(traj_file)}')

    first = "".join(lines[:5])
    last = "".join(lines[-5:])
    first_mol = qcel.molparse.from_string(first, dtype='xyz')
    last_mol = qcel.molparse.from_string(last, dtype='xyz')

    assert first_mol.get('qm').get('elem').tolist() == ['O', 'H', 'H']
    assert last_mol.get('qm').get('elem').tolist() == ['O', 'H', 'H']


#! SCF cc-pVDZ geometry optimzation, Z-matrix input, tight convergence
def test_hf_g_h2o_tight(check_iter):
    h2o = psi4.geometry(
        """
        O
        H 1 1.0
        H 1 1.0 2 104.5
    """
    )

    psi4.core.clean_options()
    psi4_options = {
        "diis": False,
        "basis": "cc-pvdz",
        "e_convergence": 10,
        "d_convergence": 10,
        "g_convergence": "gau_tight",
        "scf_type": "pk",
    }
    psi4.set_options(psi4_options)

    json_output = optking.optimize_psi4("hf")

    if _schver == 1:
        E = json_output["energies"][-1]
    elif _schver == 2:
        E = json_output["trajectory_properties"][-1]["return_energy"]
    REF_energy = -76.02705351276  # TEST
    assert psi4.compare_values(REF_energy, E, 8, "RHF energy")  # TEST
    utils.compare_iterations(json_output, 6, check_iter)


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

    psi4.core.clean_options()
    psi4_options = {
        "basis": "cc-pvtz",
        "e_convergence": 10,
        "d_convergence": 10,
        "g_convergence": "gau_tight",
        "scf_type": "pk",
    }
    psi4.set_options(psi4_options)

    json_output = optking.optimize_psi4("hf")

    if _schver == 1:
        E = json_output["energies"][-1]
    elif _schver == 2:
        E = json_output["trajectory_properties"][-1]["return_energy"]
    REF_energy = -76.05776970191  # TEST
    assert psi4.compare_values(REF_energy, E, 8, "RHF energy")  # TEST
    utils.compare_iterations(json_output, 6, check_iter)


#! SCF cc-pVDZ geometry optimzation of ketene, starting from bent structure
def test_hf_g_ketene(check_iter):
    ketene = psi4.geometry(
        """
        0 1
        H
        C 1 1.1
        C 2 1.3 1 130.0
        H 2 1.1 3 110.0 1 180.0
        O 3 1.1 4 120.0 1 180.0
    """
    )

    psi4.core.clean_options()
    psi4_options = {"basis": "cc-pvdz", "g_convergence": "gau_tight", "scf_type": "pk"}
    psi4.set_options(psi4_options)

    result = optking.optimize_psi4("scf")
    optking.logger.info(json.dumps(result, indent=2))
    if _schver == 1:
        E = result["energies"][-1]
    elif _schver == 2:
        E = result["trajectory_properties"][-1]["return_energy"]
    REF_energy = -151.7410313803  # TEST
    assert psi4.compare_values(REF_energy, E, 8, "RHF energy")  # TEST
    utils.compare_iterations(result, 7, check_iter)
