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

from numpy           import asarray, full, zeros, float64, int64
from .dl_find        import DL_FIND # this is the DL_FIND-like task class
from ..              import dl_find # the DL_FIND module
from ..dl_py2f.utils import iterutils, strutils

# now we only got DL_FIND but will have more
# we provide users some synonyms
_synons = { 'dl_find': dl_find,
            'dl-find': dl_find,
            'dlfind' : dl_find,
            'dlf'    : dl_find }

class Opt(DL_FIND):
    '''Geometry Optimisation'''

    _kwargs = { 
               'active'           :[],
               'algorithm'        :'lbfgs',
#               'buff'             : asarray([[1.0,2.0,3.0]]),
               'callback'         : lambda x:999,
               'cart_hessian'     : False,
               'cell_only'        : False,
               'coords2'          : zeros(shape=(0,3), dtype=float64),
               'conint'           :'no',
               'constraints'      : zeros((0,5), dtype=int64),
               'coordinates'      :'cartesian',
               'coupled_states'   : True,
               'cycle'            : 0,
               'delta'            : 0.01,
               'derivatives'      : zeros((1,), dtype=float),
               'dump'             : 5,
               'dimer'            : False,
               'dimer_interpolate': True,
               'distort'          : 0.0,
               'frag2'            : None,
               'fric0'            : 0.3,
               'fricfac'          : 0.95,
               'fricp'            : 0.3,
               'frozen'           : full((1), -1, dtype=int64),
               'genetic'          : False,
               'gp_c3'            : 1.0,
               'gp_c4'            : 0.9,
               'icoord'           :-1,
               'include_res'      : False,
               'init_hessian'     :-1,
               'initial_hessian'  :'external',
               'inner_atoms'      : None,
               'inner_residues'   : None,
               'iopt'             :-1,
               'lbfgs_mem'        :-1,
               'list_option'      :'medium',
               'ln_t1'            : 0.0001,
               'ln_t2'            : 1.0,
               'mass'             : zeros(shape=(0,), dtype=float64),
               'maxcycle'         : 100,
               'maxene'           : 100000,
               'maxmicrocycle'    : 100,
               'maxrot'           : 10,
               'maxstep'          : 0.5,
               'maxupdate'        : 50,
               'microiterative'   : False,
               'micro_esp_fit'    : False,
               'minstep'          : 0.00001,
               'ndimensions'      : 2,
               'neb'              :'no',
               'nebk'             : 0.01,
               'neb_cart'         : False,
               'neb_climb_test'   : 3.0,
               'neb_freeze_test'  : 1.0,
               'neb_path_sampling': -1.0,
               'ncartcons'        : 0,
               'ncons'            : 0,
               'nframes'          : 0,
               'nimages'          :-1,
               'nzero'            :-1,
               'optimiser'        :'dl_find',
               'parameterise'     : False,
               'parameters'       :[],
               'pf_c1'            : 5.0,
               'pf_c2'            : 5.0,
               'po_contraction'   : 0.9,
               'po_death_rate'    : 0.5,
               'po_distribution'  :'force_bias',
               'po_init_pop_size' : 50,
               'po_maxcycle'      : 10000,
               'po_mutation_rate' : 0.15,
               'po_nsave'         : 10,
               'po_pop_size'      : 25,
               'po_radius'        : 1.0,
               'po_reset'         : 500,
               'po_scalefac'      : 10.0,
               'po_tolerance_g'   : 0.001,
               'po_tolerance_r'   : 1.0E-8,
               'qts'              : False,
               'qtsrate'          : False,
               'rate'             : False,
               'restart'          : False,
# attribute name taken by `self.result`
#               'result'           :'dl_find.result',
               'result_file'      :'dlf_result.pun',
               'save_path'        : False,
               'scale_step'       : 1.0,
               'setParameters'    : lambda x:999,
               'soft'             : 1.0E20,
               'state_i'          : 1,
               'state_j'          : 2,
               'stochastic'       : False,
               'thermal'          : False,
               'timestep'         : 1.0,
               'tolerance'        : -1.0,
               'tolerance_e'      : -1.0,
               'tolrot'           : 5.0,
               'tsplit'           : False,
               'tsrelative'       : False,
               'trust_radius'     :'constant',
               'update'           :-1,
               'update_method'    :'bofill',
               'weights'          : None
              }

    _internals = { 
                   'coords2'      : zeros(shape=(0,3), dtype=float64),
                   '_iter'        : 0,
                   '_task'        :'standard',
                   '_convergences': zeros(shape=(0,5), dtype=float64), # convergence information: energy, max grad, max step, RMS grad, RMS step
                 }

    _synons = { 'coordstype': 'coordinates',
                'coordtype' : 'coordinates',
              }

    _priorities = 'theory', 'cell_only', 'ndimensions', 'parameters'


    @property
    def frag2(self):
        '''Auxiliary second fragment (needed by, e.g., NEB)'''

        return self._frag2

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

        try:
            self.coords2 = val.coords
            self._frag2 = val
            self.nframes = 1
        except:
            from ..objects import fragment
            merged = fragment.Fragment()
            merged.merge(*val, noself=True)
            self._frag2 = merged
            self.coords2 = merged.coords
            self.nframes = len(val)


    @property
    def neb(self):
        '''NEB method'''

        # set default when frag2 given
        if not self.frag2 and self._neb != 'no':
            self._neb = 'no'
            print(" ChemShell >>> NEB turned off since frag2 not given")

        if self.frag2 and self._neb != 'no':
            self._neb = 'frozen'

        return self._neb

    @neb.setter
    def neb(self, val):
        '''Setter of NEB method'''

        # force to switch off NEB
        if val.lower() == 'no':
            self.frag2 = None

        self._neb = val.lower()

    @property
    def coords2(self):
        '''Second coords'''

        return self._coords2

    @coords2.setter
    def coords2(self, val):
        '''Setter of second coords'''

