#  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 ..mm        import _MM
from ...dl_py2f  import py2f, utils
from ...base     import constants, errors
from ...data     import atomicdata
from ...parallel import parallel


class GULP(_MM):
    '''Interface to GULP'''

    _kwargs = {
                'cartesian'   : True,
                'c6'          : False,
                'input'       :'_gulp.gin',
                'molecule'    : False,
                'name'        : __name__,
                'ndimensions' : 3,
                'output'      :'_gulp.gout',
                'runargs'     :'_gulp',
                'potential'   : False,
                'conjugate'   : True,        # use conjugate gradients minimiser
              }

    _internals = {
                   '_qmmm_scheme'   :  None,
                   '_qmmm_embedding':  None,
                   '_syscmd'        : 'gulp',
                   '_extra_outputs' :['gulp.gradient','gulp.relshel','_gulp.frc'], # must be a list
                   '_cell_redundant':  False,
                 }

    def getInpBuff(self, task='qmmm', **kwargs):
        ''''''

        from numpy      import array2string, 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

        flag_frozen = "0 0 0"
        flag_active = "1 1 1"

        # to optimise the lattice constants only
        cell_only = task.strip().lower() in [ 'cell-only', 'cell_only' ]

        # use a reference for all actions to avoid time-consuming access, especially in loops!
        shells = self.frag.shells

        # YL 13/02/2019: rewritten the logics to cover pure MM calculations with shells
        #                shells    grads    to relax shells      not to relax shells
        #                O         O        "opti comp"          "energy gradient"
        #                O         X        "opti comp"          "energy single"
        #                X         O        "energy gradient"    "energy gradient"
        #                X         X        "energy single"      "energy single"
        # QM/MM w/ shells: switch on shell relaxation
        atom_flags = flag_frozen
        # when no shells
        if not self.frag.nshells:
            self.relax_shells = False
        # whether QM/MM or MM
        # YL 19/11/2019: "noaddshells" should always be used or GULP results in nshells crisis
        if self.relax_shells:
            strbuff     = ''
            opti_flags  = 'opti comp noaddshells '
            shell_flags = flag_active
        else:
            if cell_only:
                self._extra_outputs += [ '_gulp.frc' ]
                strbuff = 'gradient cellonly stress_out\n'
                strbuff += 'output frc _gulp.frc'
            else:
                strbuff = 'energy noaddshells '
                shell_flags = flag_frozen
                # gradients requested
                if self._gradients:
                    opti_flags = 'gradient '
                    atom_flags = flag_active
                # no gradients requested
                else:
                    opti_flags = 'single '

        # note the trailing space
        if self.molecule and self.frag.cell.natoms:
            strbuff += 'molecule '

        if not cell_only:
            strbuff += opti_flags

            if self.conjugate:
                strbuff += 'conj '
            strbuff += 'angle torsion bond '

        # no partial occupancy constraints!
        if self._qmmm_scheme == "qm/me":
            strbuff += 'noco '

        strbuff += '\ntitle\n'
        if self.frag.title.strip() in [ 'title', b'title', 'b\'title\'' ]:
            titlebuff = self.frag.formula
        else:
            titlebuff = self.frag.title.strip()
        strbuff += '    Py-ChemShell/GULP: {}\n'.format(titlebuff)
        strbuff += 'end\n'

        if task == 'charges':
            strbuff += "chemshell_mode init\n"
#        elif task.strip().lower() in [ 'cell-only', 'cell_only' ]:
#            pass
        else:
            strbuff += "chemshell_mode calc\n"

        if self.frag.cell.natoms:
            # 2D surface for QM/Me scheme
            # NB: vec[x/y/z, a/b/c]
            # either svectors or scell is OK
            vectors = self.frag.cell.vectors#*unitconvert.Bohr2Angstrom
            if self._qmmm_scheme == "qm/me" or ( self.ndimensions == 2 and cell_only):
                strbuff += 'svectors au\n'
                strbuff += '{:16.9f} {:16.9f}\n'.format(vectors[0,0], vectors[1,0])
                strbuff += '{:16.9f} {:16.9f}\n'.format(vectors[0,1], vectors[1,1])
                strbuff += '0 0 0 0\n'
