import numpy
class Field3d:
    def __init__(self,data,origin,spacing,period=(None,None,None)):
        assert len(origin)==3,  "'origin' must be of length 3"
        assert len(spacing)==3, "'spacing' must be of length 3"
        assert len(period)==3,  "'period' must be of length 3"
        self.data      = numpy.array(data)
        self.data.setflags(write=False)
        self.numpoints = self.data.shape
        self.origin    = tuple([float(x) for x in origin])
        self.spacing   = tuple([float(x) for x in spacing])
        self.period    = tuple(period)
        return

    def __str__(self):
        str = 'Field3d with\n'
        str+= ' dimension    {}, {}, {}\n'.format(*self.numpoints)
        str+= ' origin       {:G}, {:G}, {:G}\n'.format(*self.origin)
        str+= ' spacing      {:G}, {:G}, {:G}\n'.format(*self.spacing)
        str+= ' period       {}, {}, {}'.format(*self.period)
        return str

    @classmethod
    def from_chunk(cls,chunk,gridg,period=(None,None,None)):
        '''Initialize Field3d from chunk data and global grid.'''
        xg,yg,zg = gridg
        ib,jb,kb = chunk['ibeg']-1, chunk['jbeg']-1, chunk['kbeg']-1
        dx,dy,dz = xg[1]-xg[0], yg[1]-yg[0], zg[1]-zg[0]
        xo,yo,zo = xg[ib]-chunk['ighost']*dx, yg[jb]-chunk['ighost']*dy, zg[kb]-chunk['ighost']*dz
        nx,ny,nz = chunk['data'].shape
        assert (chunk['nxl']+2*chunk['ighost'])==nx, "Invalid chunk data: nxl != chunk['data'].shape[0]"
        assert (chunk['nyl']+2*chunk['ighost'])==ny, "Invalid chunk data: nyl != chunk['data'].shape[1]"
        assert (chunk['nzl']+2*chunk['ighost'])==nz, "Invalid chunk data: nzl != chunk['data'].shape[2]"
        return cls(chunk['data'],origin=(xo,yo,zo),spacing=(dx,dy,dz),period=period)

    def dim(self,axis=None):
        if axis is None:
            return tuple(self.dim(axis=ii) for ii in range(0,3))
        assert axis<3, "'axis' must be one of 0,1,2."
        return self.numpoints[axis]

    def endpoint(self,axis):
        assert axis<3, "'axis' must be one of 0,1,2."
        return self.origin[axis]+(self.numpoints[axis]-1)*self.spacing[axis]

    def derivative(self,axis,keep_origin=False):
        if not keep_origin: 
            # 2nd order FD with h = dx/2
            origin = list(self.origin)
            if self.period[axis] is not None:
                slhi = [slice(None),slice(None),slice(None)]
                sllo = [slice(None),slice(None),slice(None)]
                slmi = [slice(None),slice(None),slice(None)]
                if numpy.isclose(self.origin[axis],-self.spacing[axis]):
                    slhi[axis] = slice(1,-1)
                    sllo[axis] = slice(0,-2)
                    origin[axis] += 0.5*self.spacing[axis]
                else:
                    slhi[axis] = slice(2,None)
                    sllo[axis] = slice(1,-1)
                    origin[axis] -= 0.5*self.spacing[axis]
                slmi[axis] = slice(1,-1)
                slhi = tuple(slhi)
                sllo = tuple(sllo)
                slmi = tuple(slmi)
                data = numpy.zeros_like(self.data)
                data[slmi] = (1./self.spacing[axis])*(self.data[slhi]-self.data[sllo])
                Field3d.__update_ghostcells_periodic(data,axis)
            else: 
                slhi = [slice(None),slice(None),slice(None)]
                sllo = [slice(None),slice(None),slice(None)]
                slhi[axis] = slice(1,None)
                sllo[axis] = slice(0,-1)
                slhi = tuple(slhi)
                sllo = tuple(sllo)
                data = (1./self.spacing[axis])*(self.data[slhi]-self.data[sllo])
                origin[axis] += 0.5*self.spacing[axis]
        else:
            # 2nd order FD with h = dx, one-sided at boundaries
            # https://numpy.org/doc/stable/reference/generated/numpy.gradient.html
            origin = self.origin
            if self.period[axis] is not None: 
                data = numpy.gradient(self.data,self.spacing[axis],axis=axis,edge_order=1)
                Field3d.__update_ghostcells_periodic(data,axis)
            else:
                data = numpy.gradient(self.data,self.spacing[axis],axis=axis,edge_order=2)
        return Field3d(data,origin,self.spacing,self.period)
    
    def gradient(self,axis,keep_origin=False):
        return [self.derivative(axis,keep_origin=keep_origin) for axis in range(0,3)]

    def change_grid(self,origin,spacing,numpoints,EPS_COLLAPSE=1e-12):
        # TBD: periodicity
        assert all([origin[ii]>=self.origin[ii] for ii in range(0,3)]), "New origin is out of bounds."
        endpoint = [origin[ii]+(numpoints[ii]-1)*spacing[ii] for ii in range(0,3)]
        assert all([endpoint[ii]<=self.endpoint(ii) for ii in range(0,3)]), "New end point is out of bounds."
        data = numpy.zeros(numpoints)
        if numpy.allclose(spacing,self.spacing):
            # spacing is the same: we can construct universal weights for the stencil
            i0,j0,k0 = [numpy.floor((origin[ii]-self.origin[ii])/self.spacing[ii]).astype('int') for ii in range(0,3)]
            cx,cy,cz = [numpy.remainder(origin[ii]-self.origin[ii],self.spacing[ii]) for ii in range(0,3)]
            c = weights_trilinear((cx,cy,cz))
            for ii in range(0,2):
                for jj in range(0,2):
                    for kk in range(0,2):
                        if c[ii,jj,kk]>EPS_COLLAPSE:
                            data += c[ii,jj,kk]*self.data[
                                i0+ii:i0+ii+numpoints[0],
                                j0+jj:j0+jj+numpoints[1],
                                k0+kk:k0+kk+numpoints[2]]
        else:
            for ii in range(0,numpoints[0]):
                for jj in range(0,numpoints[1]):
                    for kk in range(0,numpoints[2]):
                        coordinate = (
                            origin[0]+ii*spacing[0],
                            origin[1]+jj*spacing[1],
                            origin[2]+kk*spacing[2])
                        data[ii,jj,kk] = self.interpolate(coordinate)
        return Field3d(data,origin,spacing,self.period)

    def interpolate(self,coordinate,EPS_COLLAPSE=1e-12):
        assert all([coordinate[ii]>=self.origin[ii] for ii in range(0,3)]), "'coordinate' is out of bounds."
        assert all([coordinate[ii]<=self.endpoint(ii) for ii in range(0,3)]), "'coordinate' is out of bounds."
        i0,j0,k0 = [numpy.floor((coordinate[ii]-self.origin[ii])/self.spacing[ii]).astype('int') for ii in range(0,3)]
        cx,cy,cz = [numpy.remainder(coordinate[ii]-self.origin[ii],self.spacing[ii]) for ii in range(0,3)]
        c = weights_trilinear((cx,cy,cz))
        val = 0.0
        for ii in range(0,2):
            for jj in range(0,2):
                for kk in range(0,2):
                    if c[ii,jj,kk]>EPS_COLLAPSE:
                        val += c[ii,jj,kk]*self.data[i0+ii,j0+jj,k0+kk]
        return val

    def move_origin_periodic(self,shift,axis):
        '''Moves origin of grid in periodic direction by shift*spacing.'''
        assert self.period[axis], "Origin can only be moved in periodic directions."
        self.data = numpy.roll(self.data,shift,axis=axis)
        self.data.setflags(write=False)
        origin = list(self.origin)
        origin[axis] += shift*self.spacing[axis]
        self.origin = tuple(origin)
        return
    @staticmethod
    def __update_ghostcells_periodic(data,axis):
        slsrc = [slice(None),slice(None),slice(None)]
        sldst = [slice(None),slice(None),slice(None)]
        slsrc[axis] = slice(1,2)
        sldst[axis] = slice(-1,None)
        data[sldst] = data[slsrc]
        slsrc = [slice(None),slice(None),slice(None)]
        sldst = [slice(None),slice(None),slice(None)]
        slsrc[axis] = slice(-2,-1)
        sldst[axis] = slice(0,1)
        data[sldst] = data[slsrc]
        return

