#  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/>.

import chemsh
from sys              import stdout
from ..mm             import _MM
from ...base          import constants, errors
from ...data          import atomicdata
from ...io            import file, files
from ...dl_py2f.utils import fileutils, miscutils, modutils

class DL_POLY(_MM):
    '''Interface to DL_POLY 4'''

    _kwargs = {
                'berendsen'       : (0.5,1.5),
                'delr'            : 0.0,
                'densvar'         : 0,
                'dl_field_charges': True,
                'dl_field_types'  : True,
                'equilibration'   : 10,
                'ewald'           : 1.0E-6,
                'index'           : True,
                'input'           :('_dl_poly.inp','CONFIG','FIELD'),
                'output'          :'_dl_poly.out',
                'potential'       : False,
                'print'           : 2,
                'rcut'            :-1.0,          # use negative default so rcut is determined by self.frag.cell.cutoff
                'rpad'            : 0.0,
                'runargs'         :'',
                'rvdw'            : 0.0,
                'scale'           : 2,
                'steps'           : 20,
                'stack'           : 50,
                'stats'           : 2,
                'restart'         :'scale',
                'timestep'        : 0.000,
#                'unit'            :'kcal',
              }

    _internals = {
                   '_amber'            : False,
# not in use
#                   '_dl_field_deffered': False,
                   '_exclude_qm'       : False,
                   '_qmmm_embedding'   : None,
                   '_syscmd'           :'DLPOLY.Z',
                 }

    # YL 06/12/2023: processing `self.ff` may require to write out a new FIELD file (input[1]) so that `input` must be defined prior to `ff`
    _priorities = 'frag', 'input', 'ff'

    def getInpBuff(self, input=input, **kwargs):
        ''''''

        from os          import path
        from numpy       import arange, full, __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
        from ...data     import unitconvert
        from ...io       import files
        from ...parallel import parallel
        from .           import tools

#        # if using DL_FIELD switch on CONFIG's index
#        if hasattr(self.ff, 'run'):
#            self.index = True

        # formerly "CONTROL" file, now "_dl_poly.inp"
        if input == self.input[0]:

            strbuff = 'ChemShell DL_POLY\n'

            if self.densvar:
                strbuff += 'densvar %d %s\n'%(self.densvar, '%')
    
            if not self.index:
                strbuff += 'no index\n'

            strbuff += 'coul\n'

            strbuff += 'temperature ' + str(self.temperature) + '\n'
#            strbuff += 'pressure ' + str(self.pressure) + '\n'
#            strbuff += 'ensemble ' + self.ensemble + ' ber ' + str(self.berendsen[0]) + ' ' + str(self.berendsen[1]) + '\n'
    
            strbuff += 'timestep ' + str(self.timestep) + '\n'
    
            strbuff += 'restart %s\n'%self.restart
    
#            strbuff += 'steps ' + str(self.steps) + '\n'
#            strbuff += 'equilibration ' + str(self.equilibration) + '\n'
    
#            strbuff += 'scale ' + str(self.scale) + '\n'
#            strbuff += 'rcut %f\n'%self.rcut
            strbuff += 'cut %f\n'%self.rcut
            strbuff += 'vdw direct\n'
# 21/06/2021: equals to rcut if not specified separately
#            strbuff += 'rvdw 150.00\n'
            strbuff += 'delr %f\n'%self.delr
#            strbuff += 'ewald precision ' + str(self.ewald) + '\n'
    
            strbuff += 'print ' + str(self.print) + '\n'
    
