#  Copyright (C) 2020 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@stfc.ukri.org'

# YL 04/03/2021: heavily revised to make it object-oriented

from sys             import stdout
from .               import _Workflows
from ..base          import errors
from ..data          import unitconvert
from ..dl_py2f.utils import fileutils, iterutils, miscutils

# we provide users some synonyms
_synons = {
          }

class Solvation(_Workflows):
    '''Parent class of workflows'''

    from ..objects    import fragbase
    from ..interfaces import namd
    from ..tasks      import md

    _kwargs = {
                'constrained_npt'       :'_constrained_npt.pdb',
                'constrained_npt_ref'   :'_constrained_npt_ref.pdb',
                'driver'                : namd.NAMD(),     # NAMD as the default MD driver
                'dryrun'                : False,
                'ff'                    :'charmm',
                'fixed_npt'             :'_fixed_npt.pdb',
                'fixed_nvt'             :'_fixed_nvt.pdb',
                'fix_extra_npt'         :([],{}),          # use a tuple to accept multiple legal types, same below
                'fix_extra_nvt'         :([],{}),          # YL 05/07/2023: in a tuple, list and dict are mutually exclusive, so that this won't initialise a list of dict ([{...}])
                'harmonic_constraint'   : 5.0,
                'length_npt'            : 5000000,         # in fs (timestep: 2 fs)
                'length_nvt'            : 2000000,         # in fs (timestep: 2 fs)
                'length_production'     : 20000000,        # in fs (timestep: 2 fs)
                'max_ncycles'           : 20,
                'md'                    : md.MD(),         # YL 04/03/2021: MD instance, allowed input from the user
                'minimisation_npt'      : 50000,
                'minimisation_nvt'      : 50000,
                'neutralise'            : True,
                'neutralise_anion'      :'Cl',
                'neutralise_cation'     :'Na',
                'neutralise_ion_segment':'',
                'nsnapshots'            : 20,
                'pbc_wrap'              : True,
                'pdb2pqr'               : True,            # this can be a str path
                'pH'                    : None,
                'prefix_neutralised'    :'_neutralised',
                'prefix_solvated'       :'_solvated',
                'prefix_npt'            :'_02_npt',
                'prefix_nvt'            :'_01_nvt',
                'prefix_production'     :'_03_production',
                # YL 27/02/2021: TODO 'nucleotides'
                'relax_extra_npt'       :([],{}),
                'relax_extra_nvt'       :([],{}),
                'relax_residues_nvt'    :['solvents','ions'],
                'relax_residues_npt'    :['amino_acids','dna','rna','solvents','ions'],
                'restart_from_md'       : 0,               # YL 04/03/2021: added restart feature
                'restart_script'        :'_continued_production_md.py',
                'solute'                : None,
                'threshold'             : 100,             # TODO
              }