# TODO
        self._coords2 = asarray(val).flatten()


    @property
    def neb_path_sampling(self):
        '''Return the sampling density of images along the reaction path'''

# AJL Aug 2020
# Unclear if this should be the private or public equivalent of this variable?
# There isn't a clear guide as to what is used where - e.g. coords2 and _coords2 are used above, as are neb/_neb
        return self._neb_path_sampling

    @neb_path_sampling.setter
    def neb_path_sampling(self, val):
        '''Setter for NEB path sampling - will also then work out a value for nimages'''

        self._neb_path_sampling = val

# It would make more sense to me if this property was a redirection of self.theory.frag, 
# rather than using different calls for each sets of coords, but I presume there is a reason this is absent?
#        print(self._coords)
#        print(self.theory.frag.coords)
#        print(self._coords2.reshape(self.theory.frag.coords.shape))
#        print(self.theory.frag.coords-self._coords2.reshape(self.theory.frag.coords.shape))
        
        from numpy.linalg import norm
        # None of this needs storing - it can be broken up if you want the variables to be more readable
        linear_interpolation_path_distance = norm(self.theory.frag.coords-self._coords2.reshape(self.theory.frag.coords.shape))
        self.nimages = int(linear_interpolation_path_distance/self.neb_path_sampling)

    @property
    def nimages(self):
        '''Return the number of images along the reaction trajectory'''

        if not self.frag2 and self._neb != 'no':
            if self._nimages == -1:
                self.neb_path_sampling = 0.5 # A very conservative default value
    
        return self._nimages
 
    @nimages.setter
    def nimages(self, val):
        '''Set the number of images, as this value is hidden by default'''

        self._nimages = val

    @property
    def active(self):
        '''List of active centres (complementary to frozen)'''

        return self._active


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

        from numpy import arange, ones

        self._active = asarray(self.theory.frag.getRange(val))

        # only if val is non-emtpy
        if self._active.size:
            self.theory.frag.frozen[:] = 0
            mask = ones(self.theory.frag.frozen.shape, dtype=bool)
            try:
                mask[self._active] = False
            except IndexError as e:
                print(e)
                print("\n >>> ERROR: it's possible theory.frag hasn't been assigned.\n")
            self.theory.frag.frozen[mask] = -1
            self._frozen = arange(self.theory.frag.natoms)[mask]


    @property
    def frozen(self):
        '''List of frozen centres (complementary to active)'''

        return self._frozen


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

        from numpy import arange, ones

        self._frozen = asarray(self.theory.frag.getRange(val))

        # only if val is non-emtpy
        if self._frozen.size:
            self.theory.frag.frozen[:] = 0
            self.theory.frag.frozen[self._frozen] = -1
            mask = ones(self.theory.frag.frozen.shape, dtype=bool)
            mask[self._frozen] = False
            self._active = arange(self.theory.frag.natoms)[mask]


    @property
    def constraints(self):
        '''Constraints specification.

        Returns a list of constraints on each atom in DL-FIND spec array format.
        '''

        return self._constraints

    @constraints.setter
    def constraints(self, val):
        '''Set constraint definitions.

        For HDLC coordinates, accepts a list of the form:

        [ ( 'bond', atom1, atom2, ),
          ( 'angle', atom1, atom2, atom3 ), 
          ... ]

        This data is saved as an array of constraints using DL-FIND spec 
        array conventions (5 integers per constraint, consisting of an 
        integer representing the type and up to 4 atom numbers.)

        For Cartesian coordinates, accepts a list of the form:

        [ ( 'x', atom ),
          ( 'yz', atom ),
          ... ]

        For consistency here we will use the same 5 integer per 
        constraint table, but when passed to DL-FIND this data will 
        actually be saved in the first nat entries of spec 
        (along with data on which atoms are entirely frozen)
        '''

        # Handle empy list default
        if len(val) == 0:
            self._constraints = zeros((0,5), dtype=int64)
            return

        cons = list(val)
        ninpcons = len(cons)

        self._constraints = zeros((ninpcons,5), dtype=int64)

        # DL-FIND integer conventions for constraint types
        #
        # TWK 19/08/2022 Internal coordinates may be used with DLC or
        # HDLC coordinates. In principle, (single axis) cartesian 
        # constraints are also accepted for DLCs and HDLCs (see the 
        # constraints option in TCS), but I'm not convinced they actually
        # work. A working test job should be committed before they are enabled.
        # For Cartesian coordinates, 1 and 2 axis cartesians constraints
        # are supported by DL-FIND, but in this case are supplied via a 
        # different part of the spec array, so need to be counted independently.
        # 
        contyps = { 'bond' : 1,
                 'angle' : 2,
                 'torsion' : 3,
                 'x' : -2,
                 'y' : -3,
                 'z' : -4,
                 'xy' : -23,
                 'xz' : -24,
                 'yz' : -34
                 }

        icon = 0
        self.ncons = 0
        self.ncartcons = 0
        for con in cons:
            # convert type string to DL-FIND type integer
            ityp = contyps[con[0]]
            self._constraints[icon, 0] = ityp
            if (ityp > 0):
                # Constraints for HDLC coordinates
                self.ncons += 1
            else:
                # Constraints for Cartesian coodinates
                self.ncartcons += 1
            # copy the atom numbers in 
            self._constraints[icon,1:len(con)] = con[1:]
            icon += 1

        if self.ncons > 0 and self.ncartcons > 0:
            print("\n ChemShell ERROR: HDLC and Cartesian constraints cannot be mixed. Constraints must match the coordinate system.\n")
            exit(112)


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

        return self._lbfgs_mem

    @lbfgs_mem.setter
    def lbfgs_mem(self, val):
        ''''''

        self._lbfgs_mem = val


    @property
    def tolerance_e(self):
        '''Convergence criterion for the energy change (a.u.)'''

        if self._tolerance_e < 0:
            return self.tolerance/450.0
        else:
            return self._tolerance_e

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

        self._tolerance_e = val


    @property
    def nframes(self):
        '''Number of frames'''

        if self._nframes > 100:
            print(" >>> ERROR: nframes cannot be greater than 100")
            quit(1)
        return self._nframes

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

        self._nframes = val


    @property
    def parameterise(self):
        '''Parameterisation mode'''

        return self.cell_only or len(self.parameters) > 0


    @property
    def parameters(self):
        '''Variables to parameterise'''

        from numpy import array, copy

        if self.cell_only:
            self._task = 'cell_only'
            if self.ndimensions == 2:
                self.theory.ndimensions = 2
                if len(self._parameters) != 3:
                    self._parameters = array([0.0,0.0,0.0])
                    self.derivatives.resize(3, refcheck=False)
            elif self.ndimensions == 3:
                self.theory.ndimensions = 3
                if len(self._parameters) != 6:
                    self._parameters = array([0.0,0.0,0.0,0.0,0.0,0.0])
                    self.derivatives.resize(6, refcheck=False)
            else:
                print("\n ChemShell ERROR: we currently only support 2D and 3D cell-only optimisation\n")
                exit(112)

        return self._parameters


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

        # flatten
        self._parameters = iterutils.getFlattenedIter(val, sort=False)


    @property
    def thermal(self):
        '''Hessian and thermal corrections only'''

        return self._thermal


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

        self._thermal = bool(val)
        if self._thermal:
            self.algorithm = 'thermal'


    def run(self, dryrun=False, info_filename='_profile_opt.npz', **kwargs):
        '''Run goemetry optimization'''

        from sys     import stdout
        from numpy   import copy, savez_compressed, vstack
        from ..base  import run

        # YL: cell-only optimisation
        if self.cell_only:
            # compute the strains
            self.theory._strains = True
            # save the original vectors: the strains being optimised are targeting the original vectors!
            self._vectors0 = copy(self.theory.frag.cell.vectors)

        # make sure gradients are always computed for optimisation
        self.theory._gradients = True

        # if path can be imported as module call run.runLib(), otherwise run.runExec() (see <dict> run.run().runTypes)
        optimiser = strutils.importAsModule(_synons[self.optimiser].__name__)

        # YL 26/01/2021: it's safe to cite the reference because callback funtions are not invoked unless during an optimisation
        optimiser.callback._taskObj        = self
        optimiser.callback._dryrun         = dryrun
        optimiser.callback._info_filename  = info_filename
        optimiser.setParameters._taskObj   = self

        # YL 26/01/2021: handles of functions
        self.callback      = optimiser.callback
        self.setParameters = optimiser.setParameters

        # run
        optimiser.run(self.theory, self.options, **kwargs)

        # report final results
        if self.cell_only:
            _printFinalCell(self)

        # convergence information: max_gradient, max_step, RMS_gradient, RMS_step
        self._convergences = vstack((self._convergences, self.theory._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     = self._convergences[:,0],
                         max_gradient = self._convergences[:,1],
                         max_step     = self._convergences[:,2],
                         RMS_gradient = self._convergences[:,3],
                         RMS_step     = self._convergences[:,4]
                        )

        print()
        stdout.flush()

        # switch off strains by default
        self.theory._strains = False

        return self


# aliases
Optimisation = Opt
Optimization = Opt
Optimise     = Opt
Optimize     = Opt


def _printFinalCell(_self):
    '''Print the final results of cell-only optimisation'''

    if _self.ndimensions in 2:
        print("\n Optimised 2D cell vectors (a.u.):\n")
        print(" X {:>15.9f} {:>15.9f}".format(*_self.theory.frag.cell.vectors[0]))
        print(" Y {:>15.9f} {:>15.9f}".format(*_self.theory.frag.cell.vectors[1]))
        print("\n Original 2D cell vectors (a.u.):\n")
        print(" X {:>15.9f} {:>15.9f}".format(*_self._vectors0[0]))
        print(" Y {:>15.9f} {:>15.9f}".format(*_self._vectors0[1]))
        print("\n Strain tensor matrix applied to the original cell:\n")
        print(" {:>15.9f} {:>15.9f}".format(*_self.theory.frag.cell.strains[0]))
        print(" {:>15.9f} {:>15.9f}".format(*_self.theory.frag.cell.strains[1]))
    elif _self.ndimensions in 3:
        print("\n Optimised 3D cell vectors (a.u.):\n")
        print(" X {:>15.9f} {:>15.9f} {:>15.9f}".format(*_self.theory.frag.cell.vectors[0]))
        print(" Y {:>15.9f} {:>15.9f} {:>15.9f}".format(*_self.theory.frag.cell.vectors[1]))
        print(" Z {:>15.9f} {:>15.9f} {:>15.9f}".format(*_self.theory.frag.cell.vectors[2]))
        print("\n Original 3D cell vectors (a.u.):\n")
        print(" X {:>15.9f} {:>15.9f} {:>15.9f}".format(*_self._vectors0[0]))
        print(" Y {:>15.9f} {:>15.9f} {:>15.9f}".format(*_self._vectors0[1]))
        print(" Z {:>15.9f} {:>15.9f} {:>15.9f}".format(*_self._vectors0[2]))
        print("\n Strain tensor matrix applied to the original cell:\n")
        print(" {:>15.9f} {:>15.9f} {:>15.9f}".format(*_self.theory.frag.cell.strains[0]))
        print(" {:>15.9f} {:>15.9f} {:>15.9f}".format(*_self.theory.frag.cell.strains[1]))
        print(" {:>15.9f} {:>15.9f} {:>15.9f}".format(*_self.theory.frag.cell.strains[2]))
        print("\n Final stress tensor matrix (GPa):\n")
        print("            X               Y               Z")
        print(" X {:>15.9f} {:>15.9f} {:>15.9f}".format(*_self.theory.frag.cell.stress_tensor[0]))
        print(" Y {:>15.9f} {:>15.9f} {:>15.9f}".format(*_self.theory.frag.cell.stress_tensor[1]))
        print(" Z {:>15.9f} {:>15.9f} {:>15.9f}".format(*_self.theory.frag.cell.stress_tensor[2]))


