#  Copyright (C) 2017 The authors of Py-ChemShell
#
#  This file is part of Py-ChemShell.
#
#  Py-ChemShell is free software: you can redistribute it and/or modify
#  it under the terms of the GNU Lesser General Public License as
#  published by the Free Software Foundation, either version 3 of the
#  License, or (at your option) any later version.
#
#  Py-ChemShell is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU Lesser General Public License for more details.
#
#  You should have received a copy of the GNU Lesser General Public
#  License along with Py-ChemShell.  If not, see
#  <http://www.gnu.org/licenses/>.

__author__ = ['Matt Farrow <m.farrow@ucl.ac.uk>', 
              'You Lu <you.lu@stfc.ac.uk>',
              'Tom Keal <thomas.keal@stfc.ac.uk>']

import chemsh
from numpy            import full
from ..qm             import _QM
from ...dl_py2f       import py2f
from ...base          import errors
from ...parallel      import parallel
from ...dl_py2f.utils import fileutils, miscutils, modutils

class GAMESS_UK(_QM):
    '''Interface to GAMESS-UK'''


    # Will try and keep keywords in alphabetical order. 
    _kwargs = { 
                'adaptation'   : False,
                'bq_grads'     :'_gamess-uk.bq_grads',
                'bypass'       :'no',
# YL: 'dipole' is a matrix tag in Tcl-ChemShell
#     see http://www.cse.scitech.ac.uk/ccg/software/chemshell/manual/gamess.html
#     this concerns 'punch' (see below)
                'dipole'       :'',
                'ed0'          :'',
                'ed2'          :'_gamess-uk.ed2',
                'ed3'          :'_gamess-uk.ed3',
                'ed4'          :'',
                'ed7'          :'_gamess-uk.ed7',
                'enter'        :'1',
                'harmonic'     : False,
                'input'        :'_gamess-uk.in',  # '_gamess-uk.inp' won't work
                'level'        :'',
                'level2'       : 0.5,
                'maxiter'      : 50,
                'memory'       : 52000000,
                'newscf'       :'',
                'output'       :'_gamess-uk.out',
                'potential'    : False,
                'punch'        :'_gamess-uk.pun',
                'punchlevel'   :'eandg',
                'quadrature'   : None,
                'runargs'      :'',
                'scf_threshold': 7,
                'shift_diis'   : False,
                'symmetry'     : False,
                'vectors'      :'atoms',
                'zmatrix'      :''
              }

    _internals = {
#                   '_ecpsection':'bq',
                   '_graphics' : False,
                   '_syscmd'   :'rungamess', # default system command to run GAMESS-UK
                 }

    # overloaded QM.runargs
    @property
    def runargs(self):
        '''Process arguments to be passed to rungamess'''

        # do nothing more than just returning! define `runargs` only in the setter! otherwise its length will grow in every cycle!
        return self._runargs


    # overloaded QM.runargs.setter
    @runargs.setter
    def runargs(self, val):
        '''Setter of runargs'''

        self._runargs = val

        selectcases = {
                        'ed2':'-t',
                        'ed3':'-k',
                        'ed4':'-k',
                        'ed7':'-t',
                      }

        # see GAMESS-UK manual: http://www.cfs.dl.ac.uk/docs/html/part13/node3.html
        # default: rungamess -k ed3 _gamess-uk
#        for edx in [ 'ed0', 'ed2', 'ed3', 'ed4', 'ed7' ]:
        for edx in selectcases.keys():
            if getattr(self, edx):
                self._runargs += ' '+selectcases.get(edx, ' ')+' ' + edx+'='+getattr(self, edx)

        self._runargs += ' ' + fileutils.getBaseName(self.input)


    # must be a procedure rather than a direct definition, otherwise frag.natoms is unknown
    def setOutKeys(self):
        '''Define {task:keyword} for searching in output file'''

        # each keyword should contain an array of integers specifying:
        #     [ row_0, # starting row
        #       row_n, # ending row
        #       row_i, # row increment
        #       col_0, # starting col
        #       col_n, # ending col
        #       col_i, # col increment
        #       skip , # skip the first (skip-1) match,
        #       coeff, # unit conversion coefficient
        #     ]
        # referencing to the line where the key is matched
#       if self._gradients:
# YL: not in use? see parseAdditionalOutput() below
        self._outkeys = {
                          'energy'    :{'  total energy      =':[0,0,1,3,3,1,1,1.0]},
                          'gradients' :{'gradient (hartree':[4,4+self.frag.natoms-1,1,6,8,1,1,1.0]},
                        }