#    _kwargs.update()

    def printHeader(self, msg):
        '''Print a workflow section header.'''

        headerwidth = 70

        print('\n\n ' + '-'*headerwidth)
        print(' Solvation: {}'.format(msg))
        print(' '+'-'*headerwidth+'\n')


    def runPDB2PQR(self):
        '''Wrapper function of self._runPDB2PQR() returning a solvated fragment object'''

        # with PDB2PQR
        if self.pdb2pqr:
            pqrFrag = self._runPDB2PQR()
            # we don't need to connect or check the contiguousness again because the
            # solute must be contiguous already, whether provided by a user-defined
            # PDB (see the PDB parser) or by Py-ChemShell
            return pqrFrag.solvate(**self._run_kwargs, reconnect=False, recentre=True)

        # without PDB2PQR (not recommended!)
        else:
            # rename the residues to find out terminals
            from ..addons import pdb2pqr
            self.solute.detectNucleotides(aliases=pdb2pqr.charmm2pqr, delete_terminal_hydrogen=True, delete_terminal_phosphate=True)
            self.solute.detectNucleotides(aliases=pdb2pqr.pqr2chemsh, delete_terminal_hydrogen=False, delete_terminal_phosphate=False)
            return self.solute.solvate(**self._run_kwargs, reconnect=False, recentre=True)


    def _runPDB2PQR(self):
        '''Run PDB2PQR and return the new fragment object'''

        from numpy     import char, isin, unique
        from os.path   import join
        from ..objects import fragbio, fragment
        from ..addons  import pdb2pqr
        from ..io      import file, format

        self.printHeader("running PDB2PQR... (pH by PROPKA)")

        # run
        pdb2pqr_input = '_pdb2pqr.pdb'
        pdb2pqr_dir   = '_pdb2pqr'
        pdb2pqr_out   = join(pdb2pqr_dir, '_pdb2pqr.out')
        pdb2pqr_pqr   = join(pdb2pqr_dir, '_pdb2pqr.pqr')

        # PDB2PQR can't assign charges to resname "TIP3"
        tip3 = self.solute.select(residues='TIP3')
        if len(tip3):
            self.solute.residues[tip3] = 'HOH'
            format.printwrap('NB: Residues \"TIP3\" renamed as \"HOH\" for PDB2PQR')

        # write out a pdb as the input file
        self.solute.save(pdb2pqr_input, version='pdb2pqr')

        ierror = 1
        while ierror:

            # allow user-defined path
            ierror = pdb2pqr.run(pdb2pqr_input, output_dir=pdb2pqr_dir, path={True:''}.get(self.pdb2pqr, self.pdb2pqr))
            fileutils.archiveFile(pdb2pqr_input, dir=pdb2pqr_dir, suffix='')

            # some raw crystal structures may not be complete
            if ierror:
                error_key = 'Too few atoms present to reconstruct or cap residue'

                # file-based PDB2PQR output
                try:
                    with open(pdb2pqr_out) as fp:
                        for lineno, line in enumerate(fp):
                            if error_key in line:
                                break
                # returned value of direct PDB2PQR call as a library
                except FileNotFoundError:
                    line = ierror.args[0]

                except:
                    format.printwrap(">>> ChemShell/PDB2PQR ERROR: PDB2PQR failed")
                    exit(999)

                # parse the line containing error message
                broken_resname, broken_chainID, broken_resID = line.partition(error_key)[-1].split()[:3]
                broken_resname, broken_chainID, broken_resID = broken_resname.strip(), broken_chainID.strip(), int(broken_resID)
                print(format.wrap("ChemShell Solvation: trying to remove the broken residue and rerun PDB2PQR..."))
                # delete the broken residue's atoms
                broken_residue = self.solute.select(residues=broken_resname, resIDs=broken_resID, chainIDs=broken_chainID)
                broken_segment = unique(self.solute.segments[broken_residue])[0]
                print("\n Residue name: {} (residue ID: {}, segment name: {}, chain ID: {})\n".format(broken_resname,
                                                                                                      broken_resID,
                                                                                                      broken_segment.decode().strip(),
                                                                                                      broken_chainID))
                # check if the broken residue is at the terminals
                selected_segment = self.solute.select(segments=broken_segment)
                min_res = self.solute.resIDs[selected_segment].min()
                max_res = self.solute.resIDs[selected_segment].max()

                # get the resIDs of residues that are connected to the current broken one and belong to the current broken segment
                unique_connected_atoms = unique(self.solute.conn[broken_residue])
                masks_connected_atoms_in_segment = isin(unique_connected_atoms, selected_segment)
                unique_connected_atoms = unique_connected_atoms[masks_connected_atoms_in_segment]
                unique_connected_resIDs = unique(self.solute.resIDs[unique_connected_atoms[isin(unique_connected_atoms, broken_residue, invert=True)]])

                # YL 27/02/2021: we issue a warning here if the broken residue might not be a terminal
                self.solute.delete(broken_residue)
#                if any(in1d(unique_connected_resIDs, range(min_res, broken_resID+1))) \
#                or any(in1d(unique_connected_resIDs, range(broken_resID, max_res+1))):
                if len(unique_connected_resIDs) > 1:
                    print(format.wrap("ChemShell Solvation WARNING: the deleted broken residue \"{}\" (ID: {}) might not be a terminal! Following calculations may be wrong!".format(broken_resname.strip(), broken_resID)))
                # shift the terminal's sequence
                else:
                    # YL FIXME 27/02/2021: the terminal resID is wrong because of self.solute.optimiseResIDs()
                    #                      and the current code may be wrong because the broken residue may not be
                    #                      part of protein/DNA
                    self.solute.terminal.resID = unique_connected_resIDs[0]
                    self.solute.terminal.residue = self.solute.resID2residue(self.solute.terminal.resID, segment=broken_segment)
                # write out a new file
                self.solute.save(pdb2pqr_input, version='pdb2pqr')

        # 27/04/2022: we create a fragment from the resulting PDB file as the reference to retain the chainIDs
        pdbFrag = fragbio.FragmentBio()
        file.parse( pdb2pqr_out+'.pdb',
                   'pdb',
                    receiver=pdbFrag,
                    ff=self.ff,
                    force_reassign_segments=True, # to make the fragment same as the PQR one we force the reassignment to be done
                    mute=True,
                    choose_locations=False,
                    attach_to_protein=False)      # to make the fragment same as the PQR one we can't check the connectivity to protein

        # read in the resulting PQR file which should contain only amino/nucleic acids and solvents
        # YL 28/01/2021: segment names are assigned
        pqrFrag = fragment.Fragment(coords=pdb2pqr_pqr)
        pqrFrag.title = self.solute.title

        # read in the indices of atoms omitted by PDB2PQR: NB indices start with 1!
        # YL NB 28/01/2021: there is an implicit interconnectivity check for the residues (excluding water)
        omitted_atoms = file.parse(pdb2pqr_pqr, 'pqr', omitted_only=True)['indices'] - 1
        # for a new frag with the PQR-processed and the PQR-omitted ones
        omittedFrag = self.solute.getSelected(omitted_atoms, flatten=False) # flatten=False is faster than True (default), though negligibly here
        # YL TODO 29/01/2022: the PDB2PQR radii for nucleotides look suspicious and cause wrong connectivity
        #                     and here is an ad hoc way to avoid reconnecting using them
        if omittedFrag.natoms:
            newFrag = pqrFrag.merge(omittedFrag, reconnect=True, reset_bfactors=True)
        else:
            newFrag = pqrFrag
        newFrag.title = pqrFrag.title
        # YL 26/11/2021: inherit the renaming records
        newFrag._renamed_residues.update(self.solute._renamed_residues)

        # YL 28/01/2021: make sure all segments are contiguous in the atom sequence
        segments_newFrag = unique(newFrag.segments)
        for segment_newFrag in segments_newFrag:
            newFrag.keepSegmentContiguous(segment_newFrag)
        # YL 29/01/2021: the omitted residues are not recognised by PDB2PQR so that they won't be
        #                accepted by PSFGEN with standard CHARMM forcefield, either; they must be
        #                in separate segements, respectively
        residues_omittedFrag = unique(omittedFrag.residues)
        for residue_omittedFrag in residues_omittedFrag:
            newFrag.segments[newFrag.select(residues=residue_omittedFrag)] = residue_omittedFrag

        # 27/04/2022: keep the original chainIDs, suggested by HMS
        newFrag.retainChainIDs(pdbFrag)

        return newFrag


    def runNeutralise(self):
        '''Neutralise the system by adding ions'''

        if self.neutralise:
            self.printHeader("neutralising the system...")
            self._psfFrag.neutralise(cation=self.neutralise_cation,
                                     anion=self.neutralise_anion,
                                     ion_segment=self.neutralise_ion_segment,
                                     # YL TODO 11/09/2021: solvent segments name (for unstrict matching); this is currently hardcoded
                                     solvent='SOL',
                                     ff=self.ff)

            # save to files which will be used for the MDs
            self._psfFrag.save(f'{self.prefix_neutralised}.pdb')
            self._psfFrag.save(f'{self.prefix_neutralised}.psf')
            # TODO: run PSFGEN again
            # in self.driver we could run psfgen again but it'd be even more complicated