#                strbuff += 'scell au\n'
#                strbuff += '{:16.9f} {:16.9f} {:16.9f}\n'.format(self.frag.cell.a, self.frag.cell.b, self.frag.cell.alpha)
            # YL 24/11/2020: is this XYZ/ABC correct???
            else:
                strbuff += 'vectors au\n'
                strbuff += '%16.9f %16.9f %16.9f\n' % (vectors[0,0], vectors[0,1], vectors[0,2])
                strbuff += '%16.9f %16.9f %16.9f\n' % (vectors[1,0], vectors[1,1], vectors[1,2])
                strbuff += '%16.9f %16.9f %16.9f\n' % (vectors[2,0], vectors[2,1], vectors[2,2])
                strbuff += '0 0 0 0 0 0\n'
                # Who? should be no need to use fractional coordinates even if they exist

# FIXME: can use just names
        if self.frag.names[0]:
            # Atom names exist
            namemode = 'names' 
        else:
            # Generate atom names
            namemode = 'symbols'

        # YL TODO: consider more than one QM or MM regions
        # FIXME GULP fails if the coordinates lines are too wide
        if self.qmatoms and self._qmmm_scheme != 'qm/me':
            strbuff += 'cartesian region 1 qm\n'
            strbuff += self.frag.coords2str(select=self.qmatoms, atoms=namemode, coords=10, charges=8, inserts=[(1,'c'), (6,atom_flags)], delimiter=' ', unit='angstrom')
            strbuff += 'cartesian region 2 mm nonrigid\n'
            strbuff += self.frag.coords2str(select=self.mmatoms, atoms=namemode, coords=10, charges=8, inserts=[(1,'c'), (6,atom_flags)], delimiter=' ', unit='angstrom')
            # charges-only mode
            if task == 'charges':
                 # TWK we can't set charges=None here as shell charges that are already set should be respected 
                strbuff += shells.coords2str(atoms=namemode, coords=10, charges=8, inserts=[(1,'s'), (6,shell_flags)], delimiter=' ', unit='angstrom')
            # normal mode
            else:
                # get an indices list of "region 4" atoms
                self.shells2freeze = shells.getRegion(4)
                # make a new list of all active ("1 1 1") flags
                shell_flags = full((self.frag.nshells,), shell_flags, dtype='S5')
                # freeze the "region 4" atoms
                shell_flags[self.shells2freeze] = flag_frozen
                strbuff += shells.coords2str(atoms=namemode, coords=10, charges=8, inserts=[(1,'s'), (6,shell_flags)], delimiter=' ', unit='angstrom')
                strbuff += 'qmmm ' + self._qmmm_embedding + '\n'
                # YL: changed `self._qmmm_embedding` as internal, so there should not be a chance to receive unrecognised values
                #     if this is not the case, use: raise errors.ChemShValError(key, selectcases.keys())

                # Forces on shells from QM calculations
                # YL 12/12/2017: removed self.frag.shells which is pretty slow. now it runs reasonably fast but may still need improvement by replacing for loops
                if(self.frag.nshells > 0):
                    if chemsh.CHEMSH_DEBUG > 3:
                        print("\n ChemShell/GULP: 1 a.u. = 51.422606 eV/Angstrom\n")
                    strbuff += "external_force\n"
                    for i in range(self.frag.nshells):
                        n = self.frag.natoms + i + 1
                        gradients = list(shells.gradients[i])
                        # conversion factor from a.u. to eV/Angstrom taken from GULP (old chemshellenergy/initdervs routines)
                        # NB: sign change needed because external_force requires a genuine "force"
                        forces_external = [ j * -51.422606 for j in gradients ]
                        strbuff += "{:5d} {:20.14f} {:20.14f} {:20.14f}\n".format(n, forces_external[0], forces_external[1], forces_external[2])

        else:
            # use section "sfractional" for QM/Me periodic embedding
            # not in use (equivalent to Cartesian but this needs cart2frac() every time)
