#  Copyright (C) 2019 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__ = 'You Lu <you.lu@stfc.ukri.org>'

from ..mm             import _MM
from ...data          import unitconvert
from ...dl_py2f.utils import fileutils, iterutils

class NAMD(_MM):
    ''''''

    _kwargs = { 
                'binary'          : True,
                'constant_area'   : False,
                'constraints'     :'',
                'constraints_ref' :'',
                'coor'            :'_namd.coor',
                'cutoff'          : 12.0*unitconvert.Angstrom2Bohr,
                'logger'          :'_namd.err',
                'exclude'         :'scaled1-4',
                'ff_dir'          :'',
                'fixed_atoms'     :'',
                'fixed_atoms_forces': True,
                'flexible_cell'   : False,
                'freq_nonbonded'  : 1,
                'freq_full_elect' : 1,
                'group_pressure'  : False,
                'input'           :'_namd.inp',
                'margin'          : 2.5*unitconvert.Angstrom2Bohr,
                'memory_per_proc' : 250,             # memory per process in MiB for NAMD on HPC machines (typically used for oversubscribing a scheduler, see self.runargs below)
                'merge_cross'     : True,
                'merge_restarted' : True,
                'nsteps_per_cycle': 1,
                'output'          :'_namd.out',
                'pairlist_dist'   : 14.0*unitconvert.Angstrom2Bohr,
                'par'             :'',
                'patch'           : True,             # needed by PSFGEN only
                'pdb'             :'',
                'pme'             : True,
                'pme_grid_sizes'  :[0,0,0],
                'prefix'          :'_namd',
                'prefix_restart'  :'_namd',
                'psf'             :'',
                'psfgen_options'  :{},
                'scaling14'       : 1.0,
                'seed'            : 12345,
                'switching'       : True,
                'switch_dist'     : 10.0*unitconvert.Angstrom2Bohr,
                'top'             :[],
                'vel'             :'_namd.vel',
                'wrap_all'        : True,
                'wrap_water'      : True,
                'xsc'             :'_namd.xsc',
                'xst'             :'_namd.xst',
              }

    _internals = { '_additional_prm'    :[],
                   '_chk'               :'_namd_checkpoint',
                   '_dcd'               :'_namd.dcd',
                   '_default_pdb'       : False,
                   '_default_psf'       : False,
                   '_exec_run_shell'    : False,
                   '_MDOptions'         : None,
                   '_par'               :'_namd.prm',        # NB: NAMD.par is not an property() and NAMD._par isn't its underlying value
                   '_runargs_print_hint': True,
                 }



    def getInpBuff(self, MDOptions=None, **kwargs):
        ''''''

        from numpy        import any, char, unique
        from os.path      import abspath, isdir, isfile, join, relpath
        from ...data      import forcefield, references, unitconvert
        from .            import psfgen
        from ...io.banner import printModuleBanner
        from ...io        import format

        ffObj = forcefield.getFFObj(MDOptions.ff)

        # back up previous exsiting folder
        self.archive(dest=self.prefix)

        selectcases = { int  :'d',
                        float:'f',
                        str  :'s',
                        bool :'r',
                      }

        strbuff = '\n# NAMD configuration for ChemShell\n'

        bools = { True : 'on',
                  False: 'off',
                }

        yorn = { True : 'on',
                 False: 'off',
               }


        ff_dir = forcefield.getFFDir(MDOptions.ff)
        par_abspaths = psfgen.getAbsPathsList(self.par, ff_dir)
        par_abspaths += self._additional_prm
        # YL 01/05/2024: there are some parameters (such as vdW for OC2D1, OC2D3, NN1...) missing
        #                in par_all36_cgenff.prm but defined in other parameter files
        if join(ff_dir, 'par_all36_cgenff.prm') in par_abspaths:
            par_abspaths += [ join(ff_dir, 'par_all36_carb.prm') ]
            par_abspaths += [ join(ff_dir, 'par_all36_lipid.prm') ]
            par_abspaths += [ join(ff_dir, 'par_all36_na.prm') ]
        par_abspaths = unique(par_abspaths).tolist()

        if issubclass(type(self.par), str):
            pass

        # YL 03/11/2021: we now have a great feature function NAMD.mergeParameters() which we can consider to use here
        if len(par_abspaths) > 1:
            print("\n Concatenating parameter files:")
        fileutils.archiveFile(self._par)
        # concatenate PAR files
        types_lower = char.strip(char.lower(char.decode(self.frag.types)))
        def _commentOut(_strbuff, _types):
            '''Determine if a line needs to be commented out'''
            # YL 26/01/2021: only comment out if the type doesn't appear in the fragment
            # YL 26/01/2021: added "OC" and "OS" for DNA
            # TODO: protein-DNA complex
            _types = [ 'oc', 'os', 'on3', 'ocl', 'obl', 'o2l', 'oc2d2', 'oc2dp', 'og2d2', 'og2p1' ]
            for _type in _types:
                if _type in _strbuff.split() and _type not in types_lower:
                    return True
            return False
        with open(self._par, 'w') as fp0:
            for par_abspath in par_abspaths:
                if len(par_abspaths) > 1:
                    print(" {}".format(abspath(par_abspath)))
                with open(par_abspath, 'r') as fp:
                    lines = fp.readlines()
                    # comment out lines that NAMD doesn't recognise
                    for l in lines:
                        ls = l.strip().lower()
                        if ls.startswith('set') \
                        or ls.startswith('if')  \
                        or ls.startswith('wrnlev') \
                        or ls.startswith('bomlev') \
                        or ls.startswith('return'):
                            fp0.write('!'+l)
                        elif _commentOut(ls, types_lower):
                            fp0.write('!'+l)
                        else:
                            fp0.write(l)

        # par
        if not self.ff_dir:
            if isdir(join(ffObj.default_location)):
                self.ff_dir = ffObj.default_location
                print("\n Using default CHARMM topology & parameter files in (version: {}):\n {}".format(ffObj.default_version, self.ff_dir))
        else:
            self.ff_dir = abspath(self.ff_dir)
            if isdir(self.ff_dir):
                print("\n Using specified topology & parameter files in:\n {}".format(self.ff_dir))
            else:
                print("\n >>> ERROR: Directory specified by `ff_dir` does NOT exist:\n {}".format(self.ff_dir))
                exit(999)

