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

# this is an interface to C. Yong's DL_FIELD
# see: http://www.ccp5.ac.uk/DL_FIELD
#      DOI: 10.1021/acs.jcim.6b00323

# you.lu@stfc.ac.uk
# May 2018


import chemsh
from sys              import modules, path
from os.path          import dirname, join
from ..driver         import _Driver
from ...io            import warnings
from ...dl_py2f.utils import fileutils, modutils

class DL_MONTE(_Driver):
    '''Interface to DL_MONTE'''

    # YL: get the HTK path inscribed at compiletime
    try:
        with open(join(dirname(__file__), 'htk_path'), 'r') as fp:
            htk_path = fp.read()
        try:
            # YL 09/07/2019: enabled importing of DL_MONTE HTK as a module 
            path.append(htk_path.strip())
            import dlmontepython.htk.sources.dlmonte as dlmonte
            print("\n >>> DL_MONTE HTK library found: {}\n".format(htk_path))
        except:
            warnings.warn('unavailable', 'DL_MONTE HTK')

    # no HTK at compiletime
    except FileNotFoundError:
        pass

    

    _kwargs = {
                'title'                 : None,
                'use'                   : [],
                'temperature'           : 0.0,
                'steps'                 : 0,
                'print'                 : 100,
                'stack'                 : 100,
                'equilibration'         : 100000,
                'noewald'               : None,
                'sample'                : None,
                'sample_frequency'      : 10000,
                'revconformat'          : None,
 
                'maxatmdist'            : 0.001,
                'maxvolchange'          : 0.01,
                'acceptatmmoveupdate'   : 100,
                'acceptvolupdate'       : 100,
 
                'move_atom'              : [],
                'move_atom_frequency'     : 100,
                'move_volume'            : None,
                'move_vector'            : None,
                'move_volume_frequency'   : 20,
                'move_gcinsertatom'     : [],
                'move_gcinsertatom_frequency' : 20,
                'move_gcinsertatom_min_dist' : 1.0,
 
                'ff'                    :'',
                'input'                 :('CONTROL','CONFIG','FIELD'),
                'pbc'                   : 3,      #defult periodic boundary conditions
 
                #if you just want to pass a string for the control, config & field use commands below
                'control'               : None,
                'config'                : None,
                'field'                 : None,
 
                #log
                'output'                :'dl_monte.out',
                
              }


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

        from numpy       import 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 os.path     import dirname, join
        from ...data     import unitconvert
        from ...parallel import parallel
        #from ...objects import fragment as fr

        # YL TODO 09/07/2019: allow users to overload the compiletime HTK location with a runtime one
#        if self.htk:
#            path.append(self.htk)
#            from htk.sources import dlmonte
#            try:
#                from htk.sources import dlmonte
#            except:
#                print(' >>> ChemShell WARNING: invalid HTK location given in input script\n htk="{}"'.format(self.htk))
#                path.pop()

        # YL TODELETE: example of using dlfield functions