#            self.driver.psfgen(return_frag=False)


    def initMD(self):
        '''Initialise the MD'''

        from ..objects import fragment
        from ..io      import file

        # reassign the driver's fragment
        try:
            self.driver.frag = self._psfFrag
        # restarted job from saved files
        except:
            self._psfFrag = fragment.Fragment(coords=f'{self.prefix_neutralised}.pdb')
            file.parse(f'{self.prefix_neutralised}.psf', 'psf', receiver=self._psfFrag)
            self.driver.frag = self._psfFrag

        self.driver.merge_cross = True
        self.driver.switching   = True
        self.driver.wrap_all    = True
        self.driver.wrap_water  = True
        # set options if not specified by the user
        # note these override any defaults set by the (NAMD) driver class
        self.driver._setattr('freq_full_elect' ,  2)
        self.driver._setattr('nsteps_per_cycle', 10)

        # use psfgen outputs if not specified otherwise by the user
        # YL 24/07/2023: improved as only self.driver.psf is set after PSFGEN
        if self.neutralise:
            if not self.driver.psf or self.driver._default_psf:
                self.driver._psf = f'{self.prefix_neutralised}.psf'
                self.driver.psf = self.driver._psf
            stdout.flush()
            if not self.driver.pdb or self.driver._default_pdb:
                self.driver._pdb = f'{self.prefix_neutralised}.pdb'
                self.driver.pdb = self.driver._pdb

        # YL 19/01/2022: removed such check because: NAMD doesn't need an RTF file; the parameters file should have been assigned by PSFGEN already
        # check if topologies/parameters are specified
#        if not len(self.driver.top):
#            raise errors.ChemShGeneralError('Please specify topology files for the NAMD instance using \"top=\"')
#        if not len(self.driver.par):
#            raise errors.ChemShGeneralError('Please specify parameter files for the NAMD instance using \"par=\"')

        # initialise the MD object
        self.md.driver = type(self.driver).__name__.lower() # must be a str driver name
        self.md.theory = self.driver
        self.md.ff     = self.ff
        self.md._setattr('langevin_damping',  5)
        self.md._setattr('rigid'           , 'all')
        self.md._setattr('timestep'        ,  2.0)


    ### 1st MD: minimisation; NVT MD with solute fixed
    @miscutils.time
    def runNVT(self):
        '''Run the 1st MD (NVT)'''

        from ..io    import format

        self.printHeader("performing the 1st MD (NVT)...")

        dryrun = self.dryrun

        # default frequency of writing/printing data/result: nsteps/1000
        # recommended: at least 50,000 steps
        self.md.minimise        = self.minimisation_nvt
        # typically 2 ns: 2 fs * 1000000
        self.md.nsteps          = int(self.length_nvt/self.md.timestep)
        self.md.nsnapshots      = 0
        self.md.boundary        = 'periodic'
        self.md.langevin        = True
        self.md.langevin_piston = False       # no piston for NVT

        # TODO: DNA/RNA
        # fix all but solvents
        print("\n Fixing atoms except for solvents (incl. ions)...")
        masks_to_fix  = self._getFixedIndices('nvt', self.getResiduesToRelax('nvt', self.relax_extra_nvt), **self.selections_to_relax_nvt)
        # YL 15/12/2021: added extra selected atoms to fix
        masks_to_fix += self._getExtraFixedIndices('nvt', self.getResiduesToFix('nvt', self.fix_extra_nvt), **self.selections_to_fix_nvt)

        # fixed: bfactor=1; relaxed: bfactor=0
        self._psfFrag.occ[:]      = 1.0
        self._psfFrag.bfactors[:] = 0.0
        self._psfFrag.bfactors[masks_to_fix] = 1.0
        fixed_nvt = self._run_kwargs.get('fixed_nvt', '_fixed_nvt.pdb')
        self._psfFrag.save(fixed_nvt)

        # restore bfactor values
