# chorismate mutase single-point QM/MM energy (AM1/AMBER)

from chemsh import *
from chemsh.utils import testutils
from chemsh.io import file

# Parse charges from AMBER prmtop
prmtop = file.parse('cm.prmtop', fmt='amber_prmtop')
for key, val in prmtop.items():
    if type(val) is dict and 'atom' in val:
        charges = prmtop[key]['atom']['data']['charge']

# Import coords from trajectory file and charges from prmtop
cm = Fragment(coords='cm.pdb', trajectory='cm.rst', charges=charges)

# QM region is chorismate (only)
indices_qm_region = cm.select(residues='cho')
frag_qm_region = cm.getSelected(indices_qm_region)

# QM charge obtained from FF partial charges
qm_charge = round(frag_qm_region.charges.sum())
print("\n QM charge =", qm_charge)

# Semi-empirical AM1
qm = MNDO(method='am1', charge=qm_charge) 

# Import AMBER forcefield into DL_POLY
# TODO consider changing default rcut so it doesn't have to be specified here
mm = DL_POLY(ff='cm.prmtop', rcut=99.99)

# Create a QM/MM instance
qmmm = QMMM(frag=cm, qm=qm, mm=mm, qm_region=indices_qm_region)

# Run a single-point energy evaluation
sp = SP(theory=qmmm)
sp.run()

ecalc = sp.result.energy

# Regression test result
eref = -14.249000257881
tolerance = 1.0e-7

testutils.validate(ecalc, eref)