#        field = dlmonte.dlfield.FIELD()



        # CONTROL file
        if input == self.input[0]:

            if self.control is not None:

                if isinstance(self.control, self.dlmonte.dlcontrol.CONTROL) == True:
                    strbuff = self.control.to_string()
                else:
                    strbuff = fileutils.file2text(self.control, prefix=parallel.prefix) + '\n'

            else:

                #first create simple title, use block, main block
                if self.title is None:
                    self.title = 'ChemShell DL_MONTE'

                strbuff = self.title + '\n'

                if len(self.use) > 0:

                    for i in range(len(self.use)):

                        dummystr = 'use {}\n'.format(self.use[i])
                        strbuff += dummystr

                    strbuff += 'finish use-block\n'   #there is no useblock processing for now

                strbuff += 'steps {}\n'.format(self.steps)
                strbuff += 'temperature {}\n'.format(self.temperature)
                strbuff += 'equilibration {}\n'.format(self.equilibration)
                strbuff += 'print {}\n'.format(self.print)
                strbuff += 'stack {}\n'.format(self.stack)

                if self.noewald is not None:
                    strbuff += 'noewald {}\n'.format(self.noewald)

                if self.sample is not None:
                    strbuff += 'sample {} {}\n'.format(self.sample, self.sample_frequency)

                if self.revconformat is not None:
                    strbuff += 'revconformat {}\n'.format(self.revconformat)

                #process the moves
                if len(self.move_atom) > 0:         #i.e. moving atoms is required

                    strbuff += 'maxatmdist {}\n'.format(self.maxatmdist)
                    strbuff += 'acceptatmmoveupdate {}\n'.format(self.acceptatmmoveupdate)
                    strbuff += 'move atom {} {}\n'.format(len(self.move_atom), self.move_atom_frequency)

                    for i in range(len(self.move_atom)):

                        dummystr = self.move_atom[i] + '\n'
                        strbuff += dummystr

                if self.move_volume is not None or self.move_vector is not None:

                    strbuff += 'maxvolchange {}\n'.format(self.maxvolchange)
                    strbuff += 'acceptvolupdate {}\n'.format(self.acceptvolupdate)

                if self.move_volume is not None:
                    strbuff += 'move volume {} {}\n'.format(self.move_volume, self.move_volume_frequency)

                if self.move_vector is not None:
                    strbuff += 'move volume vector {} \n'.format(self.move_volume_frequency)

                if len(self.move_gcinsertatom) > 0:

                    strbuff += 'move gcinsertatom {} {} {}\n'.format(len(self.move_gcinsertatom), self.move_gcinsertatom_frequency, self.move_gcinsertatom_min_dist)

                    for i in range(len(self.move_gcinsertatom)):

                        dummystr = self.move_gcinsertatom[i] + '\n'
                        strbuff += dummystr

                strbuff += 'start simulation\n'

        # CONFIG file
        elif input == self.input[1]:
            ''' 
            The CONFIG can be passed as either a string or fragment (NB this is not sufficient except for the
            simplest of systems). The easiest way for more complex systems is to create am instance of config
            in the calling scipt and pass the formatted string.
            '''
            if self.config is not None:

                if isinstance(self.config, self.dlmonte.dlconfig.CONFIG) == True:

                    strbuff = self.config.to_string()

                else:

                    strbuff = self.config

            else:                       #passed a fragment

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

                # record 2: LEVCFG and IMCON
                if self.pbc > -1:
                    strbuff += '0 {}\n'.format(self.pbc)
                else:
                    strbuff += '0 {}\n'.format(self.frag.cell.imcon)

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

#            strbuff += 'NUMMOLS {} {}\n'.format(self.nmols, [ len(i) for i in self.moltypes ])
                strbuff += 'NUMMOLS 1 1\n'
                strbuff += 'MOLECULE {} {} {}\n'.format(self.frag.title, self.frag.natoms, self.frag.natoms)

                # core atoms
                strbuff += self.frag.coords2str(atoms='names', coords=6, charges=None, delimiter=' ', unit='angstrom', inserts=[(1, 'core\n')])

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

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

            if self.field is not None:

                if isinstance(self.field, self.dlmonte.dlfield.FIELD) == True:

                    strbuff = self.field.to_string()

                else:

                    strbuff = self.field

            else:

                from ...io import file
    
                # receiver has to be a mutable type, such as bytearray
                strbuff = bytearray(b'')

                file.parse(self.ff, fmt='ff', receiver=strbuff)

                # back to str (no leading and trailing blank lines allowed by DL_MONTE)
                strbuff = strbuff.decode().lstrip().rstrip()

        return strbuff


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

        from ctypes     import c_int, byref, CDLL
        from _ctypes    import dlclose
        from ...dl_py2f import py2f
        from chemsh     import CHEMSH_LIB_DIR
        from ...data    import unitconvert

        print("\n >>> Running DL_MONTE simulation using directly linked-in library...\n")

        # load DL_MONTE library
        libdl_monte = modutils.getLinkedLib(self.__class__.__name__, chemsh.CHEMSH_LIB_DIR)

        libdl_monte.interface_dl_monte(byref(py2f(self.frag)),
                                       byref(py2f(self._result)))

        dlclose(libdl_monte._handle)


    def parseMainOutput(self):
        '''Parse result information out from output'''

        # YL placeholder