#        if isfile(self.par):
#            par_path = relpath(self.par)
#        elif isfile(join(self.ff_dir, self.par)):
#            par_path = relpath(join(self.ff_dir, self.par))
#        else:
#            print("\n >>> ERROR: cannot find parameter file {}".format(par_path))
#            exit(999)

        # use psfgen outputs if not specified otherwise by user
        if not self.psf:
            self.psf = self._psf
        if not self.pdb:
            self.pdb = self._pdb

        strbuff += '\n# input/output\n'
        strbuff += '{:24}{}\n'.format('coordinates', self.pdb)
        strbuff += '{:24}{}\n'.format('structure'  , self.psf)
        strbuff += '{:24}{}\n'.format('parameters' , self._par)
        strbuff += '{:24}{}\n'.format('outputname' , self.prefix)

        # restart from
        if self.restart:
            strbuff += '\n# restart from existing files\n'
            strbuff += '{:24}{}\n'.format('bincoordinates', self.prefix_restart+'.coor')
            strbuff += '{:24}{}\n'.format('binvelocities' , self.prefix_restart+'.vel')
            strbuff += '{:24}{}\n'.format('extendedsystem', self.prefix_restart+'.xsc')

        # binary output (default: yes)
        if not self.binary:
            strbuff += '{:24}{}\n'.format('binaryoutput' , 'off')
            strbuff += '{:24}{}\n'.format('binaryrestart', 'off')

        # MD
        if MDOptions:
            # YL 27/07/2023: save a reference for later use (self.parseMainOutput())
            self._MDOptions = MDOptions
            strbuff += '\n# molecular dynamics\n'
            strbuff += '{:24}{}\n'.format('timestep'     , MDOptions.timestep)
            strbuff += '{:24}{}\n'.format('numsteps'     , MDOptions.nsteps)
            strbuff += '{:24}{}\n'.format('stepspercycle', self.nsteps_per_cycle)
            strbuff += '{:24}{}\n'.format('firsttimestep', 0)
            if not self.restart:
                strbuff += '{:24}{}\n'.format('temperature'  , MDOptions.temperature)
            strbuff += '{:24}{}\n'.format('rigidbonds'   , MDOptions.rigid)
        
            strbuff += '{:24}{}\n'.format('restartname'  , self._chk)
            strbuff += '{:24}{}\n'.format('restartfreq'  , int(MDOptions.freq_chk))
            strbuff += '{:24}{}\n'.format('dcdfile'      , MDOptions.dcd)
            strbuff += '{:24}{}\n'.format('dcdfreq'      , int(MDOptions.freq_dcd))
            strbuff += '{:24}{}\n'.format('xstfreq'      , int(MDOptions.freq_xst))

            if MDOptions.freq_out_energy%self.freq_full_elect:
                format.printwrap("NAMD option fullelectfrequency set to 1 because outputenergies must be divisible by it")
                self.freq_full_elect = 1
            strbuff += '{:24}{}\n'.format('fullelectfrequency', int(self.freq_full_elect))
            strbuff += '{:24}{}\n'.format('nonbondedfreq'     , int(self.freq_nonbonded))

            strbuff += '{:24}{}\n'.format('outputenergies', int(MDOptions.freq_out_energy))
            strbuff += '{:24}{}\n'.format('outputpressure', int(MDOptions.freq_out_pressure))

            if MDOptions.seed:
                strbuff += '{:24}{:d}\n'.format('seed', MDOptions.seed)
            else:
                istrbuff += '{:24}{:d}\n'.format('seed', 'off')

            # PBCs
            # YL 01/03/2021: we don't write cell vectors if it's a restarted job
            if MDOptions.boundary and not self.restart:
                strbuff += '\n# periodic boundary conditions\n'
                if MDOptions.boundary.lower() in [ 'periodic' ]:
                    vectors = self.frag._xyz_spanning*unitconvert.Bohr2Angstrom
                    strbuff += '{:24}{:12.6f}{:12.6f}{:12.6f}\n'.format('cellbasisvector1', vectors[0],       0.0 ,         0.0)
                    strbuff += '{:24}{:12.6f}{:12.6f}{:12.6f}\n'.format('cellbasisvector2',        0.0, vectors[1],         0.0)
                    strbuff += '{:24}{:12.6f}{:12.6f}{:12.6f}\n'.format('cellbasisvector3',        0.0,       0.0 ,  vectors[2])
                    strbuff += '{:24}{:12.6f}{:12.6f}{:12.6f}\n'.format('cellorigin'      , *self.frag.centroid*unitconvert.Bohr2Angstrom)
                # spheric
                elif MDOptions.boundary.lower() in [ 'spheric', 'spherical' ]:
                     strbuff += '\n{:24}{}\n'      .format('sphericalbc', 'on')
                     strbuff += '{:24}{:12.6f}{:12.6f}{:12.6f}\n'.format('sphericalbccenter', *self.frag.centroid*unitconvert.Bohr2Angstrom)
                     strbuff += '{:24}{:12.6f}\n'.format('sphericalbcr1'  , MDOptions.spheric_r0*unitconvert.Bohr2Angstrom)
                     strbuff += '{:24}{:12.6f}\n'.format('sphericalbck1'  , MDOptions.spheric_k)
                     strbuff += '{:24}{}\n'      .format('sphericalbcexp1', MDOptions.spheric_exponent)
                # TODO
                elif MDOptions.boundary.lower() in [ 'cylindric', 'cylindrical' ]:
                     pass

            if self.wrap_all:
                strbuff += '{:24}{}\n'.format('wrapall'  , bools[self.wrap_all])
            if self.wrap_water:
                strbuff += '{:24}{}\n'.format('wrapwater', bools[self.wrap_water])

            if MDOptions.langevin:
                strbuff += '\n{:24}{}\n'.format('langevin'        , 'on')
                strbuff += '{:24}{}\n'.format('langevintemp'    , MDOptions.langevin_temperature)
                strbuff += '{:24}{}\n'.format('langevindamping' , MDOptions.langevin_damping)
                # NAMD default: on
                if not MDOptions.langevin_H:
                    strbuff += '{:24}{}\n'.format('langevinhydrogen', 'off')

        # single-point?
        else:
            strbuff += '{:24}{}\n'.format('timestep'     , 1)
            strbuff += '{:24}{}\n'.format('numsteps'     , 1.0)

        if self.fixed_atoms:
             strbuff += '\n# fixed atoms'
             strbuff += '\nfixedatoms       {}'.format('on')
             strbuff += '\nfixedatomsfile   {}'.format(self.fixed_atoms)
             strbuff += '\nfixedatomscol    {}'.format('b')
             strbuff += '\nfixedatomsforces {}\n'.format(bools[self.fixed_atoms_forces])

        if self.constraints:
             strbuff += '\n# harmonic constraints'
             strbuff += '\nconstraints      {}'.format('on')
             strbuff += '\nconsref          {}'.format(self.constraints_ref)
             strbuff += '\nconskfile        {}'.format(self.constraints)
             strbuff += '\nconskcol         {}\n'.format('b')

        strbuff += '\n# force field\n'
        if MDOptions.ff.lower().strip() == 'charmm':
            strbuff += '{:24}{}\n'.format('paratypecharmm', 'on')
        if self.merge_cross:
            strbuff += '{:24}{}\n'.format('mergecrossterms', 'on')

        strbuff += '{:24}{}\n'.format('exclude'      , self.exclude)
        strbuff += '{:24}{}\n'.format('1-4scaling'   , self.scaling14)
        strbuff += '{:24}{:12.3f}\n'.format('margin'       , self.margin*unitconvert.Bohr2Angstrom)
        strbuff += '{:24}{}\n'.format('switching'    , bools[self.switching])
        strbuff += '{:24}{:12.3f}\n'.format('switchdist'   , self.switch_dist*unitconvert.Bohr2Angstrom)
        strbuff += '{:24}{:12.3f}\n'.format('cutoff'       , self.cutoff*unitconvert.Bohr2Angstrom)
        strbuff += '{:24}{:12.3f}\n'.format('pairlistdist' , self.pairlist_dist*unitconvert.Bohr2Angstrom)

        # PME
        # YL 01/03/2021: to use PME the system has to be periodic, but cell vectors won't
        #                have to be written because they can be read in from .xsc file
        if self.pme and MDOptions.is_periodic:
            pme_grid_sizes = self.pme_grid_sizes
            strbuff += '\n{:24}\n'.format('# PME')
            strbuff += '{:24}{}\n'.format('pme'          , yorn[self.pme])
            strbuff += '{:24}{}\n'.format('pmegridsizex' , pme_grid_sizes[0])
            strbuff += '{:24}{}\n'.format('pmegridsizey' , pme_grid_sizes[1])
            strbuff += '{:24}{}\n'.format('pmegridsizez' , pme_grid_sizes[2])

        # NPT
        strbuff += '\n{:24}\n'.format('# NPT/NVT setting')
        strbuff += '{:24}{}\n'.format('usegrouppressure', bools[self.group_pressure])
        strbuff += '{:24}{}\n'.format('useflexiblecell' , bools[self.flexible_cell])
        strbuff += '{:24}{}\n'.format('useconstantarea' , bools[self.constant_area])
        if MDOptions.langevin_piston:
            strbuff += '\n{:24}\n'.format('# Langevin piston method')
            strbuff += '{:24}{}\n'.format('langevinpiston'      , yorn[MDOptions.langevin_piston])
            strbuff += '{:24}{}\n'.format('langevinpistontarget', MDOptions.langevin_piston_target)
            strbuff += '{:24}{}\n'.format('langevinpistonperiod', MDOptions.langevin_piston_period)
            strbuff += '{:24}{}\n'.format('langevinpistondecay' , MDOptions.langevin_piston_decay)
            strbuff += '{:24}{}\n'.format('langevinpistontemp'  , MDOptions.langevin_piston_temp)


        if MDOptions.minimise:
            strbuff += '\n# minimise the system to eliminate bad initial contacts'
            strbuff += '\n{:24}{}\n'.format('minimize', MDOptions.minimise)

        if MDOptions:
            strbuff += '\n# reinitialise the velocities to the desired target temperature'
            strbuff += '\n{:24}{}\n'.format('reinitvels', MDOptions.temperature)
            strbuff += '\n# run the MD simulation: {} fs * {:d} = {} ns'.format(MDOptions.timestep, MDOptions.nsteps, MDOptions.timestep*MDOptions.nsteps/1.E6)
            strbuff += '\n{:24}{}\n\n\n'.format('run', MDOptions.nsteps)

        return strbuff



    def parseMainOutput(self):
        ''''''

        from numpy   import all, append, arange, diff
        from ...data import unitconvert
        from ...io   import file

        self.setOutKeys()

        selectcases = { 'energies'    : unitconvert.kcal_mol2Ha,
                        'potentials'  : unitconvert.kcal_mol2Ha,
                        'kinetics'    : unitconvert.kcal_mol2Ha,
                      }

        # YL 27/02/2021: we first get the temperature points of MD data (to be distinguished with the minimisation output)
        abuff = file.getArrayFromFile(self.output,
                                          'out',
                                          (1),
                                           regex=self._outkeys['temperatures'])

        # YL 27/02/2021: check the points for MD and minimisation, respectively
        masks_md_points = abuff > 0.0
        indices_all_points = arange(len(abuff))
        if all(diff(indices_all_points[masks_md_points]) == 1) and self.temperature > 0.0:
            indices_opt_points = indices_all_points[~masks_md_points]
            indices_md_points  = indices_all_points[masks_md_points]
        else:
            print("\n ChemShell/NAMD WARNING: The output temperatures might be problematic")
            indices_opt_points = ...
            indices_md_points  = ...

        # save the temperature data
        setattr(self._result, 'temperatures', abuff[indices_md_points])

        # save the values
        for k, v in selectcases.items():

            abuff = file.getArrayFromFile(self.output,
                                              'out',
                                              (1),
                                               regex=self._outkeys[k])*v

            setattr(self._result, k, abuff[indices_md_points])

            # save the miniminsaton energies
            if k == 'energies':
                setattr(self._result, 'opt_energies', abuff[indices_opt_points])

        # don't move the files because they may be needed by a higher level of tasks
        # at last: tidy up by copying all files to a folder
