#  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 zeros, float64, full, int64, array
from .xyz   import _XYZ

# the name "adq" in Tcl-ChemShell is ambiguous
class ShiftQs(_XYZ):
    '''Class of shifted point charges'''

# YL 16/03/2021: reviewed/refactored to make the variable names more reasonable and understandable;
#                this class is used only internally so some abbreviations are OK:
#                src   - source
#                dst   - desination
#                coefs - coefficients
    _kwargs = {
                'atoms_src'   : full((1,), -1, dtype=int64),
                'atoms_dst'   : full((1,), -1, dtype=int64),
                'coords'      : zeros(shape=(1,3), dtype=float64),
                'charges'     : zeros(shape=(1,), dtype=float64),
                'factors'     : zeros(shape=(1,), dtype=float64),
                'delqs'       : zeros(shape=(1,), dtype=float64),
                'dipoles'     : zeros(shape=(1,), dtype=float64),
                'q0'          : zeros(shape=(1,), dtype=float64),
                'grads_src'   : zeros(shape=(1,3), dtype=float64),
                'grads_dst'   : zeros(shape=(1,3), dtype=float64),
                'coefs_src'   : zeros(shape=(1,), dtype=float64),
                'coefs_dst'   : zeros(shape=(1,), dtype=float64),
                'names'       : full((1,), '', dtype='S8'),
                'unique_delqs': zeros(shape=(1,), dtype=float64),
                'nshiftqs'    : 0,
                # assume: r0 = 2.5 a.u.
                'r0'          : 2.5,
                'shift0'      : 0.15
              }

    _internals = {
                   '_mmatoms': 0
                 }

    _fields = 'names', 'atoms_src', 'atoms_dst', 'coords', 'charges', 'factors', 'delqs', 'unique_delqs', 'dipoles', 'q0', 'grads_src', 'grads_dst', 'coefs_src', 'coefs_dst'

    synons = {
               'nshiftqs': 'npts',
               'coords'  : 'xyz'
             }
    _synons = {
                'nshiftqs': 'npts',
                'coords'  : 'xyz'
              }


    @property
    def nshiftqs(self):
        '''Number of shifted charges'''

        return self.natoms


    @nshiftqs.setter
    def nshiftqs(self, val):
        '''Setter of shifted charges'''

        self.natoms = val


    @property
    def unique_delqs(self):
        '''delqs before being applied repeatedly for calculating self.charges'''

        return self._master.unique_delqs


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

        self._master.setField(val, field='unique_delqs')


    @property
    def delqs(self):
        '''Scaled (by number of adjacent MM atoms) charges of origins'''

        return self._master.delqs


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

        self._master.setField(val, field='delqs')


    @property
    def dipoles(self):
        '''Dipoles based on r0'''

        self._master.setField(self.delqs*self.r0, field='dipoles')

        return self._master.dipoles


    @property
    def q0(self):
        '''Charges based on assumption of shift0'''

        self._master.setField(0.5*self.dipoles/self.shift0, field='q0')

        return self._master.q0


    @property
    def r0(self):
        '''Assumed distance (2.5 a.u. as default)'''

        return self._r0

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

        self._r0 = val


    @property
    def shift0(self):
        '''Assumption that shift0 = 0.15 as default'''

        return self._shift0


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

        self._shift0 = val


    @property
    def charges(self):
        '''Charges'''

        self._master.setField(0.5*self.unique_delqs*self.r0/self.shift0, field='charges')
        self._master.charges[::2] *= -1.0

        return self._master.charges


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

        self._master.setField(val, field='charges')


    @property
    def grads_src(self):
        '''Gradients on shiftqs multiplied by coefs_dst'''

        return self._master.grads_src


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

        self._master.setField(val, field='grads_src')

    @property
    def grads_dst(self):
        '''Gradients on shiftqs multiplied by coefs_src'''

        return self._master.grads_dst


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

        self._master.setField(val, field='grads_dst')


    @property
    def atoms_src(self):
        '''Atoms from which charges are shifted'''

        return self._master.atoms_src


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

        self._master.setField(val, field='atoms_src')


    @property
    def atoms_dst(self):
        '''Atoms to which charges are shifted'''

        return self._master.atoms_dst


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

        self._master.setField(val, field='atoms_dst')


    @property
    def atoms_dstinshiftqs(self):
        '''Indices of atoms_dst in the list of shiftqs'''

        from numpy import isin 

        return isin(self._mmatoms, self.atoms_dst[:]).nonzero()[0]


    @property
    def factors(self):
        '''1.0/float(self.atoms_dst.size)'''

        return self._master.factors


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

        self._master.setField(val, field='factors')


    # "deriv" in Tcl-ChemShell 3.6 (src/hybrid/adq.c)
    @property
    def coefs_src(self):
        '''Coeffiecients for correcting gradients on atoms_src'''

        self._master.coefs_src[0::2] = - self.shift0/self.r0
        self._master.coefs_src[1::2] =   self.shift0/self.r0

        return self._master.coefs_src


    # "deriv" in Tcl-ChemShell 3.6 (src/hybrid/adq.c)
    @property
    def coefs_dst(self):
        '''Coeffiecients for correcting gradients on atoms_dst'''

        self._master.coefs_dst[0::2] = 1.0 + self.shift0/self.r0
        self._master.coefs_dst[1::2] = 1.0 - self.shift0/self.r0

        return self._master.coefs_dst


    @property
    def _mmatoms(self):
        '''A list of MM atoms (in bqs) excluding atoms_src'''

        return self.__mmatoms


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

        self.__mmatoms = val


    @property
    def excluded(self):
        '''A list of MM atoms (atoms_src) excluded in bqs'''

        return list(set(self.atoms_src))