#        self._psfFrag.bfactors[masks_to_fix], bfactors = 1.0, self._psfFrag.bfactors[masks_to_fix]
        self._psfFrag.bfactors[:] = 0.0

        self.driver.restart            = False
        self.driver.merge_restarted    = False # YL 08/08/2023: we must set this off which is defaulted to True
        self.driver.group_pressure     = False
        self.driver.flexible_cell      = False # False for water box & True for membrane
        self.driver.constant_area      = False # False for water box & True for membrane
        self.driver.fixed_atoms        = fixed_nvt
        self.driver.fixed_atoms_forces = True

        # run
        ierror = self.md.run(dryrun=dryrun, info_filename=f'{self.prefix_nvt}.npz')

        # rename folder by archiving
        self.driver.archive(dest=self.prefix_nvt)
        fileutils.archiveFile(fixed_nvt     , dir=self.prefix_nvt, suffix='')
        fileutils.archiveFile(f'{self.prefix_nvt}.npz', dir=self.prefix_nvt, suffix='')
        print(format.wrap("{} ns NVT MD simulation has run successfully. Results saved in:\n {}".format(self.md.timestep*self.md.nsteps/1000000, self.prefix_nvt)))


    ### 2nd MD: minimisation; NPT MD with backbone constraints
    @miscutils.time
    def runNPT(self):
        '''Run the 2nd MD (NPT)'''

        from os.path import join
        from ..io    import format

        self.printHeader("performing the 2nd MD (NPT)...")

        dryrun = self.dryrun

        # minimal 50,000 steps
        self.md.minimise        =  self.minimisation_npt
        # typically 5 ns: 2 fs * 2,500,000
        self.md.nsteps          =  int(self.length_npt/self.md.timestep)
        self.md.nsnapshots      =  0
        # YL 01/03/2021: mdOBj.is_period must be True (iff self.md.boundary is not "droplet" or "sphere") to make sure PME works
        #                and self.driver.restart can suppress writing the cell vectors (but reading from .xsc file instead)
        self.md.boundary        = 'periodic'
        self.md.langevin        =  True
        self.md.langevin_piston =  True

        # fix all but relax solvents and major solute (protein, DNA, lipid) excluding cofactors
        print(format.wrap("Relaxing the major solutes and solvents (incl. ions), with the rest (e.g., cofactors/ligands) remaining fixed..."))
        masks_to_fix = self._getFixedIndices('npt', self.getResiduesToRelax('npt', self.relax_extra_npt), **self.selections_to_relax_npt)
        # YL 15/12/2021: added extra selected atoms to fix
        masks_to_fix += self._getExtraFixedIndices('npt', self.getResiduesToFix('npt', self.fix_extra_npt), **self.selections_to_fix_npt)

        # fixed: bfactor=1; relaxed: bfactor=0
        self._psfFrag.bfactors[masks_to_fix] = 1.0
#        self._psfFrag.bfactors[masks_to_fix], bfactors = 1.0, self._psfFrag.bfactors[masks_to_fix]
        fixed_npt = self._run_kwargs.get('fixed_npt', '_fixed_npt.pdb')
        self._psfFrag.save(fixed_npt)
        # restore bfactor values
