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

import chemsh
from ..dl_py2f.utils import fileutils

# YL 26/01/2021: moved here from tasks/opt.py because these functions are specially for DL_FIND;
#                this is a modernised version of the old callback() which contains global variables;
#                the reference to these callback functions are now passed as components of optionsObj
def callback(*args) -> int:
    '''Callback function for DL_FIND'''

    from sys       import stderr, stdout, exc_info
    from traceback import print_tb
    from copy      import copy
    from numpy     import append, arange, savez_compressed, vstack
    from ..base    import errors
    from ..objects import fragbio

    # callback._taskObj is defined at the <class Tasks> level (so to meet different tasks)
    dryrun        = callback._dryrun
    taskObj       = callback._taskObj
    info_filename = callback._info_filename
    theoryObj = taskObj.theory

    # archive input and output files
    archivedir = '_dl_find'
    if not taskObj._iter: 
        fileutils.archiveDir(archivedir)
    fmt = '.%0'+str(len(str(1000)))+'d'

    # convergence information: max_gradient, max_step, RMS_gradient, RMS_step
    if taskObj._iter > 0:
        taskObj._convergences = vstack((taskObj._convergences, theoryObj._result.convergences.T))
        # dump the information of convergence path to a compressed NumPy .npz file
        # we do it every step in case the optimisation job fails
        savez_compressed(info_filename,
                         energies     = taskObj._convergences[:,0],
                         max_gradient = taskObj._convergences[:,1],
                         max_step     = taskObj._convergences[:,2],
                         RMS_gradient = taskObj._convergences[:,3],
                         RMS_step     = taskObj._convergences[:,4]
                        )

    # YL 18/11/2020: in some special cases (e.g. the QM/Me scheme) with temporary input/output filenames we have to archive at the interface level
    try:
        theoryObj._qm._archive_dir    = archivedir
        theoryObj._qm._archive_suffix = fmt%taskObj._iter
    except:
        pass

    # callback always runs single-point calculation
    try:

        if chemsh.CHEMSH_DEBUG > 0:
            print("\n DL-FIND theory call", taskObj._iter, '\n')

        # at the 1st cycle `restart` may be True by user input or False by default, but afterwards it is always set True
        if taskObj._iter: 
            # YL 18/04/2018: theoryObj.run(restart=True) does not work for QM/MM geometry optimisation
            #                because in QMMM.run() we use part.run(restart=part.restart, **kwargs) to
            #                allow changing restart during shells relaxation
            # save the initial value and set restart True
            if not theoryObj.restart:
                restart = copy(theoryObj.restart)
                theoryObj.restart = True
                theoryObj.run(restart=True, dryrun=dryrun, task=taskObj._task)
                # bring back the initial value
                theoryObj.restart = restart
            # YL 05/08/2019: it's possible that theoryObj.restart isn't boolean (e.g., in DL_POLY 4)
            else:
                theoryObj.run(dryrun=dryrun, task=taskObj._task)
        else:
            theoryObj.run(dryrun=dryrun, task=taskObj._task)

        print()
        stdout.flush()
    
        # parameterisation mode
        try:
            taskObj.derivatives[:] = theoryObj._result.derivatives
            if taskObj.cell_only:
                if taskObj.ndimensions == 2:
                    theoryObj.frag.cell.strain_derivatives[:2] = taskObj.derivatives.reshape(2,3)
                # NB: this is the GULP order
                elif taskObj.ndimensions == 3:
                    theoryObj.frag.cell.strain_derivatives[0,0] = taskObj.derivatives[0]
                    theoryObj.frag.cell.strain_derivatives[1,1] = taskObj.derivatives[1]
                    theoryObj.frag.cell.strain_derivatives[2,2] = taskObj.derivatives[2]
                    theoryObj.frag.cell.strain_derivatives[1,2] = theoryObj.frag.cell.strain_derivatives[2,1] = taskObj.derivatives[3]
                    theoryObj.frag.cell.strain_derivatives[0,2] = theoryObj.frag.cell.strain_derivatives[2,0] = taskObj.derivatives[4]
                    theoryObj.frag.cell.strain_derivatives[0,1] = theoryObj.frag.cell.strain_derivatives[1,0] = taskObj.derivatives[5]
        except:
            pass
    
        # YL 13/02/2019: save the latest geometry (won't be saved if theoryObj.run() fails)
        if not dryrun:
            theoryObj.frag.save('_dl_find.pun')

            # biological format
            if isinstance(theoryObj.frag, fragbio.FragmentBio):
                theoryObj.frag.save('_dl_find.pdb')
                theoryObj.frag.save('_dl_find.pqr')
                if chemsh.CHEMSH_DEBUG or taskObj.save_path:
                    fileutils.archiveFile('_dl_find.pdb' , dir=archivedir, suffix=fmt%taskObj._iter)
                    fileutils.archiveFile('_dl_find.pqr' , dir=archivedir, suffix=fmt%taskObj._iter)

            # YL 13/02/2019: archive if debugging
            if chemsh.CHEMSH_DEBUG or taskObj.save_path:
                # full fragment
                fileutils.archiveFile('_dl_find.pun' , dir=archivedir, suffix=fmt%taskObj._iter)
                # YL 03/07/2021: also save the structure of the active region
                if len(taskObj.active):
                    activeFrag = theoryObj.frag.getSelected(taskObj.active)
                else:
                    activeFrag = theoryObj.frag.getSelected(arange(theoryObj.frag.natoms))
                activeFrag.save('_dl_find_active.pun', mute=True)
                fileutils.archiveFile('_dl_find_active.pun', dir=archivedir, suffix=fmt%taskObj._iter)
                activeFrag.save('_dl_find_active.xyz', mute=True)
                fileutils.archiveFile('_dl_find_active.xyz', dir=archivedir, suffix=fmt%taskObj._iter)

                # QM/MM
                try:
                    fileutils.archiveFile([theoryObj._qm.input, theoryObj._qm.output] , dir=archivedir, suffix=fmt%taskObj._iter)
                    fileutils.archiveFile([theoryObj._mm.input, theoryObj._mm.output] , dir=archivedir, suffix=fmt%taskObj._iter)
                    theoryObj._qm.frag.save('_dl_find_qm.pun', mute=True)
                    fileutils.archiveFile('_dl_find_qm.pun', dir=archivedir, suffix=fmt%taskObj._iter)
                    theoryObj._qm.frag.save('_dl_find_qm.xyz', mute=True)
                    fileutils.archiveFile('_dl_find_qm.xyz', dir=archivedir, suffix=fmt%taskObj._iter)
                    try:
                        fileutils.archiveFile(theoryObj._qm._extra_outputs, dir=archivedir, suffix=fmt%taskObj._iter)
                        fileutils.archiveFile(theoryObj._mm._extra_outputs, dir=archivedir, suffix=fmt%taskObj._iter)
                    except:
                        pass
                # pure QM or MM
                except:
                    fileutils.archiveFile([theoryObj.input, theoryObj.output] , dir=archivedir, suffix=fmt%taskObj._iter)
                    fileutils.archiveFile(theoryObj._extra_outputs, dir=archivedir, suffix=fmt%taskObj._iter)
                else:
                    pass
    
                # YL 25/04/2020: also archive additional files
                # YL 19/09/2020: QM/MM only
                try:
                    fileutils.archiveFile(theoryObj._qm.punch , dir=archivedir, suffix=fmt%taskObj._iter)
                except:
                    pass

        # counter
        taskObj._iter += 1
    
        return 0

    # YL 20/02/2021 if we use only `except:` it somehow doesn't raise with the traceback
    #               we hence use a list of all Python built-in exceptions to get useful information
    except errors.all_builtin_exceptions as e:

        print('\n >>> ERROR: DL_FIND failed to run callback function of energy/gradients evaluation. Exiting...')
        print(f'\n >>> Error message: {e}')
        print('\n >>> Traceback:')
        tb = e.__traceback__
        print_tb(tb)
        stdout.flush()
        stderr.flush()
        # TWK 8/2/2022 Exceptions are ignored by ctypes callback functions 
        # (at least as of Python 3.7.11). Instead, return an appropriate error
        # code for DL-FIND to handle.
        # YL 28/07/2022: Python Exceptions no longer have error codes
