# TEMPORARY REGRESSION TEST FOR THE QM/Me SCHEME
# TODO Create a cut down version (ideally an opt) that captures the essential method
# but runs in <2 mins on a single core.

# The location of the GULP potential library should be specified
# in the GULP_LIB environment variable before running this test.

from chemsh import *
from chemsh.utils import testutils
from os import environ

# read in the Pd (nonprimitive) unit cell
pd = Fragment(coords='pd_unitcell.pun', connect_mode='ionic')

# build a supercell by expanding the unit cell
# nlayers=4      : will be 4 extra cell along each axis plus the original unit cell (default: 1)
# symmetric=False: along the axes' positive directions only (default: True)
# redundant=False: centres on the (1¬00), (01¬0), and (001¬) surfaces of the supercell are removed as redundant (default: False)

supercell = pd.getSuperCell(nlayers=4, symmetric=False, redundant=False)
supercell.save('pd_cluster.pun')

# attach a O2 molecule on to the Pd surface
# distance=1.0: distance to surface, in Angstrom here (default: 3.0 Bohr)
# angle=90.0  : angle of binding (default: 180 degrees)
# surface='ab': (1,0,0) surface (default: 'ab')

o2_pd = supercell.bind('O2', distance=1.0, angle=90.0, unit='angstrom', surface='ab')
o2_pd.save('o2_pd.pun')

# get the oxygen atoms' indices
indices_O2 = o2_pd.select(atoms='O', strict=False)

# choose the first atom (note: Python starts from 0)
O1 = indices_O2[0]

# select 2D slabs to form the QM region
# dimensions=(4,8,3): select slabs 4 atoms wide, 8 atoms long, and 3 layers deep
#                     NB: the first two numbers must be even because it is face-centred cell!
# origin=O1         : choose the origin
# visualise='vmd'   : print selection clause to be used for VMD representation
slabs = o2_pd.selectPeriodic(dimensions=(4,8,3), origin=O1, surface='ab', visualise='vmd')

# merge the two lists of adsorbate and slabs to form indices of the the QM atoms
# NB: indices_O2 and slabs are NumPy arrays and they must be converted to Python list to use the '+' concatenation operator
qm_atoms = indices_O2.tolist() + slabs.tolist()

# optional keyword options:
#             basis_file       = 'BASIS_MOLOPT',   # automatically assigned by ChemShell if not specified
#             potential_file   = 'POTENTIAL',      # automatically assigned by ChemShell if not specified
#             kpoints_method   = 'Monkhorst-Pack', # this is the default
my_qm = CP2K(functional        = 'pbe',
        # Note some basis sets provide insufficient number of orbitals to use nmos_added = n. Use a suitable basis set when mos added, e.g.,
             #basis            = {'Pd':'DZVP-MOLOPT-SR-GTH-q18', 'O':'TZV2P-MOLOPT-GTH'},
             # A simple basis set is used for easy convergence of the test run. 
             #Use a higher basis set for production runs along with using smearing.
             # basis           = {'Pd':'DZVP-MOLOPT-SR-GTH-q18/SZV-MOLOPT-SR-GTH', 'O':'TZV2P-MOLOPT-GTH'},
             
             basis             = {'Pd':'SZV-MOLOPT-SR-GTH', 'O':'SZV-MOLOPT-SR-GTH'},
             pseudopotential   = 'GTH-PBE',
             # NB: if CP2K is going to be run as system command (but not from the directlly linked-in ChemShell module)
             #sysmpi           = 'mpirun -n 16',    # system MPI commmand requesting 16 MPI processes
             #path             = 'cp2k',           # system commmand to run CP2K
             restart           =  False,           # a new job
             kpoints           = (4,2,1),          # according to the 2D slabs selected above
             mixing_method     = 'BROYDEN_MIXING', # Broyden is strongly recommended to use to get better SCF convergence for adsorbate+surface systems (default: DIRECT_P_MIXING)
             #mixing_method    = 'DIRECT_P_MIXING',
             mixing_alpha      = 0.2,              # smaller mixing percentage for a smoother SCF convergence (default: 0.4)
             maxiter           = 1000,             # maximum number of SCF cycles
             nmos_added        = 20,               # add virtual orbials to describe more than just the conduction band (default: 10)
             grid_cutoff       = 400.0,            # increase the grid's accuracy (default: 280.0)
             grid_real_cutoff  = 50.0,             # increase the real space grid's accuracy (default: 40.0)

             # set smear = True , whenever smearing is used. Default is smear=False
             #smear             = True,
             #Turned off here for a quick run  with SZ.. basis set
             smear             = False,
             smear_temperature = 500.0,            # increase the smearing temperature to get easier SCF convergence (default: 293.15)
            )

my_mm = GULP(ff = 'meam_2nn')    # 2nd nearest-neighbour modified embedded atom method (2NN MEAM) forcefield

# Check that the environmental variable GULP_LIB has been set
# NB: Do not set it in this input file, as all MPI processes need to see it.
if "GULP_LIB" not in environ.keys():
    print(" >>> Unsupported test: GULP_LIB must be set first.", flush=True)
    testutils.setUnsupported()
    exit()

# we specify the fragment instance to only the QMMM instance
my_qmmm = QMMM(frag      =  o2_pd,
               qm_region =  qm_atoms,
               qm        =  my_qm,
               mm        =  my_mm,
               scheme    = 'qm/me'
              )

# created a geometry optimisation instance
my_sp = SP(theory=my_qmmm, gradients=True)
my_sp.run()

ecalc = my_sp.result.energy

# Regression test (CP2K 2022 on Ubuntu)
eref = -100.5653596942
testutils.validate(ecalc, eref)