#            strbuff += 'stack ' + str(self.stack) + '\n'
            strbuff += 'stats ' + str(self.stats) + '\n'

            # YL 21/09/2020: time elapsed of DL_POLY 4 is obtained from MPI, which is set up
            #                globally in the context of ChemShell. so the `job time` value
            #                must be long enough to allow long run of a ChemShell job
            #                this is 365*24*60*60 s
            strbuff += 'job time   {:<f}\n'.format(31536000.0)
            # YL 21/09/2020: leave a bit time for writing out before closure
            strbuff += 'close time {:<f}\n'.format(30.0)
    
            strbuff += 'finish\n'

        # CONFIG file
        elif input == self.input[1]:

            # self.ff is a filename
            if type(self.ff) is str and path.isfile(self.ff):
                # AMBER format prmtop file: requires AMBER atom types for atom names
                try:
                    # receive the forcefield data with a dictionary
                    dbuff = file.parse(self.ff, fmt='amber_prmtop')
                    for key, val in dbuff.items():
                        if type(val) is dict and 'atom' in val:
                            self.frag.types = dbuff[key]['atom']['data']['name']
                    self._amber = True
                except:
                    pass

            # record 1: title line
            strbuff = 'ChemShell DL_POLY\n'

            # record 2: LEVCFG and IMCON
            try:
                strbuff += '0 %d\n'%self.frag.cell.imcon
            except:
                strbuff += '0 0\n'

            # record 3, 4, and 5: x/y/z of vectors a, b, and c
            if self.frag.cell.natoms:
                strbuff += '%20.6f %20.6f %20.6f\n'%(self.frag.cell.vectors[0,0]*unitconvert.Bohr2Angstrom,
                                                     self.frag.cell.vectors[1,0]*unitconvert.Bohr2Angstrom,
                                                     self.frag.cell.vectors[2,0]*unitconvert.Bohr2Angstrom)
                strbuff += '%20.6f %20.6f %20.6f\n'%(self.frag.cell.vectors[0,1]*unitconvert.Bohr2Angstrom,
                                                     self.frag.cell.vectors[1,1]*unitconvert.Bohr2Angstrom,
                                                     self.frag.cell.vectors[2,1]*unitconvert.Bohr2Angstrom)
                strbuff += '%20.6f %20.6f %20.6f\n'%(self.frag.cell.vectors[0,2]*unitconvert.Bohr2Angstrom,
                                                     self.frag.cell.vectors[1,2]*unitconvert.Bohr2Angstrom,
                                                     self.frag.cell.vectors[2,2]*unitconvert.Bohr2Angstrom)

            # core atoms
            if self.index:
                indices = arange(self.frag.natoms)+1
                if self._amber or self.atom_names.lower().startswith('typ'):
                    # YL 05/11/2021: because self.frag will be refreshed every cycle of geometry optimisatin we must update the types whenever possible
                    try:
                        self.frag.types = self._dict_FIELD[0]['atom']['data']['name']
                    except:
                        pass
                    strbuff += self.frag.coords2str(atoms='types', coords=6, charges=None, delimiter=' ', unit='angstrom', inserts=[(1, indices), (2, '\n')], align='left')
                else:
                    strbuff += self.frag.coords2str(atoms='names', coords=6, charges=None, delimiter=' ', unit='angstrom', inserts=[(1, indices), (2, '\n')])
            else:
                if self._amber or self.atom_names.lower().startswith('typ'):
                    # YL 05/11/2021: because self.frag will be refreshed every cycle of geometry optimisatin we must update the types whenever possible
                    try:
                        self.frag.types = self._dict_FIELD[0]['atom']['data']['name']
                    except:
                        pass
                    strbuff += self.frag.coords2str(indent='', atoms='types', coords=6, charges=None, delimiter=' ', unit='angstrom', inserts=[(1, '\n')])
                else:
                    strbuff += self.frag.coords2str(atoms='names', coords=6, charges=None, delimiter=' ', unit='angstrom', inserts=[(1, '\n')])

            # shell atoms
            if self.frag.nshells:
                shellnames = defchararray.add(self.frag.shells.names, full((self.frag.nshells,), 's\n', dtype='S8'))
                strbuff += self.frag.shells.coords2str(atoms=None, coords=6, charges=None, delimiter=' ', unit='angstrom', inserts=[(0, '  '), (1,shellnames)])

        # FIELD file
        elif input == self.input[2]:

            # receiver has to be a mutable type, such as bytearray
            strbuff = bytearray(b'')

            # TODO: parameterise=self.parameterise, 
            # DL_POLY allows DL_FIELD instance as forcefield input
            # YL 16/02/2021: we must run DL_FIELD at every cycle of geometry optimisaton or MD simulation
            #                because if a non-template-based forcefield, such as CVFF and PCFF, is used
            #                ChemShell will have no idea how to generate the CONFIG file
            if hasattr(self.ff, 'run'):
                self.runDL_FIELD(output=input)
                # enable indices in CONFIG
                self.index = True
                # copy the content from FIELD
                if self.ff.field_by_chemsh:
                    # write out FIELD using the saved data
                    files.dlp_field.write(self.input[2], self._dict_FIELD)
                    strbuff = fileutils.file2text(self.input[2], prefix=parallel.prefix)
                else:
                    strbuff = self.ff.fftext

            # `self.ff` is file
            else:
                # remove QM/MM redundant terms (typically for AMBER format)
                if len(self.qmatoms) and self._exclude_qm:
                    exclude = getattr(self, 'qmmm_pairs_by_junctions', [])
                    tools.excludeQM(self._dict_FIELD, self.qmatoms, exclude)

                # write out a new FIELD file iff data are previously saved in self._dict_FIELD
                if self._dict_FIELD:
                    files.dlp_field.write(self.input[2], self._dict_FIELD)
                    # to keep the standardised code we read the FIELD file into strbuff
                    file.parse(path.join(parallel.prefix,self.input[2]), fmt='ff', receiver=strbuff)
                    # back to str (no leading and trailing blank lines allowed by DL_POLY)
                    strbuff = strbuff.decode().lstrip().rstrip()
                else:
                    strbuff = fileutils.file2text(self.ff, prefix=parallel.prefix) 

        return strbuff

    def parseMainOutput(self):
        '''Overload mm.MM.parseMainOutput(): only parse gulp.energy and gulp.gradient'''


    @miscutils.memusage
    def runLib(self, options=None, task='qmmm', **kwargs):
        '''Run DL_POLY 4 as a shared library (called by run.runLib())'''

        import chemsh
        from sys         import stdout
        from _ctypes     import dlclose
        from ctypes      import byref, RTLD_LOCAL, DEFAULT_MODE
        from ...base     import errors
        from ...dl_py2f  import py2f
        from ...data     import unitconvert
        from ...io       import format
        from ...parallel import parallel

        # YL 22/09/2020: check nprocs
        if parallel.nprocs_mm not in [1,2,4,8]:
            format.printwrap(f'>>> ChemShell/DL_POLY ERROR: The number of MPI processes for DL_POLY (currently {parallel.nprocs_mm}, via command line argument -npmm) has to be 1, 2, 4, or 8\n',
                             indent='', initial_indent=' ', subsequent_indent=' '*30, prepend_blanc_line=True, replace_whitespace=True)
            print()
            stdout.flush()
            exit(999)
        # YL 22/11/2021: nprocs cannot be greater than natoms, otherwise DL_POLY's parallel I/O fails
        if parallel.nprocs_mm > self.frag.natoms:
            format.printwrap(f'>>> ChemShell/DL_POLY ERROR: The number of MPI processes for DL_POLY (currently {parallel.nprocs_mm}, via command line argument -npmm) must not be greater than the number of atoms ({self.frag.natoms})\n',
                             indent='', initial_indent=' ', subsequent_indent=' '*30, prepend_blanc_line=True, replace_whitespace=True)
            print()
            stdout.flush()
            exit(999)

        # load DL_POLY library
        libdl_poly = modutils.getLinkedLib(self.__class__.__name__, chemsh.CHEMSH_LIB_DIR, mode=DEFAULT_MODE)

        # for pure MM calculation
        if not self._qmmm_embedding:
            task = 'mm'

        fragRef    = py2f(self.frag   , byref=True, debug=chemsh.CHEMSH_DEBUG)
        resultRef  = py2f(self._result, byref=True, debug=chemsh.CHEMSH_DEBUG)
        optionsRef = py2f(self.options, byref=True, debug=chemsh.CHEMSH_DEBUG)
        self.lib.interface_dl_poly(fragRef,
                                   resultRef,
                                   optionsRef)

        # the internal unit of DL_POLY is 10 J/mol!
        # DL_POLY uses 1 cal = 4.184 J
        # and in TCS/DL_POLY 1 Ha = 627.524955413 kcal/mol (!)