#        return(e.code)
        return(999)


# this function will be called by the Fortran code (DL_FIND) to update the parameters on the Python side
def setParameters(x) -> int:
    '''Update the force field data table to be used for generating inputs'''

    from sys       import stdout
    from traceback import print_tb
    from ..base    import errors

    # setParameters._tasksObj is defined at the <class Tasks> level (so to meet different tasks)
    taskObj   = setParameters._taskObj
    theoryObj = taskObj.theory

    # cell-only optimisation
    if taskObj.cell_only:
        # 2D: apply strain to get new cell
        if taskObj.ndimensions == 2:
            # strain matrix: [ 1.0+strains[0], 0.5*strains[2] ]
            #                [ 0.5*strains[2], 1.0+strains[1] ]
            theoryObj.frag.cell.strains[0,0] = 1.0+taskObj.options.parameters[0]
            theoryObj.frag.cell.strains[1,1] = 1.0+taskObj.options.parameters[1]
            theoryObj.frag.cell.strains[0,1] = 0.5*taskObj.options.parameters[2]
            theoryObj.frag.cell.strains[1,0] = 0.5*taskObj.options.parameters[2]
            # apply strains to generate new vectors
            theoryObj.frag.cell.strains2vectors(vectors=taskObj._vectors0, ndimensions=2)
        # 3D: apply strain to get new cell
        elif taskObj.ndimensions ==  3:
            # strain matrix: [ 1.0+strains[0], 0.5*strains[5], 0.5*strains[4] ]
            #                [ 0.5*strains[5], 1.0+strains[1], 0.5*strains[3] ]
            #                [ 0.5*strains[4], 0.5*strains[3], 1.0+strains[2] ]
            # NB: this is the GULP order
            theoryObj.frag.cell.strains[0,0] = 1.0+taskObj.options.parameters[0]
            theoryObj.frag.cell.strains[1,1] = 1.0+taskObj.options.parameters[1]
            theoryObj.frag.cell.strains[2,2] = 1.0+taskObj.options.parameters[2]
            theoryObj.frag.cell.strains[1,2] = theoryObj.frag.cell.strains[2,1] = 0.5*taskObj.options.parameters[3]
            theoryObj.frag.cell.strains[0,2] = theoryObj.frag.cell.strains[2,0] = 0.5*taskObj.options.parameters[4]
            theoryObj.frag.cell.strains[0,1] = theoryObj.frag.cell.strains[1,0] = 0.5*taskObj.options.parameters[5]
            # apply strains to generate new vectors
            theoryObj.frag.cell.strains2vectors(vectors=taskObj._vectors0, ndimensions=3)

    try:
        if chemsh.CHEMSH_DEBUG > 4:
            print("\n >>> ChemShell PARAMETERISE: new parameters are\n", taskObj.options.parameters)
        stdout.flush()
        return 0

    except errors.all_builtin_exceptions as e:

        print('\n >>> ERROR: DL_FIND failed to run callback function of setting parameters. Exiting...')
        print(f'\n >>> Error message: {e}')
        print('\n >>> Traceback:')
        tb = e.__traceback__
        print_tb(tb)
        stdout.flush()
        stderr.flush()
        # TWK 8/2/2022 Exceptions are ignored by ctypes callback functions
        # (at least as of Python 3.7.11). Instead, return an appropriate error
        # code for DL-FIND to handle.
        # YL 30/11/2022: Python Exception no longer has an error code
        return(999)