#        else:
#            self._outkeys = {
#                              'energy'   :{'   total energy                     ':[0,0,1,2,2,1,1,1.0]}
#                            }

    def getInpBuff(self):
        '''String strbuff containing the input information'''

        from numpy import __version__
        # TD 18/06/2024 : Numpy v2.0.0 version dependence
        from numpy.lib import NumpyVersion
        if NumpyVersion(__version__) >= '2.0.0b1':
            from numpy._core import defchararray
        else:
            from numpy.core import defchararray

        # MRF This function call should be moved somewhere better
        # YL: moved to property(self.runargs())
#        self.commandLineInput()

        strbuff = 'memory ' + str(self.memory) + '\n'
        strbuff += 'time 9999999\n'


# TODO: only apply this when nprocs > 1 and directly linked
#        strbuff += 'parallel iomode nzo\n'
        strbuff += 'noecho\n'
        
        # restart I
        if self.restart:
            # read in the geometry from input (rather than from dumpfile)
            strbuff += 'restart new\n'

            # 'bypass two' for shells relaxation
            if self._gradients == 'shells':
                self.bypass = 'two'
                strbuff += 'bypass '+self.bypass+'\n'



        strbuff += 'title\n'
        strbuff += 'ChemShell GAMESS-UK interface ' + self.frag.formula + '\n'
        strbuff += 'noprint vectors\n'
        strbuff += 'noprint basis\n'
        strbuff += 'noprint dist geom\n'
        if not self.pop_analysis:
            strbuff += 'noprint analysis\n'
        strbuff += 'noprint bq\n'
