from amuse.units.quantities import zero, as_vector_quantity, column_stack

import numpy

from amuse.datamodel import base
from amuse.datamodel import grids

# maintain for backwards compatibility, these should go..
grids.Grid.add_global_vector_attribute("position", ["x","y","z"])
grids.Grid.add_global_vector_attribute("momentum", ["rhovx","rhovy","rhovz"])
grids.Grid.add_global_vector_attribute("magnetic_field", ["B1i","B2i","B3i"])

@grids.BaseGrid.caching_function_for_set
def cellsize(grid):
    raise Exception("a {0} does not have a constant cellsize, use the cellsizes method instead".format(grid.__class__.__name__))

@grids.RegularBaseGrid.caching_function_for_set
def cellsize(grid):
    """Returns the lenght of each direction in the grid.
    Works for regular and cartesian grids.
    """
    result = []
    Ndim= len(grid.shape)
    cell1 = grid[(0,)*Ndim]
    for i in range(len( grid[grid.get_minimum_index()].position )): # shape of position not necessarily same as shape of the grid?
        if grid.shape[i] > 1:
            cell2=grid[(0,)*i+(1,)+(0,)*(Ndim-1-i)]
            result.append((cell2.position-cell1.position)[i])      
    return as_vector_quantity(result)

@grids.BaseGrid.caching_function_for_set
def get_minimum_index(grid):
    raise Exception("not implemented")

@grids.StructuredBaseGrid.caching_function_for_set
def get_minimum_index(grid):
    return tuple(numpy.zeros_like(grid.shape))
    
@grids.BaseGrid.caching_function_for_set
def get_maximum_index(grid):
    raise Exception("not implemented")

@grids.StructuredBaseGrid.caching_function_for_set
def get_maximum_index(grid):
    return tuple(grid.shape - numpy.ones_like(grid.shape))

@grids.BaseGrid.caching_function_for_set
def get_minimum_position(grid):
    raise Exception("not implemented")

@grids.RectilinearBaseGrid.caching_function_for_set
def get_minimum_position(grid):
    return grid[grid.get_minimum_index()].position - 0.5 * grid.cellsize()

@grids.BaseGrid.caching_function_for_set
def get_maximum_position(grid):
    raise Exception("not implemented")

@grids.RectilinearBaseGrid.caching_function_for_set
def get_maximum_position(grid):
    return grid[grid.get_maximum_index()].position + 0.5 * grid.cellsize()

@grids.BaseGrid.caching_function_for_set
def get_volume(grid):
    raise Exception("not implemented")
    
@grids.RectilinearBaseGrid.caching_function_for_set
def get_volume(grid):
    maximum_position = grid.get_maximum_position()
    minimum_position = grid.get_minimum_position()
    delta = maximum_position - minimum_position
    return delta.prod()
    
@grids.BaseGrid.function_for_set
def contains(grid, points):
    raise Exception("not implemented")

@grids.RectilinearBaseGrid.function_for_set
def contains(grid, points):
    return numpy.logical_and(
        numpy.all(points >= grid.get_minimum_position(), axis=len(points.shape)-1),
        numpy.all(points < grid.get_maximum_position(), axis=len(points.shape)-1)
    )
   
@grids.BaseGrid.function_for_set
def points(grid):
    raise Exception("not implemented")

@grids.RegularBaseGrid.function_for_set
def points(grid):
    shape=grid.shape
    dx = grid.cellsize()/2
    cell_centers=grid.position

    shape_with_boundary = numpy.asarray(cell_centers.shape) + 1
    shape_with_boundary[-1] -= 1
    result = numpy.zeros(shape_with_boundary)* cell_centers.flat[0]
    
    for i in range(2**len(shape)):
        slicing=()
        offset=[]
        for j in range(len(shape)):
            if i & 2**j:
                slicing+=(slice(1,None),) 
                offset.append(1)
            else:
                slicing+=(slice(None,-1),)
                offset.append(-1)
                        
        result[slicing]=cell_centers+dx*numpy.asarray(offset)    

    return result
    
@grids.BaseGrid.function_for_set
def connectivity(grid):
    raise Exception("not implemented")

