#  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 .               import _Interfaces
from ..base          import run
from ..io            import warnings
from ..dl_py2f.utils import fileutils, miscutils, modutils

# inherit <ctypes.Structure> type for convenience of making callback
class _MM(_Interfaces):
    '''A container of MM methods: all MM interfaces should be of this class'''

    _kwargs = {
                'atom_names'               :'names',
                'ensemble'                 :'npt',
                'ff'                       :'',
                'logger'                   :'',      # error logger for stderr
                'mmatoms'                  :[],
                'path'                     :'',
                'parameterise'             : False,
                'periodic'                 : None,
                'pressure'                 : 0.001,
                'qmatoms'                  :[],
                'runargs'                  :'',
                'sysmpi'                   :'',
                'temperature'              : 293.15,
                'relax_shells'             : True,   # relax shells even for a single-point energy evaluation
                'relax_shells_frozen_cores': True,   # YL 14/05/2020: relax shells of frozen cores, too
              }

    _internals = {
                   '_shells'       : False,
                   '_syscmd'       : '',
                   '_extra_outputs': [],
                   '_recalc_cell'  : True,
                 }


    # set up at the QMMM level
    @property
    def qmatoms(self):
        '''List of QM atoms'''

        return self._qmatoms


    @property
    def mmatoms(self):
        '''List of MM atoms'''

        return self._mmatoms


    # called by run.runExec() or run.runLib()
    def genInput(self, **kwargs):
        '''Generate input'''

        # single input file
        if type(self.input) in [ str ]:
            fileutils.write(self.input, self.getInpBuff(**kwargs))

        # multiple input files
        elif type(self.input) in [ list, tuple ]:
            for input in self.input:
                fileutils.write(input, self.getInpBuff(input=input, **kwargs))


    def parseMainOutput(self):
        '''Parse main output file self.output'''

        from ..io import file

        self.setOutKeys()

        self._result.energy    = file.getArrayFromFile(self.output, 'out', self._result.energy.shape   , regex=self._outkeys['energy'])
        self._result.gradients = file.getArrayFromFile(self.output, 'out', self._result.gradients.shape, regex=self._outkeys['gradients'])


    def parseAdditionalOutput(self):
        '''Parse additional output data, e.g., nwchem.bq_grads. This procedure could be overloaded, otherwise ignored.'''


    def setOutKeys(self):
        '''Set up {task:keyword} for output parsing. Shall be overloaded if needed.'''


    def parseOutput(self):
        '''Parse output'''

        self.parseMainOutput()

        self.parseAdditionalOutput()


    @miscutils.time
    def run(self, **kwargs):
        '''Single-point procedure'''

        from sys     import stdout
        from ..io    import banner, format
        from ..data  import references

        if not self.frag.natoms:
            print("\n\n >>> ChemShell ERROR: no valid fragment has been assigned to theory {}!\n\n".format(self.__class__.__name__))
            exit(999)

        fileutils.archiveFile(self.input)
        fileutils.archiveFile(self.output)

        banner.printModuleBanner(f'ChemShell/{self.__class__.__name__} Interface', level=2)