#        strbuff += 'mfile memory\n'

        if self.harmonic:
            strbuff += 'harmonic' +'\n'

        # punch file control
        selectcases = {
                        'eandg'       :'grad',
                        'hessian'     :'grad secd',
                        'wavefunction':'vect '+ self.vectors + ' occu eige',
                        'potential'   :'grid 151',
                      }

        # punch 'punch high energy coor tran title' anyway
        # YL: not sure if 'potential' would work
        try:
            strbuff += 'punch high energy coor tran title ' + selectcases[self.punchlevel.lower()] + self.dipole + '\n'
        except KeyError:
            raise errors.ChemShValError('punch', selectcases.keys())

        # symmetry
        if self.symmetry:
            strbuff += '#use symmetry to add'
        else:
            strbuff += 'nosym\n'

        if self.scftype == 'mcscf':
            strbuff += 'super off\n'
        
        if not self.adaptation:
            strbuff += 'adapt off\n'

        nelectrons = self.frag.znums.sum() - int(self.charge)
        strbuff += 'elec ' + str(nelectrons) + '\n'
        
        # multiplicity
        strbuff += 'mult ' + str(self.mult) + '\n'

        # geometry
        if self.zmatrix:
            strbuff += self.zmatrix + '\n'
        else:
            strbuff += 'geometry\n'
            # allow the user to use element symbols instead of FF names
            if self.atom_names.lower().startswith('symb'):
                strbuff += self.frag.coords2str(order='cza', coords=14, charges=None, atoms='symbols', indent=2*' ')
            else:
                strbuff += self.frag.coords2str(order='cza', coords=14, charges=None, atoms='names', indent=2*' ')

            # bq
            if self.frag.nbqs:
                # ECPs
                if self.ecps:
                    # YL 13/09/2020: BQs.select() has been removed
                    ecprange = self.frag.bqs.selectByRegions(self.ecp_region)
                    # if qmmm.separate_ecps (see below): region 2 BQ charges subtracted with self._ecp_charge
                    if self.separate_ecps:
                        self.frag.bqs.charges[ecprange] -= self._ecp_charge

                # prefix with 'bq_'
                with_bq_prefix = defchararray.startswith(defchararray.lower(self.frag.bqs.names), b'bq')
                bqnames = defchararray.add(full((self.frag.nbqs,), '', dtype='S8'), self.frag.bqs.names)
                bqnames[~with_bq_prefix] = defchararray.add(full((self.frag.nbqs,), 'bq_', dtype='S8')[~with_bq_prefix], self.frag.bqs.names[~with_bq_prefix])

                # suffix with '_s'
                if self.frag.nshells:
                    shellrange = self.frag.bqs.stacks["Shells"]
                    bqnames[shellrange] = defchararray.add(bqnames[shellrange], full((len(shellrange),), '_s', dtype='S2'))
                
                strbuff += self.frag.bqs.coords2str(coords=14, charges=8, atoms=None, indent=2*' ', inserts=[(4,bqnames)])

                # separate boundary ECPs
                if self.ecps and self.separate_ecps:
                    # restore the region 2 BQ charges
                    self.frag.bqs.charges[ecprange] += self._ecp_charge
                    ecpnames = defchararray.add(full((self.ecps.natoms,), 'bq_', dtype='S3'), self.ecps.names)
                    ecpnames = defchararray.add(ecpnames, full((self.ecps.natoms,), '_e', dtype='S2'))
                    strbuff += self.ecps.coords2str(coords=14, charges=8, atoms=None, indent=2*' ', inserts=[(4,ecpnames)])

            strbuff += 'end\n'

        # basis set
        needbasis = True
        needend   = True
        basislines = self.basis.splitlines()
        if len(basislines) <= 1:
            needend = False
        for i in basislines:
            if i.lower().strip().startswith('basis'):
                needbasis = False
            if i.lower().strip().startswith('end'):
                needend   = False
        if needbasis:
            strbuff += '\nbasis ' + self.basis
        else:
            strbuff += '\n' + self.basis
        if needend:
            strbuff += '\nend\n' 

        strbuff += '\n'

        # ECP file or text
        # AJL, Aug2017: I suggest changing this for the functionality now included in NWChem, but I won't implement it
        # as I've no way of testing the functionality. The code to insert instead of the below if statement would be:
        # YL 08/08/2017: done
        strbuff += fileutils.file2text(self.ecp, prefix=parallel.prefix) + '\n'

        # direct scf
        if self.direct:
            direct = 'direct'
        else:
            direct = ''

        # SCF type
        # use rhf for RHF/RKS/ROHF 
        # NB: ROKS is not available in GAMESS-UK
        # TODO other ab initio SCFTYPEs
        gms_scftype = 'rhf'
        if self.scftype == 'uhf' or self.scftype == 'uks':
            gms_scftype = 'uhf'
        strbuff += 'scftype %s %s\n'%(direct, gms_scftype)

        # runtype
        if self.analysis.lower() == 'nbo':
            self._gradients = False
            self.restart = False

        if self._gradients:
            strbuff += 'runtype gradient\n'
            if self.frag.nbqs and self._gradients == 'shells':
                strbuff += 'bqfo\n'

        # population analysis
        if self.pop_analysis:
            strbuff += 'mullik atom\n'

        if self._graphics and self.frag.ngridpoints:
            strbuff += 'punch high grid 151\n'
            strbuff += 'runtype analyse\n'
            strbuff += 'graphics\n'
            strbuff += 'gdef\n'
            strbuff += 'type cards %d\n'%self.frag.ngridpoints
            strbuff += self.frag.grid.coords2str(coords=14, charges=None, atoms=None, indent=2*' ')
            strbuff += 'section 150\n'
            strbuff += 'calc\n'
            strbuff += 'type pote\n'
            strbuff += 'section 151\n'
            self.vectors = '1'
        
        # method and functional
        if self.method == 'dft':
            strbuff += 'dft ' + self.functional + '\n'
            selectcases = { 'low'      :'dft quad low\n',
                            'medium'   :'dft quad medium\n',
                            'high'     :'dft quad high\n',
                            'veryhigh' :'dft quad veryhigh\n',
                            'ultrahigh':'dft quad ultrahigh\n',
                          }
            strbuff += selectcases.get(self.grid, '')