#        self._psfFrag.bfactors[masks_to_fix], bfactors = 1.0, self._psfFrag.bfactors[masks_to_fix]

        # harmonic constraints
        print("\n Applying harmonic constraints on the backbone atoms...")
        constrained_npt     = self._run_kwargs.get('constrained_npt'    , '_constrained_npt.pdb')
        constrained_npt_ref = self._run_kwargs.get('constrained_npt_ref', '_constrained_npt_ref.pdb')
        self._psfFrag.occ[:]      = 0.0
        self._psfFrag.bfactors[:] = 0.0
        indices_backbone = self._psfFrag.selectByBackbone()
        self._psfFrag.bfactors[indices_backbone] = self.harmonic_constraint
        self._psfFrag.save(constrained_npt)

        # restore and save a reference
        self._psfFrag.bfactors[:] = 0.0
        self._psfFrag.save(constrained_npt_ref)

        # apply constraints
        self.driver.constraints     = constrained_npt
        self.driver.constraints_ref = constrained_npt_ref

        # restart from step 1
        self.driver.restart            = True
        self.driver.merge_restarted    = False # YL 08/08/2023: we must set this off which is defaulted to True
        self.driver.group_pressure     = True
        self.driver.flexible_cell      = False # False for water box & True for membrane
        self.driver.constant_area      = False # False for water box & True for membrane
        self.driver.fixed_atoms        = fixed_npt
        self.driver.fixed_atoms_forces = False
        self.driver.prefix_restart     = join(self.prefix_nvt, self.driver.prefix)

        # run
        ierror = self.md.run(dryrun=dryrun, info_filename=f'{self.prefix_npt}.npz')

        # rename folder by archiving
        self.driver.archive(dest=self.prefix_npt)
        fileutils.archiveFile(fixed_npt          , dir=self.prefix_npt, suffix='')
        fileutils.archiveFile(constrained_npt    , dir=self.prefix_npt, suffix='')
        fileutils.archiveFile(constrained_npt_ref, dir=self.prefix_npt, suffix='')
        fileutils.archiveFile(f'{self.prefix_npt}.npz'     , dir=self.prefix_npt, suffix='')
        print(format.wrap("{} ns NPT MD simulation has run successfully. Results saved in:\n {}".format(self.md.timestep*self.md.nsteps/1000000, self.prefix_npt)))


    ### 3rd MD: production MD (fully relaxed except for cofactors/ligands)
    @miscutils.time
    def runProduction(self):
        '''Run the production MD (NPT)'''

        from os.path import join
        from ..io    import format

        self.printHeader("performing the 3rd MD (production NPT)...")

        dryrun = self.dryrun

        # no minimisation required
        self.md.minimise        =  0
        # typically 20 ns: 2 fs * 10,000,000
        self.md.nsteps          =  int(self.length_production/self.md.timestep)
        # also take snapshots
        self.md.nsnapshots      =  self.nsnapshots
        # YL 01/03/2021: mdOBj.is_period must be True (iff self.md.boundary is not "droplet" or "sphere") to make sure PME works
        #                and self.driver.restart can suppress writing the cell vectors (but reading from .xsc file instead)
        self.md.boundary        = 'periodic'
        self.md.langevin        =  True
        self.md.langevin_piston =  True

        # pbc wrap when taking snapshots
        if self.pbc_wrap:
            # normally non-solvent (solvent includes ions)
            indices_pbc_wrap = self._psfFrag.select(residues=self.getResiduesToRelax('nvt', self.relax_extra_nvt), strict=True, invert=True)
            self.md.pbc_wrap = indices_pbc_wrap

        # fix all, but relax solvents and major solute (protein, DNA, lipid) excluding cofactors
        print(format.wrap("Relaxing the major solutes and solvents (incl. ions), with the rest (e.g., cofactors/ligands) remaining fixed..."))
        masks_to_fix = self._getFixedIndices('npt', self.getResiduesToRelax('npt', self.relax_extra_npt), **self.selections_to_relax_npt)
        # YL 15/12/2021: added extra selected atoms to fix
        masks_to_fix += self._getExtraFixedIndices('npt', self.getResiduesToFix('npt', self.fix_extra_npt), **self.selections_to_fix_npt)

        # fixed: bfactor=1; relaxed: bfactor=0
        self._psfFrag.occ[:]      = 0.0
        self._psfFrag.bfactors[:] = 0.0
        self._psfFrag.bfactors[masks_to_fix] = 1.0
        fixed_npt = self._run_kwargs.get('fixed_npt', '_fixed_npt.pdb')
        self._psfFrag.save(fixed_npt)

        # restore bfactor values
        self._psfFrag.bfactors[:] = 0.0

        # switch off constraints
        self.driver.constraints     = ''
        self.driver.constraints_ref = ''

        # restart from step 2
        self.driver.restart            = True
        self.driver.merge_restarted    = False # YL 08/08/2023: we must set this off which is defaulted to True
        self.driver.group_pressure     = True
        self.driver.flexible_cell      = False # False for water box & True for membrane
        self.driver.constant_area      = False # False for water box & True for membrane
        self.driver.fixed_atoms        = fixed_npt
        self.driver.fixed_atoms_forces = False
        self.driver.prefix_restart     = join(self.prefix_npt, self.driver.prefix)

        # YL 24/07/2023: create an input script of restarting the production MD for the user's convenience
        self.writeRestartScript()

        # run
        ierror = self.md.run(dryrun=dryrun, info_filename=f'{self.prefix_production}.npz')

        # rename folder by archiving
        self.driver.archive(dest=self.prefix_production)
        fileutils.archiveFile(fixed_npt     , dir=self.prefix_production, suffix='')
        fileutils.archiveFile(f'{self.prefix_production}.npz', dir=self.prefix_production, suffix='')
        print("\n {} ns NPT production MD simulation has run successfully. Results saved in:\n {}".format(self.md.timestep*self.md.nsteps/1000000, self.prefix_production))


    def writeRestartScript(self):
        '''Write out an input script with which the user can conveniently restart/continue the production MD'''

        from os.path import join

        # YL 24/07/2023: always write out an extra input script for the user to continue the production MD
        # YL 06/06/2024: added options to make the restarted job consistent with the previous one
        fileutils.archiveFile(self.restart_script)
        with open(self.restart_script, 'w') as fp:
            fp.write( '# Continue the production MD (NPT) for further simulation time\n\n')
            fp.write( 'from chemsh import *\n\n')
            fp.write( '# load the solvated fragment\n')
            fp.write(f"frag = Fragment(coords='{self.prefix_neutralised}.pdb')\n\n")
            fp.write(f'# create an instance of {self.driver.__class__.__name__}\n')
            strbuff = f'''{self.driver.__class__.__name__.lower()} = {self.driver.__class__.__name__}(frag=frag,
            restart=True,
            merge_restarted=True,
            prefix_restart='{join(self.prefix_production,self.driver.prefix)}',
            pdb='{self.driver.pdb}',
            psf='{self.driver.psf}',
            par='{join(self.prefix_production,self.driver._par)}',
            group_pressure={self.driver.group_pressure},
            flexible_cell={self.driver.flexible_cell},
            constant_area={self.driver.constant_area},
            fixed_atoms='{join(self.prefix_production,self.driver.fixed_atoms)}',
            fixed_atoms_forces={self.driver.fixed_atoms_forces},
            constraints='{self.driver.constraints}',
            constraints_ref='{self.driver.constraints_ref}',
            nsteps_per_cycle={self.driver.nsteps_per_cycle},
            freq_full_elect={self.driver.freq_full_elect},
           )\n\n'''
            fp.write(strbuff)
            if self.pbc_wrap:
                residues_to_relax = self.getResiduesToRelax('nvt', self.relax_extra_nvt)
                fp.write(f'indices_pbc_wrap = frag.select(residues={residues_to_relax}, strict=True, invert=True)\n')
                pbc_wrap = 'pbc_wrap=indices_pbc_wrap,'
            else:
                pbc_wrap = ''
            strbuff = f'''md = MD(theory={self.driver.__class__.__name__.lower()},
        driver='{self.md.driver}',
        ff='{self.md.ff}',
        boundary='{self.md.boundary}',
        langevin={self.md.langevin},
        langevin_piston={self.md.langevin_piston},
        langevin_damping={self.md.langevin_damping},
        rigid='{self.md.rigid}',
        timestep={self.md.timestep},
        nsteps={self.md.nsteps},
        minimise={self.md.minimise},
        nsnapshots={self.md.nsnapshots},
        freq_chk={self.md.freq_chk},
        freq_dcd={self.md.freq_dcd},
        freq_xst={self.md.freq_xst},
        freq_out_energy={self.md.freq_out_energy},
        freq_out_pressure={self.md.freq_out_pressure},
        {pbc_wrap}
       )\n\n'''
            fp.write(strbuff)
            fp.write( 'md.run(dryrun=False)\n\n')


    @property
    def _psfFrag(self):
        '''The fragment object created from PSFGEN result'''

        return self.__psfFrag


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

        self.__psfFrag = val


    @property
    def _solvatedFrag(self):
        '''The solvated fragment object'''

        return self.__solvatedFrag


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

        self.__solvatedFrag = val


    @property
    def md(self):
        '''MD task instance'''

        return self._md


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

        self._md = val


    @property
    def driver(self):
        '''Driver object of the MD instance'''

        return self._driver


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

        self._driver = val

    
    def _runBox(self):
        '''MD for a box model with ensembles'''

        from os.path import join
        from ..io    import format

        self.printHeader("Solvent type {} cube".format(self._run_kwargs.get('solvent', 'TIP3P')))

        # run the production MD (NPT) by restarting from previous NPT
        if str(self.restart_from_md).lower().strip() in ['2','production','3rd']:
            format.printwrap("Restarting from the previously run 1st (NVT) and 2nd (NPT) MD simulations...")
            self.initMD()
            # use the parameter file from the previous NPT (we don't run PSFGEN which supplements missing data in a restarted job)
            self.driver.par = join(self.prefix_npt, self.driver._par)
            self.runProduction()

        # run the 2nd MD (NPT) by restarting from previous NVT
        elif str(self.restart_from_md).lower().strip() in ['1','npt','2nd']:
            format.printwrap("Restarting from the previously run 1st (NVT) MD simulation...")
            self.initMD()
            # use the parameter file from the previous NVT (we don't run PSFGEN which supplements missing data in a restarted job)
            self.driver.par = join(self.prefix_nvt, self.driver._par)
            self.runNPT()
            self.runProduction()

        # run from the beginning
        else:

            # run PDB2PQR and solvate
            self._solvatedFrag = self.runPDB2PQR()

            # update title and inherit terminal
            self._solvatedFrag.terminal = self.solute.terminal

            # run PSFGEN first
            # the returned PDB and PSF will be used below by NAMD so that no pdb= and psf= are needed
            self.printHeader("running PSFGEN...")
            self.driver.frag = self._solvatedFrag
            self._psfFrag = self.driver.psfgen(return_frag=True,
                                               guess_coord=True,
                                               patch=self.driver.patch,
                                               **self.driver.psfgen_options)
            # YL 17/07/2021: after PSFGEN is run the sequence should NOT be changed any more!
            self._psfFrag.title = self._solvatedFrag.title
            # YL 26/11/2021: remember the renamings
            self._psfFrag._renamed_residues.update(self.driver.frag._renamed_residues)

            # neutralise
            self.runNeutralise()

            # run MD
            self.initMD()
            self.runNVT()
            self.runNPT()
            self.runProduction()


    def _runDroplet(self):
        '''MD for a droplet without ensemble'''

        for i in range(self.max_ncycles):
            print(" Solvation cycle {}".format(i))

            self._solvatedFrag = self.solute.solvate(**self._run_kwargs)
