from chemsh import *
from chemsh.utils import testutils
import numpy as np
# Get chemshell to throw an unsupported error if pyscf not installed
try:
    import pyscf
except ModuleNotFoundError:
    from chemsh.base import errors
    testutils.setUnsupported()
    exit(errors.UNSUPPORTED)
from chemsh.interfaces.pyscf.ComputeElectrostaticEmbedding import (
    ComputeElectrostaticEmbedding,
)

ang2bohr = 1.8897259886
np.set_printoptions(precision=8, suppress=True, linewidth=1000)
qm_xyz = ang2bohr * np.array(
    [
        [0.000000, 0.000000, 0.000000],
        [0.000000, -0.751841, 0.568201],
        [-0.000000, 0.751841, 0.568201],
    ]
)
mol = pyscf.gto.M(
    atom=f"O {qm_xyz[0,0]} {qm_xyz[0,1]} {qm_xyz[0,2]}; H {qm_xyz[1,0]} {qm_xyz[1,1]} {qm_xyz[1,2]}; H {qm_xyz[2,0]} {qm_xyz[2,1]} {qm_xyz[2,2]}",
    basis="3-21G",
    unit="Bohr",
    symmetry=False,
    charge=0,
)

xyz = ang2bohr * np.array(
    [
        -2.511474,
        0.000000,
        -1.450000,
        -1.651960,
        0.000000,
        -1.063537,
        -2.374291,
        0.000000,
        -2.382362,
    ]
).reshape(-1, 3)


charges = [-0.834, 0.417, 0.417]
e, qm_grad, point_charge_grad = ComputeElectrostaticEmbedding(
    mol, mol.RHF(), xyz, charges
)
orca_qm_grad = np.array(
    [
        [0.00834968, -0.0, 0.02795187],
        [-0.00066211, 0.02602692, -0.01190286],
        [-0.00066211, -0.02602692, -0.01190286],
    ]
)

orca_grad = np.array(
    [
        [0.00766456, 0.0, 0.00114672],
        [-0.01236124, -0.0, -0.00395691],
        [-0.00232876, -0.0, -0.00133596],
    ]
)

if np.allclose(orca_grad, point_charge_grad):
    testutils.setPassed()
else:
    testutils.setFailed()