#            if self._qmmm_scheme == 'qm/me':
            if not self.cartesian or cell_only:
                frac_x = self.frag.cell.fractional[:,0]
                frac_y = self.frag.cell.fractional[:,1]
                if self.ndimensions == 2:
                    strbuff += 'sfractional\n'
                    cart_z = self.frag.coords[:,2]*unitconvert.Bohr2Angstrom
                    strbuff += self.frag.coords2str(atoms     =  namemode,
                                                    coords    =  None,
                                                    charges   =  8,
                                                    inserts   = [(1,'c'),
                                                                 (2,frac_x,'%16.10f'),
                                                                 (3,frac_y,'%16.10f'),
                                                                 (4,cart_z,'%16.10f'),
                                                                 (6,'1 1 1')           # flags are necessary!
                                                                ],
                                                    delimiter = ' ',
                                                    unit      = 'angstrom',
                                                    order     = 'aq',       # this option can switch off coordinates and charges printing
                                                   )
                elif self.ndimensions == 2:
                    strbuff += 'fractional\n'
                    frac_z = self.frag.cell.fractional[:,2]
                    strbuff += self.frag.coords2str(atoms     =  namemode,
                                                    coords    =  None,
                                                    charges   =  8,
                                                    inserts   = [(1,'c'),
                                                                 (2,frac_x,'%16.10f'),
                                                                 (3,frac_y,'%16.10f'),
                                                                 (4,frac_z,'%16.10f'),
                                                                 (6,'1 1 1')           # flags are necessary!
                                                                ],
                                                    delimiter = ' ',
                                                    unit      = 'angstrom',
                                                    order     = 'aq',       # this option can switch off coordinates and charges printing
                                                   )
                strbuff += '\n'

            else:
                # YL 12/02/2019: also freeze the shells if the associated cores are frozen
                # YL 20/05/2020: no longer true unless 
                shell_flags = full((self.frag.nshells,), shell_flags, dtype='S5')
                shell_flags[self.shells2freeze] = flag_frozen

                strbuff += 'cartesian\n'
                strbuff += self.frag.coords2str(atoms=namemode, coords=10, charges=8, inserts=[(1,'c'), (6,atom_flags)], delimiter=' ', unit='angstrom')
                strbuff += shells.coords2str(atoms=namemode, coords=10, charges=8, inserts=[(1,'s'), (6,shell_flags)], delimiter=' ', unit='angstrom')

        # GULP forcefield definition
        try:
            # return file content
            ffbuff = utils.fileutils.file2text(self.ff, error_return=False, prefix=parallel.prefix)
            strbuff += ffbuff + '\n'
        except FileNotFoundError:
            # return the string's content itself (built-in library FF)
            ffbuff = utils.fileutils.file2text(self.ff, prefix=parallel.prefix)
            if len(ffbuff.split()) > 1:
                strbuff += ffbuff + '\n'
            elif len(ffbuff.split()) == 1:
                strbuff += 'library nodump {}\n'.format(ffbuff)

        return strbuff


    # YL 13/02/2019: allow freezing/activating cores and shells at the same time or respectively
    #                usage: many situations have been considered (see dev/sandpit/gulp/h2o_dimer_shells_freeze_active.py)
    @property
    def shells2freeze(self):
        '''A list of shells to freeze (not to relax)'''

        # shellatoms not available unless `self.frag.shells` is initialised
        self.frag.shells

        # when totally switched off: all shells are frozen
        if not self.relax_shells:
            self._shells2freeze = self.frag.shellatoms[self.frag.shellatoms!=-1]
            return self._shells2freeze
    
        # update from frag.frozen
        def _setShells2Freeze(_self):

            from numpy import array

            # if not yet defined: shellatoms of frozen cores
            lbuff = _self.frag.shellatoms[_self.frag.frozen==-1]
            # exclude no-shell cores
            _self._shells2freeze = lbuff[lbuff!=-1]

            # YL 14/05/2020: by default, still relax the shells associated with frozen cores, to be consistent with single-point calculations
            if _self.relax_shells_frozen_cores:
                try:
                    return _self._shells2freeze_val
                # YL 20/05/2020: in case self.shells2freeze not defined anywhere
                except:
                    return array([], dtype=int)
            else:
                return _self._shells2freeze

        # already defined via setter
        if hasattr(self, '_shells2freeze_val'):
            # no change of frag.frozen
            if set(self.__shells2freeze_saved) == set(_setShells2Freeze(self)):
                return self._shells2freeze_val
            else:
                return _setShells2Freeze(self)
        # update from frag.frozen if not defined via setter 
        else:
            return _setShells2Freeze(self)


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

        from numpy import asarray

        self._shells2freeze_val = asarray(val)
        # remember the current list and reset it
        self._shells2freeze, self.__shells2freeze_saved = asarray([]), getattr(self, '_shells2freeze', asarray([]))


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


    def parseAdditionalOutput(self):
        '''Overload mm.MM.parseAdditionalOutput()'''

        from ...io import file

        # filenames to be parsed and receivers of data
        # units: gulp.energy    - a.u.
        #        gulp.gradients - a.u.
        #        gulp.relshel   - a.u.
        selectcases = { '_energy'   : 'gulp.energy',
                        '_gradients': 'gulp.gradient',
                        '_shells'   : 'gulp.relshel',
                        '_strains'  : '_gulp.frc',
                      }