#        self.archive(dest=self.prefix)

        # YL 27/07/2023: merge the output files (DCD and XST) for a restarted job
        if self.restart and self.merge_restarted:
            from copy    import deepcopy
            from os.path import join
            from ...io   import files, format
            if self._MDOptions:
                files.dcd.merge(self.prefix_restart+'.dcd', self._MDOptions.dcd, f'{self.prefix}_merged.dcd')
            else:
                files.dcd.merge(self.prefix_restart+'.dcd', self._dcd, f'{self.prefix}_merged.dcd')
            files.xst.merge(self.prefix_restart+'.xst', self.xst     , f'{self.prefix}_merged.xst')

            # YL 20/06/2024: also merge the MD profiles
            namdObj = deepcopy(self)
            namdObj.output = self.prefix_restart+'.out'
            namdObj.restart = namdObj.merge_restarted = False
            namdObj.parseMainOutput()
            format.printwrap(f'Merging {namdObj.output} of {namdObj._result.temperatures.size} records into {self.output} of {self._result.temperatures.size} records...')
            npts = self._result.temperatures.size
            self._result.resize(namdObj._result.temperatures.size+self._result.temperatures.size)
            # YL 20/06/2024: do NOT redefine attributes of self._result directly which will trigger resizing each time! (see <class 'ResultMD'>)
            self._result.temperatures[:] = append(namdObj._result.temperatures, self._result.temperatures[:npts])
            self._result.energies    [:] = append(namdObj._result.energies    , self._result.energies[:npts])
            self._result.potentials  [:] = append(namdObj._result.potentials  , self._result.potentials[:npts])
            self._result.kinetics    [:] = append(namdObj._result.kinetics    , self._result.kinetics[:npts])
            try:
                self._result.opt_energies[:] = append(namdObj._result.opt_energies, self._result.opt_energies[:npts])
            except:
                pass
            del namdObj


    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
        self._outkeys = {
                          'kinetics'    :{'ENERGY:':[0,0,1,10,10,1,0,1.0]},
                          'energies'    :{'ENERGY:':[0,0,1,11,11,1,0,1.0]},
                          'temperatures':{'ENERGY:':[0,0,1,12,12,1,0,1.0]},
                          'potentials'  :{'ENERGY:':[0,0,1,13,13,1,0,1.0]},
#                          'gradients' :{'gradient \(hartree':[4,4+self.frag.natoms-1,1,6,8,1,1,1.0]},
                        } 


    @property
    def path(self):
        '''Path to exec for the system command'''

        from os.path  import join
        from ...utils import testutils
        from ...base  import errors
        try:
            from . import namd_path
        except:
            # YL 27/02/2023: failure of importing namd_path means NAMD isn't compiled with Py-ChemShell
            testutils.setUnsupported()
            exit(errors.UNSUPPORTED)

        # path: exec at buildtime for system command for now
        return self._path or join(namd_path.namd_path, 'namd2')


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

        self._path = val


    @property
    def sysmpi(self):
        '''System MPI command: it should be charm++ run command for NAMD'''

        from os.path import join, isfile
        from ...     import parallel

        # import with illegal characters
        charm = getattr(parallel, 'charm++')

        if not self._sysmpi and not charm.charm_path:
            print("\n ChemShell/NAMD WARNING: charmrun not specified or found! NAMD may fail!\n")

        # use the built-in one that was installed when NAMD was compiled with ChemShell
        return self._sysmpi or charm.charm_path


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

        self._sysmpi = val


    @property
    def runargs(self):
        '''Runtime arguments'''

        from os.path import join
        from numpy   import unique
        from .       import namd_path
        from ...     import parallel
        from ...io   import format

        # see: http://www.ks.uiuc.edu/Research/namd/2.14/notes.html
        # YL 07/03/2021: for non-MPI multicore version (multithreads run)
        if parallel.parallel.nprocs_mm == 1:
            return self._runargs or f'+auto-provision {self.input}'
        # YL 07/03/2021: multiple nodes (MPI run) (assuming infiniband or other high-performance networks)
        elif parallel.parallel.nprocs_mm > 1:
            # currently for srun only (e.g., on ARCHER2)
            if 'srun' in namd_path.sysmpi:
                if self._runargs_print_hint:
                    format.printwrap('HINT: if NAMD hangs and you see errors such as "step creation temporarily disabled, retrying (Requested nodes are busy)", try to rerun ChemShell with `chemsh --mem-ratio [...]`', indent='', initial_indent=' ', subsequent_indent=' '*7)
                    self._runargs_print_hint = False
                # YL 27/07/2023: due to change of srun, we need `--overcommit --oversubscribe --overlap --mem=0` for oversubscribing