#        self._result.energy *= unitconvert.kcal_mol2Ha*10.0/4.184/1000.0
        self._result.energy *= 1.0/627.524955413*10.0/4.184/1000.0

        # beware: DL_POLY 4 outputs "forces" rather than "gradients"
        self._result.gradients *= -1.0/627.524955413*10.0/4.184/1000.0 * unitconvert.Bohr2Angstrom

        stdout.flush()


    def runDL_FIELD(self, output='FIELD'):
        '''Run DL_FIELD'''

        from sys      import stdout
        from numpy    import char, count_nonzero
        from .        import tools

        if self.ff._rerun:
            # back up ChemShell's CONFIG before DL_FIELD overwrites it
            fileutils.archiveFile(self.input[1], suffix='.OLD')
            print("\n ChemShell/DL_POLY: Generating forcefield from {} driver...\n".format(self.ff.__class__.__name__))
            # DL_FIELD generates a new CONFIG file which overwrites the ChemShell one, except for cycle >= 1 when run_only_once=True
            self.ff.run(frag     = self.frag,
                        qmatoms  = self.qmatoms,
                        nqmatoms = len(self.qmatoms),
                        # QM/MM only
                        exclude  = getattr(self, 'qmmm_pairs_by_junctions', []))
            # YL 03/11/2021: we switch it off after DL_FIELD has been run successfully for the first time (iff run_only_once=True)
            if self.ff.run_only_once:
                self.ff._rerun = False
                # because after the first cycle we will write out CONFIG with ChemShell only 
                # we need to print FF types rather than names
                self.atom_names = 'type'
            # YL 03/11/2021: save the data of FIELD by parsing it
            dbuff = files.dlp_field.parse(output)
            # YL 03/11/2021: unpack the data containing repeated molecules and molecule data containing repeated atoms
            self._dict_FIELD = tools.unravelData(dbuff)

            # YL 07/11/2021: an optional sanity check of atom types
            natoms_FIELD = self._dict_FIELD[0]['atom']['data']['name'].size
            natoms_frag  = self.frag.natoms
            # update atom types using the FIELD data
            if natoms_FIELD == natoms_frag:

                # types
                masks_consistent_types = char.equal(char.strip(char.upper(self._dict_FIELD[0]['atom']['data']['name'])),
                                                    char.strip(char.upper(char.decode(self.frag.types)))
                                                   )
                if any(~masks_consistent_types):
                    if chemsh.CHEMSH_DEBUG:
                        print()
                        for i, t in enumerate(self.frag.types[~masks_consistent_types]):
                            print(" >>> WARNING: Atom {}\'s type \"{}\" inconsistent with \"{}\" by FIELD".format(i,
                                                                                                                self.frag.types[i].decode(),
                                                                                                                self._dict_FIELD[0]['atom']['data']['name'][i]))
                    # YL 03/06/2022: allow to override the existing types by DL_FIELD
                    if self.dl_field_types:
                        self.frag.types[:] = self._dict_FIELD[0]['atom']['data']['name'][:]
                        print('\n {} inconsistent atom types overridden according to DL_FIELD'.format((~masks_consistent_types).sum()))
                    # YL 24/07/2022: the FIELD data set must be updated accordingly
                    else:
                        masks_non_blank = char.strip(self.frag.types) != b''
                        self._dict_FIELD[0]['atom']['data']['name'][:][masks_non_blank] = self.frag.types[:][masks_non_blank]

                # YL 15/06/2022: also charges
                masks_consistent_charges = self._dict_FIELD[0]['atom']['data']['charge'] == self.frag.charges
                if any(~masks_consistent_charges):
                    if chemsh.CHEMSH_DEBUG > 1:
                        print()
                        for i, t in enumerate(self.frag.charges[~masks_consistent_charges]):
                            print(" >>> WARNING: Atom {}\'s charge \"\" inconsistent with \"{}\" by FIELD".format(i,
                                                                                                                  self.frag.charges[i],
                                                                                                                  self._dict_FIELD[0]['atom']['data']['charge'][i]))
                    # YL 03/06/2022: allow to override the existing charges by DL_FIELD
                    if self.dl_field_charges:
                        self.frag.charges[:] = self._dict_FIELD[0]['atom']['data']['charge'][:]
                        print('\n {} atom charges overridden according to DL_FIELD\n'.format((~masks_consistent_charges).sum()))
                    # YL 24/07/2022: the FIELD data set must be updated accordingly
                    else:
                        if not count_nonzero(self.frag.charges):
                            format.printwrap('WARNING: the DL_POLY option dl_field_charges is switched off but all the fragment\'s charges are zeros')
                        self._dict_FIELD[0]['atom']['data']['charge'][:] = self.frag.charges[:]
                masks = (self.frag.names==b'HD21')+(self.frag.names==b'HD22')

            else:
                print('\n >>> WARNING: DL_FIELD does not agree on the number of atoms!')

            # remove QM/MM redundant terms
            if self.ff.field_by_chemsh and len(self.qmatoms):
                exclude  = getattr(self, 'qmmm_pairs_by_junctions', [])
                tools.excludeQM(self._dict_FIELD, self.qmatoms, exclude)

            # archive DL_FIELD's FIELD before writing out a new one
            fileutils.archiveFile(output, dir=self.ff.directory, suffix='')