@grids.RegularBaseGrid.function_for_set
def connectivity(grid):
    cellcenters = grid.position
    shape = numpy.asarray(cellcenters.shape)
    dim=len(shape)-1
    shape[-1] = 2**dim
    shape_with_boundary = numpy.asarray(cellcenters.shape) + 1
    shape_with_boundary = shape_with_boundary[:dim]
    indices = numpy.arange(0, numpy.prod(shape_with_boundary), dtype=numpy.int32).reshape(shape_with_boundary)
    result = numpy.zeros(shape, dtype=numpy.int32)

    for i in range(2**dim):
        slicing1=()
        slicing2=()
        for j in range(dim):
            if i & 2**j:
                slicing2+=(slice(1,None),) 
                if len(slicing1) == 0 or slicing1[-1] is not Ellipsis:
                    slicing1+=(Ellipsis,)
            else:
                slicing2+=(slice(None,-1),)
                if len(slicing1) == 0 or slicing1[-1] is not Ellipsis:
                    slicing1+=(Ellipsis,)

        slicing1+=(i,)

        result[slicing1]=indices[slicing2]    

    return result

@grids.BaseGrid.function_for_set
def overlaps(grid, grid1,eps=None):
    raise Exception("not implemented")

@grids.RectilinearBaseGrid.function_for_set
def overlaps(grid, grid1,eps=None):
    """simple test for overlap
       optional keyword parameter:
       
       [eps]: size of buffer (to ignore just touching regions)
    """
    minp=grid.get_minimum_position()
    maxp=grid.get_maximum_position()
    minp1=grid1.get_minimum_position()
    maxp1=grid1.get_maximum_position()
    if eps is not None:
      minp+=eps
      maxp-=eps
      minp1+=eps
      maxp1-=eps
    if (maxp<=minp1).sum()>0 or (minp>=maxp1).sum()>0:
      return False
    return True
    
@grids.BaseGrid.function_for_set
def get_overlap_with(grid, grid1,eps=None):
    raise Exception("not implemented")

@grids.RegularBaseGrid.function_for_set
def get_overlap_with(grid, grid1,eps=None):
    """return overlapping subgrid"""
    if not grid.overlaps(grid1,eps):
      return None
    minindex=grid.get_minimum_index()
    maxindex=grid.get_maximum_index()
    cellsize=grid.cellsize()
        
    minp=grid.get_minimum_position()
    maxp=grid.get_maximum_position()
    minp1=grid1.get_minimum_position()
    maxp1=grid1.get_maximum_position()
    if eps is not None:
      minp1+=eps
      maxp1-=eps
    
    index_of_minp1=numpy.maximum( numpy.array((minp1-minp)/cellsize,'int'), minindex[:len(cellsize)])
    index_of_maxp1=numpy.minimum( numpy.array((maxp1-minp)/cellsize,'int'), maxindex[:len(cellsize)])
    slices=()
    for i,j in zip(index_of_minp1,index_of_maxp1):
      slices+=(slice(i,j+1),)
    if len(slices)!= len(minindex): slices+=(Ellipsis,)
    return grid[slices]


#@grids.AbstractGrid.function_for_set
#def select_fully_inside(grid, cellsizes=(), coordinates=()):
#    """returns boolean array with cells with cellsizes centered on 
#    coordinates fully inside the grid """
#    gridminx,gridminy=sys.grid.get_minimum_position()
#    gridmaxx,gridmaxy=sys.grid.get_maximum_position()

@grids.BaseGrid.function_for_set
def get_index(grid, pos=None, **kwargs):
    raise Exception("not implemented for a {0} grid".format(grid.__class__.__name__))

@grids.RegularBaseGrid.function_for_set
def get_index(grid, pos=None, **kwargs):
    pos=grid._get_array_of_positions_from_arguments(pos=pos,**kwargs)
    offset = pos - grid.get_minimum_position()
    indices = (offset / grid.cellsize())
    return numpy.floor(indices).astype(numpy.int32)

@grids.BaseGrid.function_for_set
def _get_array_of_positions_from_arguments(grid, **kwargs):
    return grids._get_array_of_positions_from_arguments(grid.get_axes_names(), **kwargs)

           
      

