#  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, dtype
from .xyz   import _XYZ
from ..data import atomicdata

class Shells(_XYZ):
    '''Shells'''

    __displace = 0.0001

    _kwargs = {
                'charges'  : zeros((1,)             , dtype=float64),
                'coords'   : zeros((1,3)            , dtype=float64),
                'coreatoms': full ((1)  , -1        , dtype=int64)  ,
                'displace' : full ((1,3), __displace, dtype=float64),
                'gradients': zeros((1,3),             dtype=float64),
                'names'    : full ((1  ), ''        , dtype=dtype('S8')),
                'natoms'   : 0,
                'map'      : zeros(shape=(1,), dtype=int64)
              }

    # YL 18/03/2021: an indices array of nonzero shells for merging gradients (BQs incl. shells of zero charge are deleted in QM.frag.bqs)
    _internals = { '_effective_shells': ...,
                 }

    # fields of <nputils.RecArray _master>: the order matters!
    _fields = 'names', 'coreatoms', 'displace', 'coords', 'charges', 'gradients', 'map'

    _synons = {
                'natoms':'npts',
                'coords':'xyz',
              }

    synons = { 'natoms':'npts' }


    @property
    def coords(self):
        '''Absolute positions'''

        return self._master.xyz


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

        from numpy import count_nonzero, subtract
        from .     import xyz

        # unit conversion should be controlled by only I/O
        buff = xyz.xyzarray(val, npts=self.natoms)

        # if frag.coords are only zeros: take values from input and set displacement all zeros 
        if not count_nonzero(self.coords):
            self._master.setField(buff, field='xyz')
            self._master.setField(xyz.xyzarray(0.0, npts=self.natoms), field='displace')
            self.__displace = 0.0
        else:
            self._master.setField(subtract(buff, subtract(self.coords, self.displace)), field='displace')


    @property
    def __cell(self):
        '''A dummy cell object'''

        from .cell import Cell

        # use a dummy cell object to convert fractionals to cartesians
        try:
            self.___cell
        except:
            self.___cell = Cell()

        return self.___cell


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

        self.___cell = val


    @property
    def fractional(self):
        '''Fractional positions'''


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

        from numpy import asarray, copy, zeros

        cell = self.__cell
        # total natoms
        # TODO: include also nbqs
        natoms = cell.natoms + self.natoms

        # fractional coordinates of all centres (core atoms, shells, and bqs)
        arraybuff = zeros((natoms,3))
        arraybuff[:cell.natoms] = cell.fractional[:]
        arraybuff[cell.natoms:] = asarray(val).reshape(self.coords.shape)[:]

        # resize
        cell.natoms = natoms

        # cell._cartesian here is not a pointer
        cell._cartesian.resize((natoms,3))

        # assign fractional coords to cell
        setattr(cell, 'fractional', arraybuff)

        # copy cartesian coords
        setattr(self, 'coords', self.__cell.cartesian[natoms-self.natoms:])


    @property
    def charges(self):
        '''Charges (natoms,1)'''

        return self._master.charges


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

        self._master.charges = val


    @property
    def gradients(self):
        '''Gradients on shells'''

        # gradients on shell atoms will be an image of part of qmmm.frag.bqs.shells
        return self._master.gradients


    @gradients.setter
    def gradients(self, val):
        '''Setter of gradients on shells'''

        # gradients on shell atoms will be an image of part of qmmm.frag.bqs.shells
        self._master.setField(val, field='gradients')


    @property
    def displace(self):
        '''Displacements'''

        return self._master.displace


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

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


    @property
    def coreatoms(self):
        '''Associated core atoms'''

        from numpy import count_nonzero

        if not count_nonzero(self._master.coreatoms >= 0) and self.natoms:
            self.assignCoreAtoms()

        return self._master.coreatoms


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

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


    @property
    def map(self):
        '''Mapping from periodic to fragment'''
        return self._master.map

    
    @map.setter
    @_XYZ._resize1D
    def map(self,val):
        ''''setter of map'''
        self._master.setField(val,field='map')


# END OF CLASS PROPERTIES
# CLASS METHODS


    # associate a shell centre to a core atom if the distance is smaller than tolerance
    # YL 29/01/2019: changed to 0.2 to cover more situations (this GULP is 0.8 Angstrom)
    from ..data import unitconvert
    def assignCoreAtoms(self, tolerance=0.2*unitconvert.Angstrom2Bohr):
        ''''''

        from numpy import add, arange, argmin, count_nonzero, nditer, linalg, where
        
        try:
            natoms  = self.__corecoords.shape[0]
        except:
            return
        nshells = self.natoms

        if not nshells*natoms or not count_nonzero(self.coords) or not count_nonzero(self.__corecoords):
            return

        shellsiter = nditer(arange(nshells), ['zerosize_ok'])

        for ishell in shellsiter:

            cores = arange(natoms)
            diffs = self.coords[ishell] - self.__corecoords[cores]
            dists = linalg.norm(diffs, axis=1)
            # filter using the tolerance
            associateds = where(dists<=tolerance)[0]
            # associateds is a list of perhaps more than one element, so only take the smallest distance
            try:
                mins = argmin(dists[associateds])
                self._master.coreatoms[ishell] = associateds[mins]
            except:
                print(" >>> WARNING: cannot find an associated core atom for shell {} within {} a.u.".format(ishell, tolerance))
                pass


# END OF CLASS METHODS