# PARAMETERS #

    @property
    def rcut(self):
        '''Cutoff'''

        from ...data import unitconvert

        # user-defined
        if self._rcut >= 0.0:
            return self._rcut
        # default
        else:
            return self.frag.cell.cutoff*unitconvert.Bohr2Angstrom


    @rcut.setter
    def rcut(self, val):
        '''Setter of rcut'''

        # self.frag.cell.cutoff is not yet ready when the current setter is called by initialisation
        self._rcut = val


    @property
    def ff(self):
        '''Force field definition (plain text or filename or DL_FIELD instance)'''

        from .. import dl_field

        return self._ff


    @ff.setter
    def ff(self, val):
        '''Setter of ff'''

        from os          import path
        from .           import tools
        from ...parallel import parallel

        self._ff = val

        # DL_FIELD instance
        if hasattr(val, 'run'):
            return

        # is a filename
        # YL 06/12/2023: added consideration of taskfarmed calculation which was missing before
        if path.isfile(path.join(parallel.prefix, val)):
            if val.strip() == 'FIELD':
                print('\n ERROR: "FIELD" cannot be used as the name of a given forcefield file as it will be overwritten!')
                exit(112)
            filename = path.join(parallel.prefix, val)
            # case 1: AMBER format prmtop file
            try:
                # receive the forcefield data with a dictionary
                self._dict_FIELD = file.parse(filename, fmt='amber_prmtop')
                self._exclude_qm = True

            # case 2: DL_POLY FIELD format
            except:
                # DL_POLY Classic format
                try:
                    # read in the file's content
                    dbuff = files.dlp_field.parse(filename)
                    # YL 05/12/2023: unfold data folded by molecules
                    self._dict_FIELD = tools.unravelData(dbuff)
                    # YL 05/12/2023: it's impossible to set the values to self.frag.types here because self.frag will be redefined by QMMM
                # later DL_POLY format
                # YL REVIEWED 06/12/2023: our parser still can't read in this type of FIELD (due to different number of columns)
                except:
                    pass

        # string containing the contents of FIELD
        else:
            # write it out to a file so that at "runtime" of DL_POLY we can treat every scenarios in a uniform way
            with open('_dl_poly.ff', 'w') as fp:
                fp.write(val)
            self._ff = '_dl_poly.ff'
        # YL 24/11/2023: we run DL_FIELD immediately as pre-processing (especially for patching residues)
