"""
Efficient implementation of unstructured mesh operations

Created Feb 8th, 2018 by C. Lapeyre (lapeyre@cerfacs.fr)
"""
import numpy as np
from scipy import sparse


class UnstructuredAdjacency():
    """Efficient scipy implementation of unstructured mesh adjacency ops

    The connectivity is stored in a sparse adjencency matrix A of shape
    (nnode, ncell). The gather operation on vector X (nnode) yields the
    scattered vector Y (ncell), and the scatter operation yields the
    filtered vector X' (nnode). This writes:
        Y  = 1/nvert . A . X
        X' = 1/bincount . A^t . Y
    where ^t is the transpose operation.

    The gatter-scatter operation resulting in filtering X can be performed
    efficiently by storing:
        F  = 1/bincount . A^t . 1/nvert . A
        X' = F . X
    """
    
    def __init__(self, connectivity):
        """Store connectivity as adjacency matrix, extract global parameters

        connectivity: a dictionary-like object with
            {number_of_vertex: connectivity}
          where:
            - number_of_vertex : an integer (number of vertex per element)
            - connectivity     : a list of node indices describing elements of
                                 length number_of_vertex * number_of_cells
          One entry is expected per element type. Ex:
              {3: [triangle_cells_node_list],
               4: [quad_cells_node_list]}
        """
        self._connectivity = connectivity
        concat = np.hstack(list(connectivity.values()))
        self.nnode = np.unique(concat).size
        min_node_number = concat.min()
        # Ensure node numbering starts at 0
        self._connectivity = {nvert: conn - min_node_number
                              for nvert, conn in self._connectivity.items()}
        self._bincount = np.bincount(concat - min_node_number)
        self._adjacency_matrix = self._make_adjacency()

    @property
    def ncell(self):
        """Total number of cells"""
        return sum(self.ncell_per_nvert.values())

    @property
    def ncell_per_nvert(self):
        """Dictionary of {nvert: ncell}

        For each type of element with nvert vertices, stores the nubmer of
        cells ncell
        """
        return {nvert: conn.size // nvert
                for nvert, conn in self._connectivity.items()}

    def _make_adjacency(self, weights=None):
        """Create adjacency matrix by reading connectivity"""
        if weights is None:
            weights = {nvert: np.ones(self._connectivity[nvert].shape)
                       for nvert in self._connectivity}
        adjacency = []
        for nvert in sorted(self._connectivity):
            conn = self._connectivity[nvert]
            ncell = self.ncell_per_nvert[nvert]
            indptr = np.arange(ncell + 1) * nvert
            adjacency.append(sparse.csr_matrix((weights[nvert], conn, indptr),
                                               shape=(ncell, self.nnode)))
        return sparse.vstack(adjacency)

    @property
    def _node2cell_matrix(self):
        """The (ncell, nnode) matrix"""
        node_weights = {nvert:
                        np.ones(self.ncell_per_nvert[nvert] * nvert) / nvert
                        for nvert in sorted(self.ncell_per_nvert)}
        return self._make_adjacency(weights=node_weights)

    @property
    def _cell2node_matrix(self):
        """The (nnode, ncell) matrix"""
        mat = self._make_adjacency().T
        weights = sparse.diags(1. / self._bincount)
        return weights * mat

    def get_node2cell(self):
    #def get_node2cell(self, node_vector):
        """Return the node2cell function"""
        return lambda node_vector: self._node2cell_matrix * node_vector

    def get_cell2node(self):
        """Return the cell2node function"""
        return lambda cell_vector: self._cell2node_matrix * cell_vector

    def get_filter(self, times=1):
        """Return the full gather + scatter filter operation

        If you need to perform the operation N times, you can use the
        times attribute.

        """
        gatherscatter = self._cell2node_matrix * self._node2cell_matrix

        def filt(node_vector):
            """Function to return by get_filter method"""
            for _ in range(times):
                node_vector = gatherscatter * node_vector
            return node_vector
        return filt
