#  Copyright (C) 2018 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'

from .dl_find        import DL_FIND
from ..dl_py2f.utils import fileutils

class Parameterise(DL_FIND):
    '''Molecular Dynamics'''

    from numpy import zeros

    _kwargs = {
                'atommode'   : False,
                'callback'   : lambda x:0,
                'coordinates':'cartesian',
                'gradients'  : zeros((1,), dtype=float),
                'method'     :'',
                'mm'         : None,
                'objective'  : zeros((1,), dtype=float),
                'objfunc'    : lambda x:0.0,
                'parameters' : zeros((1,), dtype=float),
                'qm'         : None,
                'ref'        : 0.0,
                'setParameters' : lambda x:0,
              }

#    _kwargs.update(opt.Opt._kwargs)
#    _kwargs2 = {
#               }
#    _kwargs.update(_kwargs2)

    def run(self, **kwargs):
        ''''''

        selectcases = { 'GAAMP': self.GAAMP }

        selectcases[self.method](**kwargs)


    # Roux's method: J. Chem. Theory Comput. 2013, 9, 3543−3556
    # doi: 10.1021/ct4003477
    def GAAMP(self, **kwargs):
        '''General Automated Atomic Model Parameterization'''
        # GAAMP → DL_FIND → GAAMP.callback → tasks.finitediff.2pfd → GAAMP.displaceFunc
        #                                              ↓
        #                                         GAAMP.objfunc
        #                                              ↓
        #                                          tasks.scan

        from sys    import stdout
        from numpy  import asarray, ones, set_printoptions
        from ..     import dl_find
        from ..data import forcefield, unitconvert

        set_printoptions(threshold=50)

        _kwargs = {
                  }

        for key, val in _kwargs.items():
            setattr(self.options, key, val)

        fragObj = self.mm.frag

        # function to get the 1D energy profile around all dihedrals
        # YL TODO: consider the situation that some QM points fail
        def __scan(theory=self.qm):

            from numpy  import append
            from .      import scan

            print("\n >>> ChemShell PARAMETERISE: performing 1D scan...\n")
            stdout.flush()

            energies = asarray([])

            # QM reference: 1D-scan of all dihedrals
            for dihedral in fragObj.dihedrals:
                scanInst = scan.Scan(theory=theory, coord=dihedral, range=0.1, step=0.05).run(mute=True, taskfarm=False)
                energies = append(energies, scanInst.result.energies)
                #print(" >>> energies of dihedral "+str(dihedral)+":\n"+str(scanInst.result.energies)+"\n")

            #print("\n >>> ChemShell PARAMETERSIE: energies of the 1D scan are:\n", energies)
            print("\n Energies of the 1D scan:\n", energies)
            stdout.flush()

            return energies


        self.__qm_ref = asarray([-229.2883674 , -229.28934977, -229.28960874, -229.28914246, -229.28794617, 
                                 -229.28830324, -229.28933634, -229.28960874, -229.28911862, -229.2878609 ,
                                 -229.0440124 , -229.21835629, -229.28960874, -229.21836786, -229.04416427,
                                 -229.28952626, -229.28958503, -229.28960874, -229.2895969 , -229.28954978,
                                 -229.28852087, -229.28947552, -229.28960874, -229.2889251 , -229.2874252 ,
                                 -229.27456526, -229.28548877, -229.28960874, -229.28698793, -229.25562313,
                                 -229.28097101, -229.28960874, -229.28097514, -229.255405  , -229.28954785,
                                 -229.28959646, -229.28960874, -229.28958441, -229.27455781, -229.2854726 ,
                                 -229.28960874, -229.28703037, -229.28851189, -229.28946738, -229.28960874,
                                 -229.28894056, -229.28746354, -229.25509368, -229.2808955 , -229.28960874,
                                 -229.28089169, -229.25531558, -229.28768133, -229.28899862, -229.28960874,
                                 -229.28951109, -229.2887021 , -229.28767907, -229.28900216, -229.28960874,
                                 -229.28949812, -229.28866659, -229.03971863, -229.21642385, -229.28960874,
                                 -229.21557286, -229.0334173 , -229.0342397 , -229.21497901, -229.28960874,
                                 -229.21494816, -229.03304793, -229.28828253, -229.28930077, -229.28960874,
                                 -229.28920753, -229.2880957 , -229.28828692, -229.2893034 , -229.28960874,
                                 -229.289204  , -229.28808769, -229.27664929, -229.28638715, -229.28960874,
                                 -229.28641712, -229.27671029, -229.28763588, -229.28906033, -229.28960874,
                                 -229.28928757, -229.28809905, -229.27495823, -229.28583403, -229.28960874,
                                 -229.28651436, -229.27651978, -229.25960414, -229.2820236 , -229.28960874,
                                 -229.28202061, -229.25960274, -229.27631278, -229.2864599 , -229.28960874,
                                 -229.28578861, -229.27475819, -229.2877169 , -229.28908045, -229.28960874,
                                 -229.28930762, -229.25802092, -229.28161551, -229.28960874, -229.28161853,
                                 -229.25802202, -229.27679621, -229.28636158, -229.28960874, -229.28663867,
                                 -229.27735923, -229.28803735, -229.28919102, -229.28960874, -229.28929211,
                                 -229.28823998, -229.28824943, -229.28926679, -229.28960874, -229.28927646,
                                 -229.28826882, -229.28804392, -229.28919421, -229.28960874, -229.28928908,
                                 -229.28823404, -229.27676903, -229.28635039, -229.28960874, -229.28664588,
                                 -229.28826724, -229.28927607, -229.28960874, -229.28926639])