#        if hasattr(val, 'run'):
# YL 25/11/2023: not in use
#            # run now
#            if self._hasFrag:
#                self.runDL_FIELD(output=self.input[2])
#                self._dl_field_deferred = False
#            # deferred if fragment hasn't been ready
#            else:
#                self._dl_field_deferred = True


    @property
    def _dict_FIELD(self):
        '''A dictionary holding the FIELD data'''

        try:
            return self.__dict_FIELD
        except:
            self.__dict_FIELD = {}
            return self.__dict_FIELD


    @_dict_FIELD.setter
    def _dict_FIELD(self, val):
        '''Setter of _dict_FIELD'''

        self.__dict_FIELD = val


# YL 25/11/2023: not in use at the moment; created intendedly for triggering DL_FIELD as soon as DL_POLY is defined (with `frag`)
#    @property
#    def frag(self):
#        '''Fragment'''
#
#        return self._frag
#
#
#    # YL 24/11/2023: overriding the default one under <class '_MM'>
#    @frag.setter
#    def frag(self, val):
#        '''Setter of fragment'''
#
#        self._frag = val
#        if self._hasFrag:
#            # it could be possible that DL_FIELD hasn't been employed
#            if self._dl_field_deferred:
#                self.runDL_FIELD(output=self.input[2])
#                self._dl_field_deferred = False