# YL TODO: to implement    
#            driver = mdCls(frag=self._solvatedFrag,
#                           psf='ubq_ws.psf',
#                           par='par_all27_prot_lipid.inp',
#                           pdb='_solvated.pdb',
#                           exclude='scaled1-4',
#                           path='namd2',
#                           binary=False,
#                          )
#           
#            mdObj = MD(driver=driver,
#                       freq_dcd=1,
#                       freq_chk=1,
#                       minimise=1,
#                       nsteps=10,
#                       nsnapshots=0,
#                       timestep=2.0,
#                       temperature=310.0,
#                       boundary='spherical',
#                      )
#           
#            mdObj.run(dryrun=False)


    @miscutils.memusage
    def run(self, **kwargs):
        ''''''

        from ..io.banner  import printModuleBanner

        printModuleBanner('ChemShell Solvation Workflow')

        # all keyword arguments at the self.run() step will be temporary
        # NOTE: do not use `__dict__` but only dir(), as all attibutes of <property> are not contained in the latter
        for key in set(kwargs) & set(dir(self)):
            # swap values (syntax a,b=c,d does not apply)
            buffer = getattr(self, key)
            setattr(self, key, kwargs[key])
            kwargs[key] = buffer
  
        if self._run_kwargs.get('shape', 'cube').lower() in [ 'box', 'cube' ]:
            self._runBox()

        elif self._run_kwargs.get('shape', 'cube').lower() in [ 'droplet', 'sphere' ]:
            self._runDroplet()

        else:
            print(" >>> ChemShell ERROR: unknown solvent shape \"{}\"\n".format(self._run_kwargs['shape']))
            exit(999)


        # swap back run arguments
        for key in set(kwargs) & set(dir(self)):
            # swap values back (syntax a,b=c,d does not apply)
            buffer = getattr(self, key)
            setattr(self, key, kwargs[key])
            kwargs[key] = buffer

        self.printHeader('Solvation workflow completed!')

        return self


