from .xyz import _XYZ

class Grid(_XYZ):
    '''Grid for charge fitting'''

    from numpy import zeros, float64, full

    _kwargs = {
                'coords'    : zeros(shape=(1,3), dtype=float64),
                'dists'     : zeros(shape=(1,0), dtype=float64),
                'bqs_dists' : None,
                'potentials': zeros(shape=(1,),  dtype=float64),
                'names'     : full((1,), '',     dtype='S8'),
                'natoms'    : 0,
                'nlayers'   : 3,
                'npoints'   : 2,
                'type'      :'shell',               # 'mmatoms' | 'shell' | 'surface'
                'vdw_scale' : 1.2,
                'vdw_incr'  : 0.0,
              }

    _internals = { '_frag': None }

    _priorities = '_frag', 'type', 'vdw_scale', 'vdw_incr', 'nlayers', 'natoms'

    # fields of <nputils.RecArray _master>: the order matters!
    _fields = 'names', 'coords', 'potentials', 'dists'

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

#    @property
#    def npoints(self):
#        '''Number of points'''
#
#        return self._npoints
#
#    @npoints.setter
#    def npoints(self, val):
#        ''''''
#
#        self._npoints = val

#        self.__buildShell()


    # overloaded _XYZ property
    @property
    def natoms(self):
        '''Number of atoms in the original fragment'''

#        self.__buildShell()

        return self._natoms


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

        from numpy import zeros, float64

        self._natoms = val

        # self.__dtype is of <class list> and generated according to self._kwargs
        self._kwargs['dists'] = zeros(shape=(1,self._natoms), dtype=float64)


    @property
    def npoints(self):
        '''Number of points'''

        return self._master.shape[0]


    @npoints.setter
    def npoints(self, val):
        '''Number of atoms'''

        from numpy           import zeros, float64
        from ..dl_py2f.utils import nputils

        # if _master exists
        try:
            # once natoms is changed self._master has to be recreated to suit the shape of self.dists
            if self._master['dists'].shape[1] != self._natoms:
                self._master = nputils.RecArray(shape=(val,), dtype=self._XYZ__dtype) # note the name mangling of __dtype in XYZ
            self._master.expand(val)

        # create one if not
        except:
            self._master = nputils.RecArray(shape=(val,), dtype=self._XYZ__dtype) # note the name mangling of __dtype in XYZ


    @property
    def names(self):
        '''Grid point names'''

        return self._master.names


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

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


    @property
    def potentials(self):
        '''Potentials'''

        return self._master.potentials


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

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


    @property
    def dists(self):
        '''Distances to the original fragment atoms'''

        return self._master.dists


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

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


    def getBqsDists(self):
        '''Distances between the grid points and bq charges'''

        from numpy import zeros, sqrt, dot, linalg, subtract

        grid_points = self.coords
        bqs_points = self._frag.bqs.coords

        # Calculate the distances
        diff = subtract(grid_points[:,None], bqs_points)
        bqs_dists = linalg.norm(diff, axis=2)

        return bqs_dists


    def getShellsDists(self):
        '''Distances between the grid points and shell charges'''

        from numpy import zeros, sqrt, dot, linalg, subtract

        grid_points = self.coords
        shells_points = self._frag.shells.coords

        # Calculate the distances
        diff = subtract(grid_points[:,None], shells_points)
        shells_dists = linalg.norm(diff, axis=2)

        return shells_dists


    # TODO
    def buildMMAtoms(self, density=1.0, task="init", npoints=0):
        '''Build a grid of the npoints BQ charges's positions nearest to the atoms'''

        

    # TODO
    def buildSurface(self, density=1.0, task="init", npoints=0):
        '''Build a shell of grid points (this "shell" is not the polarisable MM shell!)'''


    def buildShell(self, density=1.0, task="init", npoints=0):
        '''Build a grid of points uniformly distributed on a shell around the atoms (this "shell" is not the polarisable MM shell!)'''

        from numpy  import arange, array, append, count_nonzero, diagonal, einsum, float16, hstack, linalg, meshgrid, multiply, sin, cos, pi, square, sqrt, subtract, tile, vstack, zeros
        from .bqs   import BQs
        from ..data import atomicdata

        frag     = self._frag
        natoms   = frag.natoms
        nspecies = frag.nspecies*self.nlayers

        # Gaussian method: see IOp(6/43)
        if not self.vdw_incr:
            self.vdw_incr = 0.4/self.nlayers**0.5

        # multi-layer model
        radii_species = zeros((0,))
        for layer in range(self.nlayers):
            radii_species = hstack((radii_species, atomicdata.vdw_radii[frag.zspecies]*self.vdw_scale*(1.0+layer*self.vdw_incr)))
#            radii_atoms   = hstack((radii_species, atomicdata.vdw_radii[frag.znums]   *self.vdw_scale*(1.0+layer*self.vdw_incr)))
            
        # dthetas only depends on radii (i.e. species), therefore we do not need to loop over atoms but only species