#            if self.quadrature:
#                strbuff += 'dft quadrature %s\n'%self.quadrature

        # max cycles: default is 50, see http://www.cfs.dl.ac.uk/docs/html/part4/node6.html
        strbuff += 'maxcyc ' + str(self.maxiter) + '\n'

        # threshold
        strbuff += 'thresh ' + str(self.scf_threshold) + '\n'

        # YL 05/05/2020: level shift
        if self.level_shift:
            strbuff += 'level {} {} {}\n'.format(self.level_e1, self.level_cycles, self.level_e2)

        # newscf
        if self.newscf:
            strbuff += fileutils.file2text(self.newscf, prefix=parallel.prefix) + '\n'

        # see: http://www.cfs.dl.ac.uk/docs/html/part4/node13.html
        selectcases = { 'rhf'   :'1',   # closed-shell SCF
                        'uhf'   :'1 2', # UHF SCF
                        'rohf'  :'4 5', # open-shell RHF
                        'gvb'   :'4 5', # GVB
                        'casscf':'6 7', # CASSCF 
                      }

        self.enter = selectcases[self.scftype]

        # restart II
        if self.restart:
            # vectors and enter
            self.vectors = selectcases[self.scftype]

            if self.ed4:
                self.vectors = 'getq ed4 1 2'

        # once self.restart is set True and False again, self.vectors must be reset to default value
        else:
            self.vectors = self._kwargs['vectors']

        # MO sections to update
        if not self.guess:
            strbuff += 'vectors ' + self.vectors + '\n' 
            strbuff += 'enter ' + self.enter + '\n'
        else:
            # YL 20/02/2021: it was a bug the guess method was always applied even when it was a restarted job
            if not self.restart:
                # guess file or string
                strbuff += fileutils.file2text(self.guess, prefix=parallel.prefix) + '\n'

        # runtype analyse
        if self.analysis.lower() == 'nbo':
            strbuff += 'runtype analyse\n'
            strbuff += 'scftype direct\n'
            strbuff += 'nbo\n'
            strbuff += 'vectors 1\n'
            strbuff += 'enter 1\n'
            strbuff += '$nbo $end\n'

        return strbuff


    def parseAdditionalOutput(self):
        '''Overloaded: Parse punch file'''

        from os    import remove, getcwd
        from ...io import file

        # energy
        self._result.energy = file.getArrayFromFile(self.punch, 'punch', self._result.energy.shape, regex='total_energy')

        # gradients
        if self._gradients:
            self._result.gradients = file.getArrayFromFile(self.punch, 'punch', self._result.gradients.shape, regex='gradients')

        # bq gradients
        if self._gradients and self.frag.nbqs:
            # array shape with ECPs
            if self.ecps:
                shape = (self.frag.bqs.natoms+self.ecps.natoms,3)
                self.frag.bqs.gradients = file.getArrayFromFile(self.punch, 'punch', shape, regex='bq_gradients')[:self.frag.bqs.natoms,:]
                if self.separate_ecps:
                    self.ecps.gradients = file.getArrayFromFile(self.punch, 'punch', shape, regex='bq_gradients')[self.frag.bqs.natoms:,:]
            else:
                self.frag.bqs.gradients = file.getArrayFromFile(self.punch, 'punch', self.frag.bqs.gradients.shape, regex='bq_gradients')

        if self._graphics and self.frag.ngridpoints:
            from ...data.unitconvert import kcal_mol2Ha
            self.frag.grid.potentials = file.getArrayFromFile(self.punch, 'punch', self.frag.grid.potentials.shape, regex='grid_data')*kcal_mol2Ha


    @miscutils.memusage
    def runLib(self, options=None, **kwargs):
        '''Run GAMESS-UK as a shared library (called by run.runLib())'''

        from os       import path, environ, remove
        from ctypes   import c_int, byref, CDLL
        from chemsh   import CHEMSH_LIB_DIR

        # 'ftn058' for library call and '_gamess-uk.pun' for system call
        self.punch = 'ftn058'

        # prevent reading result from old results
        fileutils.archiveFile(self.punch)

        # GAMESS-UK uses environment variable `ed3` when command line is not available
        # being cross-platform may be a problem but GAMESS-UK probably does not work in Windows anyway
        environ['ed2'] = self.ed2
        environ['ed3'] = self.ed3
        environ['ed7'] = self.ed7

        # load GAMESS-UK library
        libgamess_uk = modutils.getLinkedLib(self.__class__.__name__, chemsh.CHEMSH_LIB_DIR)

        fragRef = py2f(self.frag)

        if libgamess_uk.interface_gamess_uk(byref(fragRef)):
            print(" >>> GAMESS-UK: libgamess-uk.interface_gamess_uk failed.")
            return 1
        else:
            return 0

#        # delete integral file
#        try:
#            remove('_gamess-uk.ed2000')
#            remove('_gamess-uk.ed2001')
#            remove('_gamess-uk.ed2002')
#            remove('_gamess-uk.ed2003')
#            remove('_gamess-uk.ed2004')
#            remove('_gamess-uk.ed2005')
#            remove('_gamess-uk.ed2006')
#            remove('_gamess-uk.ed2007')
#            remove('_gamess-uk.ed2008')
#            remove('_gamess-uk.ed2009')
#        except:
#            pass