#        if self.__class__.__name__.lower() in references.references:
#            print(' Please cite:')
#            format.printwrap(references.references[self.__class__.__name__.lower()], prepend_blanc_line=False)

        # YL: we allow users to override linked-in libs with `path`: check path first then isLinked then _syscmd
        cmdpath = self.path or modutils.getLinkedLib(self.__class__.__name__, chemsh.CHEMSH_LIB_DIR) or self._syscmd
        if chemsh.CHEMSH_DEBUG > 2:
            print("\n >>> _MM.run(): cmdpath has been determined to be:", cmdpath)
            print("                by path       = {}".format(self.path))
            print("                by linked lib = {}".format(modutils.getLinkedLib(self.__class__.__name__, chemsh.CHEMSH_LIB_DIR)))
            print("                by _syscmd    = {}".format(self._syscmd))

        # don't use higher level task routines here (need to keep lower level separate)
        # all keyword arguments at the run() step will be temporary
        # QM/Me scheme
        if self._qmmm_scheme == "qm/me":
            ierror = self.__runQMMe(cmdpath, **kwargs)
        # standard methods
        else:
            ierror = run.run(self, cmdpath, stderr=self.logger, **kwargs)

        return ierror


    @property
    def logger(self):
        '''Error logger filename'''

        if not self._logger:
            return '_{}.err'.format(self.__class__.__name__.lower())
        else:
            return self._logger


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

        self._logger = val


    def __runQMMe(self, cmdpath, **kwargs):
        '''Run QM/Me scheme for periodic embedding'''

        from sys    import stdout
        from numpy  import arange, array
        from ..io   import format
   
        # choose a surface
        try:
            # should be already determined by <class QMMM>
            surface = self._surface
        except:
            print("\n WARNING: no surface has been chosen so \"{}\" will be used as default\n".format('ab'))
            surface = 'ab'

        # we only do this for the first cycle of geometry optimisation (the cell size is fixed during the optimisation)
        if self._recalc_cell:
            self.frag.cell.consts = array(self.frag.getLatticeOfSelectedPeriodic(arange(self.frag.natoms), surface=surface, recalculate=True))

            # recalculate fractional coordinates
            self.frag.cell.cart2frac()

        # calculate padding distance (unit vector length) along periodic dimensions if the code doesn't allow redundant centres in cell (e.g., CP2K)
        # NB: in most cases we should NOT contain the redundant centres because the system is usually asymmetric especially for adsorbate+surface
        if not self._cell_redundant and self._recalc_cell:
            # always reset vectors before calculating constants! otherwise constants are wiped according to vectors!
            self.frag.cell.vectors[:] = 0.0

            # calculate the unit vector lengths: do NOT include adsorbate atoms!
            padded_consts = array(self.frag.getLatticeOfSelectedPeriodic(arange(self.frag.natoms),
                                                                         surface     = surface,
                                                                         accuracy    = 0.5,
                                                                         recalculate = True,
                                                                         redundant   = self._cell_redundant))

            # always reset vectors before calculating constants! otherwise constants are wiped according to vectors!
            self.frag.cell.vectors[:] = 0.0

            # apply the padding in all dimensions because MM is calculated in 3D
            self.frag.cell.consts = [ p for i, p in enumerate(padded_consts) ]

            # recalculate fractional coordinates
            self.frag.cell.cart2frac()

        print(format.wrap("Lattice constants (a.u.) of the selected periodic MM region automatically assigned:"))
        print(" a        = {:>12.6f}".format(self.frag._cell.a))
        print(" b        = {:>12.6f}".format(self.frag._cell.b))
        print(" c        = {:>12.6f}".format(self.frag._cell.c))
        print(" alpha    = {:>12.6f}".format(self.frag._cell.alpha))
        print(" beta     = {:>12.6f}".format(self.frag._cell.beta))
        print(" gamma    = {:>12.6f}".format(self.frag._cell.gamma))
#        print(" vector a = {:>12.6f} {:>12.6f} {:>12.6f}".format(*self.frag._cell.vectors[:,0]))
#        print(" vector b = {:>12.6f} {:>12.6f} {:>12.6f}".format(*self.frag._cell.vectors[:,1]))
#        print(" vector c = {:>12.6f} {:>12.6f} {:>12.6f}".format(*self.frag._cell.vectors[:,2]))

        ierror = run.run(self, cmdpath, stderr=self.logger, **kwargs)

        # switch off recalculation of cell constants from here on
        self._recalc_cell = False

        stdout.flush()
        return ierror


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

        if hasattr(self, '_runargs'):
            return self._runargs
        else:
            # default is just self.input
            return self.input

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

        self._runargs = val


    @property
    def restart(self):
        '''Restart from a previous calculation'''

        return self._restart


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

        self._restart = val


    @property
    def _underscored_name(self):
        ''''''

        return '_' + self.__class__.__name__.lower()


    @property
    def _extrafiles(self):
        '''Additional output files'''

        return []


    @property
    def path(self):
        '''Path of command to execute the code'''

        return self._path


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

        self._path = val


    @property
    def sysmpi(self):
        '''System MPI command'''

        return self._sysmpi


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

        self._sysmpi = val


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

        return self._ff


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

        self._ff = val 


    # YL 11/08/2019: this is for fixing a very very big bug that the MM input/output puts atoms in the order of QM first and MM then
    #                but Py-ChemShell takes gradients in the original order
    @property
    def _qmmm_order(self):
        '''QM and MM blocks are sequential but QM region can be discontinuous'''

        from numpy import array, arange, concatenate

        qmatoms = array(self.qmatoms)
        mmatoms = array(self.mmatoms)
        if qmatoms.size and mmatoms.size:
            _order = concatenate((qmatoms, mmatoms), axis=None)
        else:
            _order = arange(self.frag.natoms)

        return _order




