"""Construct sparse matrix from a local stencil"""

__docformat__ = "restructuredtext en"

import numpy
import scipy

__all__ = ['stencil_grid']

def stencil_grid(S, grid, dtype=None, format=None):
    """Construct a sparse matrix form a local matrix stencil 
    
    Parameters
    ----------
    S : ndarray
        matrix stencil stored in rank N array
    grid : tuple
        tuple containing the N grid dimensions
    dtype :
        data type of the result
    format : string
        sparse matrix format to return, e.g. "csr", "coo", etc.

    Returns
    -------
    A : sparse matrix
        Sparse matrix which represents the operator given by applying
        stencil S at each vertex of a regular grid with given dimensions.

    Notes
    -----
    The grid vertices are enumerated as arange(prod(grid)).reshape(grid).
    This implies that the last grid dimension cycles fastest, while the 
    first dimension cycles slowest.  For example, if grid=(2,3) then the
    grid vertices are ordered as (0,0), (0,1), (0,2), (1,0), (1,1), (1,2).

    This coincides with the ordering used by the NumPy functions 
    ndenumerate() and mgrid().

    Examples
    --------
    >>> stencil = [-1,2,-1]  # 1D Poisson stencil
    >>> grid = (5,)          # 1D grid with 5 vertices
    >>> A = stencil_grid(stencil, grid, dtype=float, format='csr')   
    >>> A.todense()
    matrix([[ 2., -1.,  0.,  0.,  0.],
            [-1.,  2., -1.,  0.,  0.],
            [ 0., -1.,  2., -1.,  0.],
            [ 0.,  0., -1.,  2., -1.],
            [ 0.,  0.,  0., -1.,  2.]])
    
    >>> stencil = [[0,-1,0],[-1,4,-1],[0,-1,0]] # 2D Poisson stencil
    >>> grid = (3,3)                            # 2D grid with shape 3x3
    >>> A = stencil_grid(stencil, grid, dtype=float, format='csr')   
    >>> A.todense()
    matrix([[ 4., -1.,  0., -1.,  0.,  0.,  0.,  0.,  0.],
            [-1.,  4., -1.,  0., -1.,  0.,  0.,  0.,  0.],
            [ 0., -1.,  4.,  0.,  0., -1.,  0.,  0.,  0.],
            [-1.,  0.,  0.,  4., -1.,  0., -1.,  0.,  0.],
            [ 0., -1.,  0., -1.,  4., -1.,  0., -1.,  0.],
            [ 0.,  0., -1.,  0., -1.,  4.,  0.,  0., -1.],
            [ 0.,  0.,  0., -1.,  0.,  0.,  4., -1.,  0.],
            [ 0.,  0.,  0.,  0., -1.,  0., -1.,  4., -1.],
            [ 0.,  0.,  0.,  0.,  0., -1.,  0., -1.,  4.]])
    
    """

    S = numpy.asarray(S, dtype=dtype)
    grid = tuple(grid)

    if not (numpy.asarray(S.shape) % 2 == 1).all():
        raise ValueError('all stencil dimensions must be odd')
    
    if len(grid) != numpy.rank(S):
        raise ValueError('stencil rank must equal number of grid dimensions')
    
    if min(grid) < 1:
        raise ValueError('grid dimensions must be positive')
    
    N_v = numpy.prod(grid)  # number of vertices in the mesh
    N_s = (S != 0).sum()    # number of nonzero stencil entries

    # diagonal offsets 
    diags = numpy.zeros(N_s, dtype=int)  

    # compute index offset of each DoF within the stencil
    strides = numpy.cumprod( [1] + list(reversed(grid)) )[:-1]
    indices = S.nonzero()
    for i,s in zip(indices,S.shape):
        i -= s // 2
    for stride,coords in zip(strides, reversed(indices)):
        diags += stride * coords

    data = S[ S != 0 ].repeat(N_v).reshape(N_s,N_v)

    indices = numpy.vstack(indices).T

    # zero boundary connections
    for index,diag in zip(indices,data):
        diag = diag.reshape(grid)
        for n,i in enumerate(index):
            if i > 0:
                s = [ slice(None) ]*len(grid)
                s[n] = slice(0,i)
                diag[s] = 0
            elif i < 0:
                s = [ slice(None) ]*len(grid)
                s[n] = slice(i,None)
                diag[s] = 0

    # remove diagonals that lie outside matrix
    mask = abs(diags) < N_v
    if not mask.all():
        diags = diags[mask]
        data  = data[mask]
    
    # sum duplicate diagonals
    if len(numpy.unique(diags)) != len(diags):
        new_diags = numpy.unique(diags)
        new_data  = numpy.zeros( (len(new_diags),data.shape[1]), dtype=data.dtype)

        for dia,dat in zip(diags,data):
            n = numpy.searchsorted(new_diags,dia)
            new_data[n,:] += dat
        
        diags = new_diags
        data  = new_data

    return scipy.sparse.dia_matrix((data,diags), shape=(N_v,N_v)).asformat(format)


if __name__ == '__main__':
    D = 2
    
    if D == 1:
        # 1D Laplacian
        S = array([-1,2,-1])
        grid = (4,)
    
    if D == 2:
        # 2D Laplacian
        S = array([[ 0,-1, 0],
                   [-1, 4,-1],
                   [ 0,-1, 0]])
        #S = array([[-1,-1,-1],
        #           [-1, 8,-1],
        #           [-1,-1,-1]])
        grid = (2,1)
    
    if D == 3:
        S = array([[[ 0, 0, 0],
                    [ 0,-1, 0],
                    [ 0, 0, 0]],
                   [[ 0,-1, 0],
                    [-1, 6,-1],
                    [ 0,-1, 0]],
                   [[ 0, 0, 0],  
                    [ 0,-1, 0],
                    [ 0, 0, 0]]])
        grid = (3,4,5)               
    
    A = stencil_grid( S, grid )
    
    print A.todense()



