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

from chemsh import *
from os import environ

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

# save as .XYZ for visualising
pd.save('pd_unitcell.xyz')

# 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')
# save as .XYZ for visualising
supercell.save('pd_cluster.xyz')

# 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')
# save as .XYZ for visualising
o2_pd.save('o2_pd.xyz')

# 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',
             #SZV-MOLOPT-SR-GTH/DZVP-MOLOPT-SR-GTH-q18, using doubleZ basis to be used with smear=on
             basis             = {'Pd':'DZVP-MOLOPT-SR-GTH-q18', 'O':'TZV2P-MOLOPT-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           = (2,4,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_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)
             smear             = True,             # smear is deafault to False. so switch on is necessary now
             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(" >>> GULP_LIB must be set before running this calculation.", flush=True)
    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',
#               adsorbate =  indices_O2,    # the values of indices_O2, assigned internally when binding, will be used by default if not specified otherwise
#               surface   = 'ab',           # this is the default
               coupling  = 'ionic',
              )

# save as .XYZ for visualising
my_qmmm.qm.frag.save('_qm.xyz')

# created a geometry optimisation instance
my_opt = Opt(theory=my_qmmm, maxcycle=2000, maxene=2000)

# dryrun=False: this is the default, only input files are generated if dryrun=True
my_opt.run(dryrun=False)

# HINT: the performance of parallel CP2K can vary, depending on platforms. we sugget not using
#       OpenMP multithread parallelism for the current calculation by, for example:
# chemsh -np 16 -nt 1 o2_pd_cp2k.py

# on 16 modern CPU cores the calculation can typically finish within 3 hours
# the energy/gradients profile during the optimisaiton is saved in a file "_profile_opt.npz"
# which can be visualised by command (if matplotlib is installed on your computer):
# chemsh --plot _profile_opt.npz --units kcal/mol