#                self.sysmpi = f'{namd_path.sysmpi} --oversubscribe --overlap --overcommit --mem={int(self.memory_per_proc*parallel.parallel.nprocs_mm)}M -n {parallel.parallel.nprocs_mm} -c 1'
                self.sysmpi = f'{namd_path.sysmpi} --overcommit --overlap --oversubscribe --mem=0 -n {parallel.parallel.nprocs_mm} -c 1'
                return self._runargs or f'{self.input}'
            else:
                return self._runargs or f'+p{parallel.parallel.nprocs_mm} {self.input}'


    @runargs.setter
    def runargs(self, val):
        ''''''

        self._runargs = val


    @staticmethod
    def mergeParameters(files, **kwargs):
        '''A wrapper function of .tools.mergeParamters()'''

        from .tools import mergeParameters

        mergeParameters(files, **kwargs)


    # see NAMD User's Guild Chapter 4
    def psfgen(self, dryrun=False, return_frag=False, output_dir='_psfgen', **kwargs):
        '''Run psfgen and generate PSF from PDB'''

        from sys        import stdout
        from os         import path
        from .psfgen    import psfgen
        from ...objects import fragbio
        from ...io.file import parse

        # topologies are mandatory to psfgen but not in NAMD
        if len(self.top):
            self._psf, self._pdb, self._additional_prm, self._solvents = psfgen(self.frag, dryrun=dryrun, topologies=self.top, output_dir=output_dir, **kwargs)

        else:
            self._psf, self._pdb, self._additional_prm, self._solvents = psfgen(self.frag, dryrun=dryrun, output_dir=output_dir, **kwargs)
        # YL 24/07/2023: we need a switch in Solvation
        if self._pdb:
            self._default_pdb = True
        if self._psf:
            self._default_psf = True

        # read in PDB and return the frag
        if return_frag:
            fragObj = fragbio.FragmentBio(coords=self._pdb)
            # assign attributes such as charges, types
            parse(self._psf, 'psf', receiver=fragObj)
            # YL 26/11/2021: inherit all historic renaming
            fragObj._renamed_residues.update(self.frag._renamed_residues)
            return fragObj
        else:
            # a dummy object to hold the PSF information
            fragObj = fragbio.FragmentBio()
            # assign attributes such as charges, types
            parse(self._psf, 'psf', receiver=fragObj)

        print('\n '+'-'*70)
        print(' Summary of PSFGEN run')
        print(' '+'-'*70)
        print(' Number of atoms')
        print(f' Input : {self.frag.natoms}')
        print(f' Output: {fragObj._psf.natoms}')
        print(f' New PSF saved as: {self._psf}')
        print(' '+'-'*70)
        stdout.flush()


    def archive(self, dest='_namd', source='.'):
        '''Archive exsiting directory and files'''

        from os.path  import join

        # archive old files
        fileutils.archiveDir(dest)
        fileutils.archiveFile(join(source, self.input ), dir=dest, suffix='')
        fileutils.archiveFile(join(source, self.output), dir=dest, suffix='')
        fileutils.archiveFile(join(source, self.coor  ), dir=dest, suffix='')
        fileutils.archiveFile(join(source, self.vel   ), dir=dest, suffix='')
        fileutils.archiveFile(join(source, self.xsc   ), dir=dest, suffix='')
        fileutils.archiveFile(join(source, self.xst   ), dir=dest, suffix='')
        fileutils.archiveFile(join(source, self._dcd  ), dir=dest, suffix='')
        fileutils.archiveFile(join(source, self._par  ), dir=dest, suffix='')
        fileutils.archiveFile(join(source, self._chk+'.coor'), dir=dest, suffix='')
        fileutils.archiveFile(join(source, self._chk+'.vel' ), dir=dest, suffix='')
        fileutils.archiveFile(join(source, self._chk+'.xsc' ), dir=dest, suffix='')
        fileutils.archiveFile(join(source, self._chk+'.coor.old'), dir=dest, suffix='')
        fileutils.archiveFile(join(source, self._chk+'.vel.old' ), dir=dest, suffix='')
        fileutils.archiveFile(join(source, self._chk+'.xsc.old' ), dir=dest, suffix='')
        # YL 04/09/2021: also save a PSF file
        self.frag.save(dest+'.psf')
        fileutils.archiveFile(join(source, dest+'.psf'), dir=dest, suffix='')


    @property
    def pme_grid_sizes(self):
        '''PME grid X/Y/Z sizes'''

        from numpy   import count_nonzero
        from math    import ceil
        from ...data import unitconvert

        if not count_nonzero(self._pme_grid_sizes):
            def _ceil_235(f):
                '''Return the ceiling integer of a given float which is divisible by 2, 3, and 5'''
                c = ceil(f)
                while c%2 and c%3 and c%5:
                    c += 1
                return c

            vectors = self.frag._xyz_spanning*unitconvert.Bohr2Angstrom
            _pme_grid_sizes = [ _ceil_235(vectors[f]) for f in range(3) ]
            print('\n PME grid sizes determined automatically:', *_pme_grid_sizes)

        return _pme_grid_sizes


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

        # default: unique=True
        val = iterutils.getFlattened(val, unique=False)

        if len(val) != 3:
            print('\n ChemShell/NAMD ERROR: pme_grid_sizes must be a list/tuple/array of size 3')
            exit(112)

        self._pme_grid_sizes = val


    # YL 23/02/2023: a mechanism to read in the PSF data into fragment
    @property
    def psf(self):
        '''PSF file'''

        return self._psf


    @psf.setter
    def psf(self, val):
        '''Setter of PSF file'''

        from ...io.file import parse

        if val.strip():
            # skip if frag._psf is already initialised
            try:
                self.frag._psf.titles
            # skip if frag._psf is already initialised
            # load the PSF data into the fragment
            except:
                parse(val, 'psf', receiver=self.frag)

        self._psf = val



