#  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 ..dl_py2f.utils import fileutils, iterutils

def displaceFunc(frag, displacement, active=None, index=None):
    '''This is the default function for making displacements to cartesian coordinates'''

    frag.coords[active][index] += displacement


def calcGradients(theory, method='', **kwargs):
    '''Execute finite-difference gradients calculation'''

    selectcases = {
                    ''                : lambda *a,**k: -1,
                    'analytic'        : lambda *a,**k:  1,
                    'analytical'      : lambda *a,**k:  1,
                    '1pfd'            : one_point_finite_diff_gradients,
                    '2pfd'            : two_point_finite_diff_gradients,
                    'onepointfd'      : one_point_finite_diff_gradients,
                    'twopointfd'      : two_point_finite_diff_gradients,
                    'fd'              : two_point_finite_diff_gradients,
                    'finitediff'      : two_point_finite_diff_gradients,
                    'finitedifference': two_point_finite_diff_gradients,
                  }

    try:
        method = method.replace('-', '').replace('_', '').lower()
    except:
        pass

    method = selectcases.get(method, lambda *a,**k: -1)

    return method(theory, **kwargs)


def calcHessian(theory, method='', displace=1.0E-5, **kwargs):
    '''Execute finite-difference Hessian calculation'''

    selectcases = {
                    ''                : lambda *a,**k: -1,
                    'analytic'        : lambda *a,**k:  1,
                    'analytical'      : lambda *a,**k:  1,
                    'numeric'         : two_point_finite_diff_hessian,
                    'numerical'       : two_point_finite_diff_hessian,
                    '2pfd'            : two_point_finite_diff_hessian,
                    'twopointfd'      : two_point_finite_diff_hessian,
                    'fd'              : two_point_finite_diff_hessian,
                    'finitediff'      : two_point_finite_diff_hessian,
                    'finitedifference': two_point_finite_diff_hessian,
                  }

    try:
        method = method.replace('-', '').replace('_', '').lower()
    except:
        pass

#    method = selectcases.get(method, doNothing) 
#    if method != doNothing:
#        print("\n\n "+"="*48)
#        print("\n ChemShell Module for Finite-Difference Hessian and Thermal Properties--driven by DL_FIND\n")
#        print(" "+"="*48)
    method = selectcases.get(method, lambda *a,**k: -1)

    return method(theory, displace=displace, **kwargs)


def doNothing(theory, **kwargs):
    '''Do nothing'''

    return 1


def one_point_finite_diff_gradients(theory, **kwargs):
    '''Task-farmed one-point finite-difference gradients'''

    print(" "+"="*64)
    print("\n ChemShell One-Point Finite-Difference Gradients Module\n")
    print(" "+"="*64)

    # return 0 to skip analytical gradients in the next step
    return 0


def two_point_finite_diff_hessian(theory, displace, **kwargs):
    '''Task-farmed two-point finite-difference Hessian (invokes DL_FIND)'''

    from ..tasks import opt

    print(" "+"="*64)
    print("\n ChemShell Two-Point Finite-Difference Hessian Module\n")
    print(" "+"="*64)

    thermal = opt.Opt(theory=theory,
                      algorithm='thermal',
                      **kwargs)

    thermal.run()

    # return 0 to skip analytical hessian in the next step
    return 0


def two_point_finite_diff_gradients(theory,
                                    energy=None,               # `energy` is normally a function
                                    displace=1.0E-5,
                                    displaceFunc=displaceFunc,
                                    active='all',
                                    variables=[],
                                    gradients=[],
                                    **kwargs):
    '''Task-farmed two-point finite-difference gradients'''

    from sys        import stdout
    from ctypes     import byref
    from numpy      import arange, argwhere, array_split, asarray, ndarray
    from ..dl_py2f  import py2f
    from ..parallel import parallel
    from ..objects  import result
    from ..io       import banner

    banner.printModuleBanner('ChemShell Two-Point Finite-Difference Gradients Module')

    global resultObj

    # handles
    frag      = theory.frag
    variables = asarray(variables)

    # archive
    archivedir = '_finite_diff'
    fileutils.archiveDir(archivedir)
    stdout.flush()

    # workgroup ID and number of workgroups
    workgroup   = parallel.workgroup
    nworkgroups = parallel.nworkgroups
    nprocesses  = parallel.nprocs_per_wg

    # TODO: active region of list/array
    active = kwargs.get('active', 'all')
    if type(active) is str:
        if active.lower() == 'all':
            active = arange(frag.natoms)
        elif active.lower() == 'qm':
            active = theory.qm_region
    elif type(active) is list or type(active) is ndarray:
        try:
            active = iterutils.getFlattened(active)
        except:
            active = arange(frag.natoms)
    else:
        active = theory.qm_region

