# MgO+CO2 cluster geometry optimisation

# Initialise ChemShell
from chemsh import *
from chemsh.utils import testutils
from numpy import append as np_append

# Merge MgO surface cluster and CO2 geometries
mgo = Fragment(coords='mgo_shells_opt.pun', connect_mode=None)
co2 = Fragment(coords='co2.xyz', connect_mode=None)
mgo.append(co2)
mgo.save('mgoco2start.pun','pun')

# QM region contains MgO cluster region 1 and CO2
co2_atoms = [2370, 2371, 2372]
qm_region = np_append(mgo.getRegion(1), co2_atoms)
print("QM region is: ", qm_region)

# Get MgO charge from input fragment
# CO2 is neutral so does not affect QM charge
qm_charge = round(mgo.totalcharge)
print("QM charge: ", qm_charge)

# Active atoms for optimisation are regions 1-3 plus CO2
active_region = np_append(mgo.getRegion(1,2,3), co2_atoms)
print("Active region is: ", active_region)

# QM/MM settings
qm_theory = NWChem(method='dft',
                   functional='b3lyp',
                   basis='mgoco2_nwchem.basis',
                   ecp='mgo_nwchem.ecp',
                   charge=qm_charge,
	    	   path="mpirun -np 4 nwchem")

mm_theory = GULP(ff='mgo_2body.ff',
                 conjugate=True)

qmmm = QMMM(qm=qm_theory,
            mm=mm_theory,
            frag=mgo,
            qm_region=qm_region,
            coupling='ionic',
            shl_maxcycles=5)

# Run QM/MM optimisation
opt = Opt(theory=qmmm,
          algorithm="lbfgs", trust_radius="energy",
          maxcycle=200, maxene=300, active=active_region)
opt.run()

ecalc = opt.result.energy

# Save optimised geometry
mgo.save('mgoco2end.pun', 'pun')
mgo.save('mgoco2end.xyz', 'xyz')

# Validate result against pre-calculated reference energy
eref = -1272.007579375473
tolerance = 1.0e-6

testutils.validate(ecalc, eref, tol=tolerance)
