import numpy

from amuse.units import units

from amuse.units.quantities import is_quantity, value_in, to_quantity

from amuse.datamodel import UnstructuredGrid, StructuredGrid,StructuredBaseGrid

try:
  import matplotlib
  from matplotlib import tri
  if not hasattr(tri, "LinearTriInterpolator"):
      raise Exception("LinearTriInterpolator not in matplotlib.tri")
  matplotlib_available=True
except:
  matplotlib_available=False

class interpolating_2D_remapper(object):
    def __init__(self, source, target,axes_names=None):
        """ this class maps a source grid to a target grid using linear 
            interpolation on a triangulation generated by adding a 
            midpoint to every cell (source should be a structured grid) 
            and thus generating 4 triangles for each cell. Values of the 
            midpoints are averaged from the corners. 
        """
        if len(source.shape) !=2:
            raise Exception("source grid is not 2D")
        if not isinstance(source, StructuredBaseGrid):
            raise Exception("source grid is not instance of StructuredBaseGrid")

        self.source=source
        self.target=target
        self._axes_names=list(axes_names or source.get_axes_names())
        self.generate_triangulation()

    def _generate_nodes(self,grid,attributes):

        Nx,Ny=grid.shape

        x,y=numpy.mgrid[0:Nx,0:Ny]
        x1,y1=numpy.mgrid[0:Nx-1,0:Ny-1]
                
        x_=x.flatten()
        y_=y.flatten()
        x1_=x1.flatten()
        y1_=y1.flatten()

        l1=Nx*Ny

        i=numpy.arange(Nx*Ny).reshape((Nx,Ny))
        i1=(numpy.arange((Nx-1)*(Ny-1))+l1).reshape((Nx-1,Ny-1))

      
        nodes=UnstructuredGrid(len(x_)+len(x1_))
        for name in attributes:
          values1=getattr(grid,name)[x_,y_]
          values2=getattr(grid,name)[x1_,y1_]+getattr(grid,name)[x1_+1,y1_]+\
                  getattr(grid,name)[x1_,y1_+1]+getattr(grid,name)[x1_+1,y1_+1]
          setattr(nodes[0], name, 0.*values1[0])
          setattr(nodes[:l1], name, 1.*values1)
          setattr(nodes[l1:], name, values2/4)
        return nodes

    def _generate_elements_and_boundaries(self,grid):
        Nx,Ny=grid.shape

        l1=Nx*Ny

        i=numpy.arange(Nx*Ny).reshape((Nx,Ny))
        i1=(numpy.arange((Nx-1)*(Ny-1))+l1).reshape((Nx-1,Ny-1))

        e1=numpy.zeros(((Nx-1)*(Ny-1),3),dtype='i')
        e2=numpy.zeros(((Nx-1)*(Ny-1),3),dtype='i')
        e3=numpy.zeros(((Nx-1)*(Ny-1),3),dtype='i')
        e4=numpy.zeros(((Nx-1)*(Ny-1),3),dtype='i')
                
        e1[:,0]=i[:-1,:-1].flatten()
        e1[:,1]=i[1:,:-1].flatten()
        e1[:,2]=i1[:,:].flatten()
        
        e2[:,0]=i[1:,:-1].flatten()
        e2[:,1]=i[1:,1:].flatten()
        e2[:,2]=i1[:,:].flatten()
        
        e3[:,0]=i[1:,1:].flatten()
        e3[:,1]=i[:-1,1:].flatten()
        e3[:,2]=i1[:,:].flatten()
        
        e4[:,0]=i[:-1,:-1].flatten()
        e4[:,1]=i1[:,:].flatten()
        e4[:,2]=i[:-1,1:].flatten()

        elements=numpy.zeros((4*(Nx-1)*(Ny-1),3),dtype='i8')
        elements[0::4,:]=e1
        elements[1::4,:]=e2
        elements[2::4,:]=e3
        elements[3::4,:]=e4
      
        boundaries=[xx.flatten() for xx in [i[:,0],i[-1,:],i[::-1,-1],i[0,::-1]] ]
      
        elem=UnstructuredGrid(len(elements))
        elem.nodes=elements

        return elem,boundaries

    def convert_grid_to_nodes_and_elements(self, grid, attributes=None):
        
        if attributes is None:
            attributes=grid.get_attribute_names_defined_in_store()

        nodes=self._generate_nodes(grid, attributes)
        elements,boundaries=self._generate_elements_and_boundaries(grid)
      
        return nodes,elements,boundaries

    def generate_triangulation(self):

        nodes,elements,boundaries=self.convert_grid_to_nodes_and_elements(self.source, self._axes_names)

        xpos=to_quantity(getattr(nodes,self._axes_names[0]))
        ypos=to_quantity(getattr(nodes,self._axes_names[1]))
        
        self._xpos_unit=xpos.unit
        xpos=xpos.number
        self._ypos_unit=ypos.unit
        ypos=ypos.number

        n1=elements.nodes[:,0]
        n2=elements.nodes[:,1]
        n3=elements.nodes[:,2]
        elem=numpy.column_stack((n1,n2,n3))

        self._triangulation=tri.Triangulation(xpos,ypos,elem)
        
    def sample(self, values, xpos, ypos):
        interpolator=tri.LinearTriInterpolator(self._triangulation,values)
        return interpolator(xpos,ypos)

    def forward_mapping(self, attributes, target_names=None):
        if attributes is None:
            attributes=self.source.get_attribute_names_defined_in_store()
        if target_names is None:
            target_names=attributes
        
        source=self.source.empty_copy()
        channel1=self.source.new_channel_to(source)
        target=self.target.empty_copy()
        channel2=self.target.new_channel_to(target)
        channel3=target.new_channel_to(self.target)
        
        channel1.copy_attributes(attributes)
        channel2.copy_attributes(self._axes_names)
        
        nodes=self._generate_nodes(source,attributes)
                
        xpos=value_in( getattr(target,self._axes_names[0]), self._xpos_unit)
        ypos=value_in( getattr(target,self._axes_names[1]), self._ypos_unit)
                        
        for attribute, target_name in zip(attributes, target_names):
            values=to_quantity( getattr(nodes,attribute) ) 
            unit=values.unit
            values=values.number
            samples=self.sample(values,xpos,ypos)
            setattr(target, target_name, (samples if unit is units.none else (samples | unit)))

        channel3.copy_attributes(target_names)    