#        radii_species = multiply(atomicdata.vdw_radii[frag.zspecies], factors)
#        radii_species = atomicdata.vdw_radii[frag.zspecies]*self.vdw_scale*(1.0+layer*self.vdw_incr)
        radii_atoms = atomicdata.vdw_radii[frag.znums]*self.vdw_scale*(1.0+layer*self.vdw_incr)
    
        # shape: (nspecies,)
        dthetas = 1.0/(density**0.5*radii_species)
    
        # positions
        pos   = arange(0).reshape(0,3)
        dists = arange(0).reshape(0,natoms)

        # loop over speices (to generate dthetas)
        for i in range(nspecies):

            dtheta  = dthetas[i]
            # radius of current species
            radius  = [radii_species[i]]
            # current atomic species
            species = tile(frag.species, self.nlayers)[i]
            # atoms list of current species
            homogeneous    = frag.getRange(species)
            # atoms list of other species
            heterogeneous  = list(set(range(natoms)) - set(homogeneous))
            nhomogeneous   = len(homogeneous)

            # shape: (nthetas,)
            thetas  = arange(dtheta, pi, dtheta)
            nthetas = thetas.size

            # fixed dtheta, shape: (nthetas,)
            phi0s      = dtheta/sin(thetas)*0.5
            phi0s[::2] = 0.0

            # generate all theta/phi mesh grids by stacking them together
            # this is 3-4 times faster than including all the below operation in the for loop over thetas
            ts = arange(0).reshape(0,1)
            ps = arange(0).reshape(0,1)
            # loop over thetas/phi0s (to generate phis)
            for j in range(nthetas):
                # each single value of theta corresponds to a series of phi values
                phis  = arange(phi0s[j], 2.0*pi-dtheta/sin(thetas[j])*0.5+phi0s[j], dtheta/sin(thetas[j]))
                # mesh grid of the current theta and associated phi values
                dts, dps = meshgrid(thetas[j], phis)
                ts = vstack((ts,dts))
                ps = vstack((ps,dps))

            # use Einstein summation
            # as of NumPy 1.14, the argument `optimize=True` (default) fails due to a NumPy error. See:
            # https://github.com/numpy/numpy/issues/10343
            # however the fix of using optimize=False does not affect the efficiency as optimisation is not attempted for two input argements
            dx = einsum('ij,i->ij', cos(ps)*sin(ts), radius, optimize=False).ravel()
            dy = einsum('ij,i->ij', sin(ps)*sin(ts), radius, optimize=False).ravel()
            dz = einsum('ij,i->ij', cos(ts)        , radius, optimize=False).ravel()

            # copy=False argument slows it down?
            dx, x = meshgrid(dx, frag.x[homogeneous])
            dy, y = meshgrid(dy, frag.y[homogeneous])
            dz, z = meshgrid(dz, frag.z[homogeneous])
            x += dx
            y += dy
            z += dz

            nxs, nys, nzs = x.shape[1], y.shape[1], z.shape[1]

            # positions of new points
            points = vstack((x.ravel(), y.ravel(), z.ravel())).T

            # displacement from the points to all original coordinates
            diff = subtract(points[:,None], frag.coords)
 
            # Calculate the distances 
            working_dists = linalg.norm(diff, axis=2)

            # is the distances greater than the radii of all atoms?
            comparison = working_dists > radii_atoms

            # comparison of current species
            # have to merge homo- and heterogeneous, because diagonal(comparison[:,homogeneous])
            # implicitly makes a copy while diagonal(comparison_homo) doesn't
            comparison_homo   = comparison[:,homogeneous]

            # comparison of other species
            comparison_hetero = comparison[:,heterogeneous]

            # exclude distances from points to the origin atoms (from which the points are generated)
            exclusion = diagonal(comparison_homo.T.reshape(nhomogeneous, nhomogeneous, nxs))
            # read-only
            exclusion.setflags(write=True)
            exclusion[:] = True

            # no point should be within the radii from the atoms other than the origin atoms
            comparison = hstack((comparison_homo, comparison_hetero))
            taken = comparison.all(axis=1)
            npoints += count_nonzero(taken)
## slow implementation using explicit for loop
##                for k in range(nhomogeneous):
##                    comparison[nxs*k:nxs*(k+1), homogeneous[k]] = True
##                npoints += count_nonzero(comparison.all(axis=1))

            if task == 'pos':
                pos   = vstack((pos  , points[taken]))
                dists = vstack((dists, working_dists[taken]))

        # run over again with new npoints to get the final points
        if task == 'init':

            density = float(self.npoints)/float(npoints)

            self.buildShell(density=density, task='pos', npoints=0)

        # save the positions to bqs
        elif task == 'pos':

            self.npoints = npoints
            self.coords  = pos
            self.dists   = dists

            print("\n >>> Shell-type grid of %d points has been built\n"%npoints)


    def buildGrid(self):
        '''Build the grid according to chosen type'''

        selectcases = { 
                        'mmatoms': self.buildMMAtoms,
                        'shell'  : self.buildShell,
                        'surface': self.buildSurface,
                      }

        selectcases[self.type](density=1.0, task='init', npoints=0)

        # Calculate the distances from grid points to bqs, if we have them
        if self._frag.nbqs > 0:
            self.bqs_dists = self.getBqsDists()






    __all__ = 'coords', 'potentials'