# YL: FIXME: move out of interface
        # concatenate grads and BQ gradients
        for task, tpl in selectcases.items():
            # YL 25/01/2021: improved the if statement
            if not getattr(self, task, False):
                continue
            # file "gulp.gradient" follows the order of given coords, namely first QM then MM (see getInpBuff)
            if task == '_shells':
                self.frag.shells.coords = file.getArrayFromFile(selectcases[task], 'punch', self.frag.shells.coords.shape)
            elif task == '_energy':
                self._result.energy = file.getArrayFromFile(selectcases[task], 'punch', (1,1))[0]
                # optimisation of cell parameters
            elif task == '_gradients':
                if self._qmmm_scheme == "qm/me":
                    self._result.gradients[:] = 0.0
                    self._result.gradients[self._non_adsorbate] = file.getArrayFromFile(selectcases[task], 'punch', self._result.gradients.shape)
                else:
                    # YL 11/08/2019: a serious bug found that the output gradients of GULP
                    #                are in order of QM and MM but the receiver self._result.gradients
                    #                was in the original order of the whole fragment
                    self._result.gradients[self._qmmm_order] = file.getArrayFromFile(selectcases[task], 'punch', self._result.gradients.shape)
            # do not enquire the data receiver directly through a dict value, otherwise the dict rather than the receiver is redefined!
            # e.g.: d = { 'task':frag.coords }
            #       d['task'] = val
#            selectcases[task][1] = file.getArrayFromFile(selectcases[task][0], selectcases[task][1].shape, 'punch')
            elif task == '_strains':
                if self.ndimensions == 2:
                    self._result.derivatives = file.getArrayFromFile(tpl, 'gulp', (3,))
                elif self.ndimensions == 3:
                    self._result.derivatives = file.getArrayFromFile(tpl, 'gulp', (6,))

    def runLib(self, options=None, **kwargs):
        '''Run GULP as a shared library (called by run.runLib())'''

        from resource import getrusage, RUSAGE_SELF
        from time     import time
        from ctypes   import byref
        from os       import environ
        from sys      import stdout
        from .gulp_libs_path import GULP_LIBS_PATH

        environ['GULP_LIB'] = GULP_LIBS_PATH

        # load GULP library
        libgulp = utils.modutils.getLinkedLib(self.__class__.__name__, chemsh.CHEMSH_LIB_DIR)

        #print("\n     >>> memory used: %.3f MB\n"%(getrusage(RUSAGE_SELF).ru_maxrss/1024))
        ss = time()
        if libgulp.interface_gulp(byref(py2f(self.frag))):
            stdout.flush()
            raise errors.ChemShModuleError('GULP')

        #print("\n    ### time used for GULP:", time()-ss)
        #print("\n     <<< memory used: %.3f MB\n"%(getrusage(RUSAGE_SELF).ru_maxrss/1024))


    def processErr(self, filename):
        '''Postrun process of error output'''

        with open(filename) as fp:
            sbuff = fp.read()
            if 'STOP GULP terminated with an error' in sbuff:
                print("\n ChemShell/GULP ERROR:\n {}".format(sbuff))
                return 999
            else:
                return 0