class ChunkIterator:
    '''Iterates through all chunks. 'snapshot' must be an instance
    of a class which returns a Field3d from the method call
    snapshot.field_chunk(rank,key,keep_ghost=keep_ghost).
    One example implementation is UCFSnapshot from suspendtools.ucf.'''
    def __init__(self,snapshot,key,keep_ghost=True):
        self.snapshot   = snapshot
        self.key        = key
        self.keep_ghost = keep_ghost
        self.iter_rank  = 0
    def __iter__(self):
        self.iter_rank = 0
        return self
    def __next__(self):
        if self.iter_rank<self.snapshot.nproc():
            field = self.snapshot.field_chunk(
                self.iter_rank,self.key,keep_ghost=self.keep_ghost)
            self.iter_rank += 1
            return field
        else:
            raise StopIteration

def weights_trilinear(rel_dist):
    assert len(rel_dist)==3, "len(rel_dist) must be 3."
    assert all([rel_dist[ii]>=0.0 for ii in range(0,3)]), "'rel_dist' must be >=0"
    assert all([rel_dist[ii]<=1.0 for ii in range(0,3)]), "'rel_dist' must be <=1"
    cx,cy,cz = rel_dist
    c = numpy.zeros((2,2,2))
    c[0,0,0] = 1.0-(cx+cy+cz)+(cx*cy+cx*cz+cy*cz)-(cx*cy*cz)
    c[1,0,0] = cx-(cx*cy+cx*cz)+(cx*cy*cz)
    c[0,1,0] = cy-(cx*cy+cy*cz)+(cx*cy*cz)
    c[0,0,1] = cz-(cx*cz+cy*cz)+(cx*cy*cz)
    c[1,1,0] = (cx*cy)-(cx*cy*cz)
    c[1,0,1] = (cx*cz)-(cx*cy*cz)
    c[0,1,1] = (cy*cz)-(cx*cy*cz)
    c[1,1,1] = (cx*cy*cz)
    return c