class bilinear_2D_remapper(object):
    def __init__(self, source, target, check_inside=True):
        """ this class maps a source grid to a target grid using bilinear 
            interpolation. If check_inside=True, raise exception if any 
            target point outside source grid. 
        """
        if len(source.shape) !=2:
            raise Exception("source grid is not 2D")
        if not isinstance(source, StructuredBaseGrid):
            raise Exception("source grid is not instance of RegularBaseGrid")

        self.source=source
        self.target=target
        self._axes_names=source.get_axes_names()
        self.check_inside=check_inside
        self._weights=None
        self._indices=None

    def _calculate_weights(self):
        x0=getattr(self.source[0,0], self._axes_names[0])
        x1=getattr(self.source[1,1], self._axes_names[0])
        y0=getattr(self.source[0,0], self._axes_names[1])
        y1=getattr(self.source[1,1], self._axes_names[1])
        dx=x1-x0
        dy=y1-y0
                
        x=getattr(self.target, self._axes_names[0])
        y=getattr(self.target, self._axes_names[1])

        ix=numpy.floor((x-x0)/dx).astype(int)
        iy=numpy.floor((y-y0)/dy).astype(int)
        if self.check_inside:
          if numpy.any(ix<0) or numpy.any(ix>self.source.shape[0]-2) or \
             numpy.any(iy<0) or numpy.any(iy>self.source.shape[1]-2):
               raise Exception("target not fully inside (restricted) source grid as required")
        ix=numpy.clip(ix,0, self.source.shape[0]-2)
        iy=numpy.clip(iy,0, self.source.shape[1]-2)

        wx=(x0+(ix+1)*dx-x)/dx
        wy=(y0+(iy+1)*dy-y)/dy
        wx=numpy.clip(wx,0.,1.)
        wy=numpy.clip(wy,0.,1.)

        self._weights=[wx,wy]
        self._indices=[ix,iy]
      
    def _evaluate(self, values):
        ix,iy=self._indices
        wx,wy=self._weights
        result=wx*wy*values[ix,iy]+(1.-wx)*wy*values[ix+1,iy]+ \
               wx*(1.-wy)*values[ix,iy+1]+(1.-wx)*(1.-wy)*values[ix+1,iy+1]
        return result

    def forward_mapping(self, attributes, target_names=None):
        if attributes is None:
            attributes=self.source.get_attribute_names_defined_in_store()
        if target_names is None:
            target_names=attributes
        if self._weights is None:
            self._calculate_weights()
        
        for attribute, target_name in zip(attributes, target_names):
            values=getattr(self.source,attribute)
            samples=self._evaluate(values)
            setattr(self.target, target_name, samples)

class nearest_2D_remapper(object):
    def __init__(self, source, target, check_inside=True):
        """ this class maps a source grid to a target grid getting closest
        grid value. If check_inside=True, raise exception if any 
            target point outside source grid. 
        """
        if len(source.shape) !=2:
            raise Exception("source grid is not 2D")
        if not isinstance(source, StructuredBaseGrid):
            raise Exception("source grid is not instance of RegularBaseGrid")

        self.source=source
        self.target=target
        self._axes_names=source.get_axes_names()
        self.check_inside=check_inside
        self._indices=None

    def _calculate_weights(self):        
        x=getattr(self.target, self._axes_names[0])
        y=getattr(self.target, self._axes_names[1])

        kwargs={self._axes_names[0]: x, self._axes_names[1]:y}
        indices=self.source.get_index(**kwargs)

        ix=indices[...,0]
        iy=indices[...,1]
        if self.check_inside:
          if numpy.any(ix<0) or numpy.any(ix>self.source.shape[0]-1) or \
             numpy.any(iy<0) or numpy.any(iy>self.source.shape[1]-1):
               raise Exception("target not fully inside source grid as required")
        ix=numpy.clip(ix,0, self.source.shape[0]-1)
        iy=numpy.clip(iy,0, self.source.shape[1]-1)
        
        self._indices=[ix,iy]

    def _evaluate(self, values):
        return values[self._indices[0], self._indices[1]]

    def forward_mapping(self, attributes, target_names=None):
        if attributes is None:
            attributes=self.source.get_attribute_names_defined_in_store()
        if target_names is None:
            target_names=attributes
        if self._indices is None:
            self._calculate_weights()
        
        for attribute, target_name in zip(attributes, target_names):
            values=getattr(self.source,attribute)
            samples=self._evaluate(values)
            setattr(self.target, target_name, samples)

def conservative_spherical_remapper(*args,**kwargs):
    raise Exception("conservative_spherical_remapper has moved to omuse.ext")