#    @property
#    def residues_to_relax_nvt(self):
#        '''A names list of residues to be relaxed for NPT simulations (2nd and production MD)'''
#
#        from ..data import forcefield
#
#        # get the forcefield object
#        ffObj = forcefield.getFFObj(self.ff)
#
#        lbuff = []
#
#        for item in self.relax_residues_nvt:
#            item = forcefield.synonyms.get(item.lower().strip(), item.lower().strip())
#            list_resnames = getattr(ffObj, item, [])
#            lbuff += list_resnames
#
#        return lbuff


    @property
    def selections_to_relax_nvt(self):
        '''Return selecting options for NVT if they are a dictionary'''

        if type(self.relax_extra_nvt) is dict:
            return self.relax_extra_nvt
        else:
            return {}


    @property
    def selections_to_relax_npt(self):
        '''Return selecting options for NPT if they are a dictionary'''

        if type(self.relax_extra_npt) is dict:
            return self.relax_extra_npt
        else:
            return {}


    @property
    def selections_to_fix_nvt(self):
        '''Return selecting options for NVT if they are a dictionary'''

        if type(self.fix_extra_nvt) is dict:
            return self.fix_extra_nvt
        else:
            return {}


    @property
    def selections_to_fix_npt(self):
        '''Return selecting options for NPT if they are a dictionary'''

        if type(self.fix_extra_npt) is dict:
            return self.fix_extra_npt
        else:
            return {}


    # YL 05/07/2023: beware that this function only obtains residue names from the extra inputt `fix_extra`;
    #                flexible selections passed in by a dict type parameter will be handled by _getExtraFixedIndices()
    def getResiduesToRelax(self, md, relax_extra=[]):
        '''Return a list of residue names to be relaxed for NPT simulations (2nd and production MD)'''

        from ..data  import forcefield

        # get the forcefield object
        ffObj = forcefield.getFFObj(self.ff)

        lbuff = []

        if md == 'nvt':
            relax_extra_self    = self.relax_extra_nvt
            relax_residues_self = self.relax_residues_nvt
        else:
            relax_extra_self    = self.relax_extra_npt
            relax_residues_self = self.relax_residues_npt

        # WARNING: iterutils.getFlattened() CANNOT directly deal with a single string
        if type(relax_extra) in [ str, bytes ]:
            relax_extra = [relax_extra]

        # check synonyms
        for item in relax_residues_self:
            item = forcefield.synonyms.get(item.lower().strip(), item.lower().strip())
            list_resnames = getattr(ffObj, item, [])
            lbuff += list_resnames

        # YL 26/11/2021: allow the user to add extra residues to relax
        # YL 14/12/2021: more flexible selection (see selections_to_relax_nxt())
        try:
            extra = iterutils.getFlattened(relax_extra) or iterutils.getFlattened(relax_extra_self)
            # the default blank value ([],{}) will be converted to ({},)
            if extra == ({},):
                extra = []
        except:
            extra = []
        # NB: this function is only for list type of `extra`
        if type(extra) is dict:
            extra = []
        for item in extra:
            # allow residue name from the original or automatically renamed one (e.g., either 'OXY' or 'O2' reduces to 'O2' if such a renaming exists)
            lbuff += [ self._psfFrag._renamed_residues.get(item, item) ]

        return lbuff


    # YL 05/07/2023: beware that this function only obtains residue names from the extra inputt `fix_extra`;
    #                flexible selections passed in by a dict type parameter will be handled by _getExtraFixedIndices()
    def getResiduesToFix(self, md, fix_extra=[]):
        '''A list of residue names to be fixed for NPT simulations (2nd and production MD)'''

        from ..data  import forcefield

        # get the forcefield object
        ffObj = forcefield.getFFObj(self.ff)

        lbuff = []

        if md == 'nvt':
            fix_extra_self    = self.fix_extra_nvt
        else:
            fix_extra_self    = self.fix_extra_npt

        # WARNING: iterutils.getFlattened() CANNOT directly deal with a single string
        if type(fix_extra) in [ str, bytes ]:
            fix_extra = [fix_extra]

        # YL 14/12/2021: more flexible selection (see selections_to_fix_nxt())
        try:
            extra = iterutils.getFlattened(fix_extra) or iterutils.getFlattened(fix_extra_self)
            # the default blank value ([],{}) will be converted to ({},)
            if extra == ({},):
                extra = []
        except:
            extra = []
        # NB: this function is only for list type of `extra`
        if type(extra) is dict:
            extra = []
        for item in extra:
            # allow residue name from the original or automatically renamed one (e.g., either 'OXY' or 'O2' reduces to 'O2' if such a renaming exists)
            lbuff += [ self._psfFrag._renamed_residues.get(item, item) ]

        return lbuff


    # YL 15/12/2021: _residues= for list of residues and **kwargs for free selection
    def _getFixedIndices(self, md, _residues=[], _return_masks=True, **kwargs):
        '''Return masks/indices of residues to be fixed'''

        from numpy import arange, char, unique
        from ..io  import format

        # YL 15/12/2021: retrieve residues and remove from dict of keyword args
        residues_in_dict  = kwargs.pop('residues', [])
        # YL 15/12/2021: free selection
        masks_to_fix      = self._psfFrag.select(residues=self.getResiduesToRelax(md, residues_in_dict),
                                                 union=True,        # NB: must use union=True to make all keywords effective because invert=True!
                                                 invert=True,
                                                 return_masks=True,
                                                 **kwargs)
        masks_to_fix     += self._psfFrag.select(residues=_residues, strict=True, invert=True, return_masks=True)
        resnames_to_relax = unique(self._psfFrag.residues[~masks_to_fix])
        resnames_to_fix   = unique(self._psfFrag.residues[masks_to_fix])
        print("\n Names of relaxed residues:")
        format.printwrap(' '.join(char.decode(char.strip(resnames_to_relax))), prepend_blanc_line=False)
        print("\n Names of fixed residues:")
        format.printwrap(' '.join(char.decode(char.strip(resnames_to_fix))), prepend_blanc_line=False)
        if _return_masks:
            return masks_to_fix
        else:
            return arange(self._psfFrag.natoms)[masks_to_fix]


    # YL 15/12/2021: _residues= for list of residues and **kwargs for free selection
    def _getExtraFixedIndices(self, md, _residues=[], _return_masks=True, **kwargs):
        '''Return masks/indices of residues to be fixed'''

        from numpy import arange, char, unique
        from ..io  import format

        # YL 15/12/2021: retrieve residues and remove from dict of keyword args
        residues_in_dict = kwargs.pop('residues', [])
        # YL 15/12/2021: free selection
        masks_to_fix     = self._psfFrag.select(residues=self.getResiduesToFix(md, residues_in_dict),
                                                invert=False,
                                                return_masks=True,
                                                **kwargs)
        if masks_to_fix.sum():
            frag = self._psfFrag
            indices_to_fix = arange(self._psfFrag.natoms)[masks_to_fix]
            resnames_to_fix   = unique(self._psfFrag.residues[masks_to_fix])
            print('\n Extra atoms to fix:')
            print(' {:>7}  {:4}  {:7}  {:5}  {:7}'.format('index', 'name', 'residue', 'resID', 'segment'))
            for i in indices_to_fix:
                print(' {:7}  {:4}  {:7}  {:5}  {:7}'.format(i, frag.names[i].decode(), frag.residues[i].decode(), frag.resIDs[i], frag.segments[i].decode()))
        masks_to_fix    += self._psfFrag.select(residues=_residues, strict=True, invert=False, return_masks=True)
        if len(_residues):
            resnames_to_fix   = unique(self._psfFrag.residues[masks_to_fix])
            print("\n Names of extra fixed residues:")
            format.printwrap(' '.join(char.decode(char.strip(resnames_to_fix))), prepend_blanc_line=False)
        if _return_masks:
            return masks_to_fix
        else:
            return arange(self._psfFrag.natoms)[masks_to_fix]