#        self.__qm_ref = __scan(theory=self.qm)

        npoints = self.__qm_ref.size

        self.mm.parameterise = 'dihedrals'

        ndihedrals = fragObj.dihedral_ztypes.shape[0]
        
        # U(phi) = A*[1 + cos(m*phi − sigma)]
        # take "m" and "sigma" as constants and fit only "A"
        # number of parameters: ndihedrals plus E0
        self.options.parameters.resize((ndihedrals+1,), refcheck=False)
        self.options.gradients.resize((ndihedrals+1,), refcheck=False)

        # weights
#        self.__weights = ones(shape=(npoints,), dtype=float)
        self.options.weights.resize((ndihedrals+1,), refcheck=False)
        # TODO: weights by the user
        self.options.weights[:] = 1.0

        # E0 (the constant gap between QM and MM energies to fit)
        self.__E0 = self.options.parameters[-1]

        iznums = fragObj.dihedral_ztypes[:,0]
        jznums = fragObj.dihedral_ztypes[:,1]
        kznums = fragObj.dihedral_ztypes[:,2]
        lznums = fragObj.dihedral_ztypes[:,3]
#        iznums = fragObj.znums[fragObj.alldihedrals[:,0]]
#        jznums = fragObj.znums[fragObj.alldihedrals[:,1]]
#        kznums = fragObj.znums[fragObj.alldihedrals[:,2]]
#        lznums = fragObj.znums[fragObj.alldihedrals[:,3]]

        # must convert the parameters to a.u.
        try:
            unit = mm.unit
        except:
            unit = 'kcal/mol'
        unit2au = unitconvert.unit2au.get(unit, 1.0)
        au2unit = unitconvert.au2unit.get(unit, 1.0)
#        unit2au = 1.0
#        au2unit = 1.0

        # get the default parameters
        self.options.parameters[:-1] = forcefield.dihedrals[iznums,
                                                            jznums,
                                                            kznums,
                                                            lznums,
                                                            1]*unit2au

#        self.options.parameters[-1] = self.__E0

        def displaceFunc(_frag, displacement, active=None, index=None, **kwargs):
            '''Function to make the displacement (namely, variation to the parameters) for calculating finite-difference gradients'''

            from ..data import forcefield

            # variation by percentage
            self.options.parameters[index] *= (1.0+displacement)

            # parameters excluding E0
            if index < self.options.parameters.size-1:
                forcefield.dihedrals[iznums,jznums,kznums,lznums,1] = self.options.parameters[:-1]*au2unit

            return 0
            

        # 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 ..data import forcefield

            try:
                forcefield.dihedrals[iznums,jznums,kznums,lznums,1] = self.options.parameters[:-1]*au2unit
                print("\n >>> ChemShell PARAMETERISE: new parameters are (in %s)\n"%unit)
                print(" "*3, forcefield.dihedrals[iznums,
                                                  jznums,
                                                  kznums,
                                                  lznums,
                                                  1])
                print("\n     E0 =", self.options.parameters[-1])
                stdout.flush()
                return 0
            except:
                return 999


        # objective function with type hinting (self.ref: QM reference energy)
        def objfunc(_self) -> float:
            '''Objective function'''

            from ..data import unitconvert

            mm_energies = __scan(theory=_self.mm)

            weights = _self.options.weights
            qm_ref  = _self.__qm_ref
            E0      = _self.__E0

            print

            # TODO: weights
            #result = (((qm_ref - mm_energies - E0)**2*weights).sum()/weights.sum())**0.5
            result = ((((qm_ref - mm_energies - E0)*unitconvert.Ha2kcal_mol)**2).sum()/qm_ref.size)**0.5

            print("\n >>> ChemShell PARAMETERISE: objective function's result =", result)
            stdout.flush()

            return result


        def callback(x) -> int:
            '''Callback function for DL_FIND'''

            from sys     import stdout
            from copy    import copy
            from numpy   import append
            from ..base  import finitediff
       
            # callback always runs single-point calculation
            try:

                mm = self.mm
       
                print("\n >>> DL_FIND iteration cycle", self.__iter, '\n\n')
      
#                nvariables = self.options.parameters.size

                ierror = finitediff.calcGradients(mm,
                                                  energy=lambda:objfunc(self), # this way allows passing the function with parameters
                                                  method='2pfd',
                                                  displace=0.01,
#                                                  active=self.active,
                                                  variables=self.options.parameters,
                                                  gradients=self.options.gradients,
                                                  displaceFunc=displaceFunc,
                                                  **kwargs)

                # archive input and output files
                archivedir = '_dl_find'
                fmt = '.%0'+str(len(str(1000)))+'d'

                try:
                    fileutils.archiveFile([mm.input, mm.output] , dir=archivedir, suffix=fmt%self.__iter)
                except:
                    print("\n >>> fileutils.archiveFile ERROR.")
                    stdout.flush()
       
                # counter
                self.__iter += 1

                return 0
       
            except:
                return 1


        self.__iter = 0 

        # YL 26/01/2021: it's sure this parameterisation module no longer works but it can be refactored as <class Opt> in the future
        self.setParameters = setParameters
        self.objfunc   = objfunc
        self.callback  = callback

        dl_find.run(self.mm, self.options, parameterise=True)


Parameterize = Parameterise
Param        = Parameterise
