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

from numpy import float64, full, int64, ones, zeros
from .     import _Tasks

# this is a parent class for the DL_FIND-based tasks, e.g., Parameterise and Opt
class DL_FIND(_Tasks):
    ''''''

    _kwargs = { 
                'active'           :[],
                'algorithm'        :'lbfgs',
                'cart_hessian'     : False,
                'coords2'          : zeros(shape=(0,3), dtype=float64),
                'conint'           :'no',
                'coordinates'      :'cartesian',
                'coupled_states'   : True,
                'cycle'            : 0,
                # YL 06/03/2021: we can't set the default value as just chemsh.CHEMSH_DEBUG because since 2059382 CHEMSH_DEBUG
                #                is temporarily set as -1 during the classes/packages scan of ChemShell initialisation so the
                #                current class object will always hold a -1 value; but it's better to use a property() (see below)
                'debug'            : 0,
                'delta'            : 0.01,
                'dump'             :-1,
                '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,
                'residues'         : {},                # YL 05/07/2023: changed the keyword to residues which takes a dict from the user,
                                                        #                to avoid confusion with `residues` which passed a list of residue names to Fortran
                '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,
                'neb'              : False,
                'nebk'             : 0.01,
                'neb_cart'         : False,
                'neb_climb_test'   : 3.0,
                'neb_freeze_test'  : 1.0,
                'nframes'          : 0,
                'nimages'          :-1,
                'nresidues'        : 0,                 # YL 05/07/2023: read-only
                'nzero'            :-1,
                'optimiser'        :'dl_find',
                'parameterise'     : False,
                '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,
                'resIDs'           : full((1,), -1, dtype=int64), # YL 05/07/2023: read-only
                'restart'          : False,
                'result_file'      :'dlf_result.pun',
                'scale_step'       : 1.0,
                'soft'             : 0.005,
                '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'          : ones(shape=(1,), dtype=float),
              }


    @property
    def iopt(self):
        '''Internal switch iopt'''

        selectcases = {
                        'sd'          :  0,
                        'steepest'    :  0,
                        'cg'          :  1,
                        'conjugate'   :  1,
                        'cg-restart'  :  2,
                        'lbfgs'       :  3,
                        'l-bfgs'      :  3,
                        'prfo'        : 10,
                        'p-rfo'       : 10,
                        'thermal'     : 11,
                        'qts'         : 12,
                        'nr'          : 20,
                        'newton'      : 20,
                        'quasi-newton': 20,
                        'dyn'         : 30,
                        'dynamics'    : 30,
                        'ln'          : 40,
                        'rs'          : 51,
                        'random'      : 51,
                        'stochastic'  : 51,
                        'ga'          : 52,
                        'genetic'     : 52,
                      }

        return selectcases[self.algorithm.lower()]

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

        # shall not be set through keyword


    @property
    def algorithm(self):
        '''Algorithm'''

        return self._algorithm


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

        self._algorithm = val.lower()


    @property
    def icoord(self):
        '''Internal switch icoord'''

        selectcases_coords = {
                               'cart'     : 0,
                               'cartesian': 0,
                               'hdlc'     : 1,
                               'hdlc-tc'  : 2, # total connection
                               'dlc'      : 3,
                               'dlc-tc'   : 4,
                             }
        selectcases_neb    = {
                                False         :   0,
                                None          :   0,
                               'free'         : 100,
                               'perp'         : 110,
                               'perpendicular': 110,
                               'frozen'       : 120,
                             }
        selectcases_dimer  = {
                                False         :   0,
                                True          : 220,
                             }

        return   selectcases_coords[self.coordinates] \
               + selectcases_neb.get(self.neb, 0)            \
               + selectcases_dimer.get(self.dimer, 0)

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

        # shall not be set through keyword
        pass


    @property
    def debug(self):
        '''Debug level'''

        from .. import CHEMSH_DEBUG

        return CHEMSH_DEBUG


    @property
    def trust_radius(self):
        '''Trust radius method'''

        return self._trust_radius


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

        self._trust_radius = val.strip().lower()

        if self._trust_radius in [ 'const', 'constant' ]:
            self._trust_radius = 'constant'
        elif self._trust_radius in [ 'e', 'energy' ]:
            self._trust_radius = 'energy'
        elif self._trust_radius in [ 'g', 'grad', 'gradient' ]:
            self._trust_radius = 'gradient'
        else:
            print("\n ChemShell/DL_FIND error: wrong value for keyword option \"trust_radius\"\n")
            exit(112)


    # YL 05/07/2023: moved here from opt.py
    @property
    def residues(self):
        '''HDLC residues data held in a dict'''

        return self._residues


    @residues.setter
    def residues(self, val):
        '''Setter of HDLC residue specification.

        Accepts a dictionary of the form:

        { 'res1' : [atom1, atom2, atom3, ...],
          'res2' : [atom4, atom5, atom6, ...] }

        This data is saved as follows:

        nresidues : number of HDLC residues                                needed by DL_FIND
        _resnames : list of residue names (size: nresidues)                not needed by DL_FIND
        _resIDs   : list of residue number for each atom (size: natoms)    needed by DL_FIND
                    (set to -1 if the atom is not in any residue)

        '''

        from numpy           import asarray
        from ..dl_py2f.utils import iterutils

        self._residues = {}

        # YL 05/07/2023: io/keyword.py should only assign the value as dict type
        if type(val) is dict:
            resnames = list(val.keys())
        # YL 05/07/2023: legacy--list/array/tuple
        else:
            resnames = list(val[0].keys())

        # YL 05/07/2023: this is the only way to set self.nresidues which is read-only to the user
        self._nresidues = len(resnames)

        # TWK 17/08/2022 Note that DL_PY2F does not currently accept dictionaries
        # so the residue data is stored as flat arrays. The resIDs list is in the
        # same format as (part of) the spec array in DL-FIND's API.
        # YL 05/07/2023: not in use by DL_FIND
        self._resnames = asarray(resnames, dtype='S16')

        # YL 05/07/2023: this is the only way to set self.resIDs which is read-only to the user
        self._resIDs = full((self.theory.frag.natoms,), -1, dtype=int64)
        ires = 0
        # YL 05/07/2023: io/keyword.py should only assign the value as dict type
        if type(val) is dict:
            for res, atoms in val.items():
                atoms_res = iterutils.getFlattened(atoms)
                for atom in atoms_res:
                    self._resIDs[atom] = ires 
                # flatten
                self._residues.update({res:atoms_res})
                ires += 1
        # YL 05/07/2023: legacy--list/array/tuple
        else:
            for res, atoms in val[0].items():
                for atom in atoms:
                    self._resIDs[atom] = ires 
                ires += 1


    # YL 05/07/2023: set resIDs read-only
    @property
    def resIDs(self):
        '''HDLC residue members.

        Returns a list of residue numbers for each atom in the fragment to be optimised.
        '''

        return self._resIDs


#    @resIDs.setter
#    def resIDs(self, val):
#        '''Set HDLC residue number for each atom in the fragment to be optimised.
#
#        This setter is provided for completeness but in practice the resIDs data
#        will be supplied via the dictionary accepted by the residues option.
#        '''
#
#        self._resIDs = asarray(val)


    # YL 05/07/2023: set nresidues read-only likewise
    @property
    def nresidues(self):
        '''Number of HDLC residues'''

        return self._nresidues


#    @nresidues.setter
#    def nresidues(self, val):
#        ''''''
#