#    print("\n >>> Two-Point Finite-Difference Gradients\n")
    print(" Atoms to calculate are:\n", active)

    if nworkgroups > 1:
        print("\n >>> Task-farmed: workgroup %d of %d\n"%(workgroup, nworkgroups))

    # parameter mode
    if variables.size:
        # in parameter mode we need a separate <class Result> instance to avoid clashing with `theory._result`
        resultObj = result.Result()
        # a preliminary size withouth active region
        nfreedoms = variables.size
        print("\n Displacement: {}%\n".format(displace*100.0))
    # atomistic mode
    else:
        resultObj = theory._result
        # a preliminary size withouth active region (a list of atoms rather than of gradients)
        nfreedoms = frag.natoms
        print("\n Displacement: %16.12f a.u.\n"%displace)

    # pass `result` object to Fortran: py2f(resultObj) is volatile (though resultObj is not), so save it in a variable!
    resultPtr = py2f(resultObj)

    # create a list of active freedoms
    abuff = arange(nfreedoms)
    selectcases = { 'all': abuff,
                    'qm' : getattr(theory, 'qm_region', abuff),
                  }
    # text definitions
    try:
        active = selectcases.get(active.lower(), abuff)
    # list/array
    except:
        active = asarray(active)

    # parameter mode
    if variables.size:
        nfreedoms = variables[active].size
        entry = 'energies'
        getattr(resultObj, entry).resize(nfreedoms)
        # use the `energies` entry to smuggle the 1D array of sequential gradients
        resultObj.title = entry
        if nworkgroups*nprocesses > 1:
            parallel.createPyObject(resultObj,
                                    nfreedoms,
                                    0,
                                    0)
    # atomistic mode
    else:
        natoms    = asarray(active).size
        nfreedoms = natoms*3
        resultObj.title = 'gradients'
        if nworkgroups*nprocesses > 1:
            parallel.createPyObject(resultObj,
                                    0,
                                    3,
                                    frag.natoms)

    one_dim_indices = arange(nfreedoms)
    # distribute the above list in an [natoms,3] array:
    #   x y z
    # [[0,1,2], atom 0
    #  [3,4,5], atom 1
    #  [6,7,8], atom 2
    #   ...   ]
    if not variables.size:
        two_dim_indices = one_dim_indices.reshape((natoms,3))

    # construct indices of all single-point calculations: [0,1,2,3,...]

    # indices assigned to current workgroup to deal with
    indices_of_wg = array_split(one_dim_indices, nworkgroups)[workgroup]
    if chemsh.CHEMSH_DEBUG > 3:
        print(" >>> ChemShell 2PFD: indices of current workgroup are", indices_of_wg)

    selectcases = { 0:'X', 1:'Y', 2:'Z' }

    # counter
    nenergies = 0

    # archive input and output files
    fmt = '.%0'+str(len(str(nfreedoms*2)))+'d'

    directions = { -1.0:'l',  1.0:'r' }

    for i in indices_of_wg:

        # get the current atom and the x/y/z (position of current atom in the above 2D-array)
#        print("\n "+"-"*48)
        print("\n\n "+">"*70)
        if variables.size:
            atom, axis = 0, i
            print("\n Evaluating gradient w.r.t. variable %d..."%(i))
        else:
            atom, axis = argwhere(two_dim_indices==i)[0]
            print("\n Evaluating gradient of centre %d along %s-axis..."%(atom, selectcases[axis]))
#        print(" "+"*"*70)
        stdout.flush()
#        print(" "+"-"*48)

        energy_diff = 0.0

        # left and right
        for direction in [ -1.0, 1.0 ]:

            # displace: the argument "active" will be ignored in the parameter mode
#            frag.coords[active[atom]][axis] += displace*direction

            displaceFunc(frag, displace*direction, active=active[atom], index=axis)
            #displaceFunc(frag, 0.0, active=active[atom], index=axis)

            # run the default way to get the energy
            if energy == None:

                if chemsh.CHEMSH_DEBUG > 3:
                    print("\n >>> ChemShell 2PFD: running", theory)
                    stdout.flush()

                # force the subsequent single-point calculations to restart from the first of each workgroup
                restart = nenergies > 0 or theory.restart
                theory.run(_gradients=False, restart=restart, **kwargs)
            
                eresult = resultObj.energy[0]

            else:

                #print("\n >>> evaluating energy function...\n")
                eresult = energy()

            # TODO: create a function for this
            # archiving
            try:
                fileutils.archiveFile([theory._qm.input, theory._qm.output] , dir=archivedir, suffix=fmt%i+directions[direction])
                fileutils.archiveFile([theory._mm.input, theory._mm.output] , dir=archivedir, suffix=fmt%i+directions[direction])
            except:
                fileutils.archiveFile([theory.input, theory.output] , dir=archivedir, suffix=fmt%i+directions[direction])
            else:
                pass

            nenergies += 1

            # dE
            energy_diff += eresult*direction

            # revert to orignial coordinate
#            frag.coords[active[atom]][axis] -= displace*direction
            displaceFunc(frag, -displace*direction, active=active[atom], index=axis)
    
        grad = energy_diff/(displace*2.0)

#        print("\n "+"*"*70)
            #print("\n Energy evaluation %d: %20.12f a.u."%(nenergies, eresult))
        # update gradients array
        if variables.size:
            print(' {} single points evaluated\n'.format(nenergies))
            print(' Energy  : {:20.12f} a.u.'.format(eresult))
            print(" Gradient: {:20.12f} a.u.".format(grad))
            # TODO: resultObj.energies[active] returns a copy? resultObj.energies[active][i] doesn't change the values of resultObj.energies
            resultObj.energies[i] = grad
        else:
            resultObj.gradients[active[atom]][axis] = grad
            print(' {} single points evaluated\n'.format(nenergies))
            print(' Energy  : {:20.12f} a.u.'.format(eresult))
            print(" Gradient: {:20.12f} a.u.".format(grad))
        print('\n '+'<'*70+'\n')
        stdout.flush()

    # synchronise across workgroups
    if variables.size:
        # use `energies` for the sequential gradients
        gradients[:] = resultObj.energies[:]
        if nworkgroups*nprocesses > 1:
            parallel.syncPyObject(entry='energies')
    else:
        if nworkgroups*nprocesses > 1:
            parallel.syncPyObject(entry='gradients')

    # clean up
    if nworkgroups*nprocesses > 1:
        parallel.cleanPyObject()

    stdout.flush()

    if variables.size:
        print("\n Gradients of the sequential parameters:\n", resultObj.gradients)
        stdout.flush()

    # return 0 to skip analytical gradients in the next step
    return 0

