from __future__ import division
from __future__ import unicode_literals
from builtins import range
from builtins import str
from future.utils import text_to_native_str
__docformat__ = 'restructuredtext'

__all__ = []

from PyTrilinos import Epetra
from PyTrilinos import EpetraExt

from fipy.matrices.sparseMatrix import (_SparseMatrix, _RowMesh2Matrix,
                                        _ColMesh2Matrix, _RowColMesh2Matrix)
from fipy.tools import numerix

# Current inadequacies of the matrix class:

# 1) Adding matrices - the matrix with fewer non-zeros gets added into the one
# that has more; this works as long as it's nonzero entries are a subset of the
# larger one's nonzero entries. Is true for all cases in fipy, but is not true
# in the general case - this isn't a general matrix class like the pysparse
# matrix class is.
#
# 2) addAt currently not guaranteed to work for fill-completed matrices, if
# elements are being added in new spots.
#
# 3) put currently not guaranteed to work for non-empty matrices that do not
# have all the target spots occupied.
#
# None of these situations currently come up in FiPy; tests do not reveal any of
# the warnings that guard for those, and all tests pass. Because of the way
# FiPy constructs its matrices, I do not anticipate any of these occurring.

class _TrilinosMatrix(_SparseMatrix):
    """class wrapper for a PyTrilinos `Epetra.CrsMatrix`.

    Allows basic python operations __add__, __sub__ etc.
    Facilitate matrix populating in an easy way.
    """
    def __init__(self, matrix, bandwidth=None):
        """
        Parameters
        ----------
        matrix : Epetra.CrsMatrix
            The internal Trilinos matrix
        bandwidth : int
            The proposed band width of the matrix.
        """
        self.matrix = matrix

        self.comm = matrix.Comm()
        if bandwidth is None:
            self.bandwidth = ((matrix.NumGlobalNonzeros() + matrix.NumGlobalRows() - 1)
                              // matrix.NumGlobalRows())
        else:
            self.bandwidth = bandwidth

        super(_TrilinosMatrix, self).__init__()

    def _setMatrix(self, m):
        self._matrix = m

    matrix = property(lambda self: self._matrix, _setMatrix)

    @property
    def rowMap(self):
        return self.matrix.RowMap()

    @property
    def colMap(self):
        return self.rowMap

    @property
    def domainMap(self):
        """Distribution of input vector `x` in `y = A*x`

        The domain Map must be one-to-one: that is, each global index in
        the Map must be uniquely owned by one and only one process.
        """
        return self.colMap

    @property
    def rangeMap(self):
        """Distribution of output vector `y` in `y = A*x`

        The range Map must be one-to-one: that is, each global index in
        the Map must be uniquely owned by one and only one process.
        """
        return self.rowMap

    # All operations that require getting data out of the matrix may need to
    # call FillComplete to make sure they work.  There will be no warnings when
    # FillComplete is implicitly called; there will only be warnings when
    # insertions fail.
    def copy(self):
        self.fillComplete()

        return _TrilinosMatrix(matrix=Epetra.CrsMatrix(self.matrix))


    def __getitem__(self, index):
        self.fillComplete()

        return self.matrix[index]

    def __str__(self):
        self.fillComplete()

        s = _SparseMatrix.__str__(self)

        comm = self.matrix.Map().Comm()
        if comm.NumProc() > 1:
            from fipy.tools import parallelComm
            return ''.join(parallelComm.allgather(s))
        else:
            return s

    @property
    def _range(self):
        return (list(range(self.rowMap.NumGlobalElements())), self.rowMap.MyGlobalElements())

    def __setitem__(self, index, value):
        self.matrix[index] = value


    # Addition is tricky.
    # Trilinos interface is as such: A can be added into B, but A has to be
    # Filled() beforehand. If B is filled beforehand, this may or may not
    # crash, depending on whether things are being added into spots in B that
    # were not there before.  Have put in some order-of-operands twiddling to
    # make it look like two things can be added in any order.

    # Though not guaranteed to work for arbitrary matrices, it should work for
    # all those generated by FiPy and will give warnings if it encounters
    # trouble (unless Trilinos runs into an error and aborts instead of
    # returning an error code)

    def __iadd__(self, other):
        if other != 0:
            other.fillComplete()

            # Depending on which one is more filled, pick the order of operations
            if self.matrix.Filled() and other.matrix.NumGlobalNonzeros() \
                                            > self.matrix.NumGlobalNonzeros():
                tempBandwidth = other.matrix.NumGlobalNonzeros() / self.matrix.NumGlobalRows()+1

                tempMatrix = Epetra.CrsMatrix(Epetra.Copy, self.rowMap, tempBandwidth)

                if EpetraExt.Add(other.matrix, False, 1, tempMatrix, 1) != 0:
                    import warnings
                    warnings.warn("EpetraExt.Add returned error code in __iadd__, 1",
                                  UserWarning, stacklevel=2)

                if EpetraExt.Add(self.matrix, False, 1, tempMatrix, 1) != 0:
                    import warnings
                    warnings.warn("EpetraExt.Add returned error code in __iadd__, 2",
                                  UserWarning, stacklevel=2)

                self.matrix = tempMatrix

            else:
                if EpetraExt.Add(other.matrix, False, 1, self.matrix, 1) != 0:
                    import warnings
                    warnings.warn("EpetraExt.Add returned error code in __iadd__",
                                  UserWarning, stacklevel=2)

        return self


    # To add two things while modifying neither, both must be FillCompleted
    def _add(self, other, sign=1):
        self.fillComplete()
        other.fillComplete()

        # make the one with more non-zeros the right-hand operand
        # so addition is likely to succeed
        if self.matrix.NumGlobalNonzeros() > other.matrix.NumGlobalNonzeros():
            tempMatrix = self.copy()
            tempMatrix.__iadd__(other*sign)
        else:
            tempMatrix = other.copy()
            tempMatrix.__iadd__(self*sign)

        return tempMatrix

    def __add__(self, other):
        """
        Add two sparse matrices. The nonempty spots of one of them must be a
        subset of the nonempty spots of the other one.

            >>> L = _TrilinosMatrixFromShape(rows=3, cols=3)
            >>> L.addAt((3., 10., numerix.pi, 2.5), (0, 0, 1, 2), (2, 1, 1, 0))
            >>> L.addAt([0, 0, 0], [0, 1, 2], [0, 1, 2])
            >>> print(L + _TrilinosIdentityMatrix(size=3))
             1.000000  10.000000   3.000000  
                ---     4.141593      ---    
             2.500000      ---     1.000000  

            >>> print(L + 0)
                ---    10.000000   3.000000  
                ---     3.141593      ---    
             2.500000      ---        ---    

            >>> print(L + 3)
            Traceback (most recent call last):
            ...
            AttributeError: 'int' object has no attribute 'fillComplete'
        """

        if other is 0:
            return self
        else:
            return self._add(other)

    __radd__ = __add__

    def __sub__(self, other):
        if other is 0:
            return self
        else:
            return self._add(other, sign=-1)

    def __mul__(self, other):
        """
        Multiply a sparse matrix by another sparse matrix.

        >>> L1 = _TrilinosMatrixFromShape(rows=3, cols=3)
        >>> L1.addAt((3, 10, numerix.pi, 2.5), (0, 0, 1, 2), (2, 1, 1, 0))
        >>> L2 = _TrilinosIdentityMatrix(size=3)
        >>> L2.addAt((4.38, 12357.2, 1.1), (2, 1, 0), (1, 0, 2))

        >>> tmp = numerix.array(((1.23572000e+05, 2.31400000e+01, 3.00000000e+00),
        ...                      (3.88212887e+04, 3.14159265e+00, 0.00000000e+00),
        ...                      (2.50000000e+00, 0.00000000e+00, 2.75000000e+00)))

        >>> L = (L1 * L2).numpyArray

        >>> print(numerix.allclose(tmp, L)) # doctest: +SERIAL
        True

        or a sparse matrix by a vector

        >>> tmp = numerix.array((29., 6.28318531, 2.5))
        >>> print(numerix.allclose(L1 * numerix.array((1, 2, 3), 'd'), tmp)) # doctest: +SERIAL
        True

        or a vector by a sparse matrix

        >>> tmp = numerix.array((7.5, 16.28318531,  3.))
        >>> print(numerix.allclose(numerix.array((1, 2, 3), 'd') * L1, tmp))  # doctest: +SERIAL
        True


        """
        N = self.matrix.NumMyCols()

        if isinstance(other, _TrilinosMatrix):
            if isinstance(other.matrix, Epetra.RowMatrix):
                self.fillComplete()
                other.fillComplete()

                result = Epetra.CrsMatrix(Epetra.Copy, self.rowMap, 0)

                EpetraExt.Multiply(self.matrix, False, other.matrix, False, result)
                copy = self.copy()
                copy.matrix = result
                return copy
            else:
                raise TypeError

        else:
            shape = numerix.shape(other)
            if shape == ():
                result = self.copy()
                result.matrix.Scale(other)
                return result
            elif shape == (N,):
                self.fillComplete()

                y = Epetra.Vector(self.domainMap, other)
                result = Epetra.Vector(self.rangeMap)
                self.matrix.Multiply(False, y, result)
                return numerix.array(result)
            else:
                raise TypeError

    def __rmul__(self, other):
        if isinstance(numerix.ones(1, 'l'), type(other)):
            self.fillComplete()

            y = Epetra.Vector(self.rangeMap, other)
            result = Epetra.Vector(self.domainMap)
            self.matrix.Multiply(True, y, result)
            return numerix.array(result)
        else:
            return self * other

    @property
    def _shape(self):
        N = self.matrix.NumGlobalRows()
        return (N, N)



    def put(self, vector, id1, id2):
        """
        Put elements of `vector` at positions of the matrix corresponding to (`id1`, `id2`)

            >>> L = _TrilinosMatrixFromShape(rows=3, cols=3)
            >>> L.put((3., 10., numerix.pi, 2.5), (0, 0, 1, 2), (2, 1, 1, 0))
            >>> print(L)
                ---    10.000000   3.000000  
                ---     3.141593      ---    
             2.500000      ---        ---    
        """

        if hasattr(id1, 'dtype') and id1.dtype.name == 'int64':
            id1 = id1.astype('int32')
        if hasattr(id2, 'dtype') and id2.dtype.name == 'int64':
            id2 = id2.astype('int32')

        if self.matrix.Filled():
            if self.matrix.ReplaceGlobalValues(id1, id2, vector) != 0:
                import warnings
                warnings.warn("ReplaceGlobalValues returned error code in put",
                              UserWarning, stacklevel=2)
                # Possible different algorithm, to guarantee success:
                #
                # Make a new matrix,
                # Use addAt to put the values in it,
                # Use replaceGlobalValues in the original matrix to zero out the terms
                # And add the old one into the new one,
                # Replace the old one.
                #
                # Would incur performance costs, and since FiPy does not use
                # this function in such a way as would generate these errors,
                # I have not implemented the change.

        else:

            # This guarantees that it will actually replace the values that are there,
            # if there are any
            if self.matrix.NumGlobalNonzeros() == 0:
                self.matrix.InsertGlobalValues(id1, id2, vector)
            else:
                self.matrix.InsertGlobalValues(id1, id2, numerix.zeros(len(vector), 'l'))
                self.fillComplete()
                if self.matrix.ReplaceGlobalValues(id1, id2, vector) != 0:
                    import warnings
                    warnings.warn("ReplaceGlobalValues returned error code in put",
                                  UserWarning, stacklevel=2)
                    # Possible different algorithm, to guarantee that it does not fail:
                    #
                    # Make a new matrix,
                    # Use addAt to put the values in it,
                    # Use replaceGlobalValues in the original matrix to zero out the terms
                    # And add the old one into the new one,
                    # Replace the old one.
                    #
                    # Would incur performance costs, and since FiPy does not use
                    # this function in such a way as would generate these errors,
                    # I have not implemented the change.




    def putDiagonal(self, vector):
        """
        Put elements of `vector` along diagonal of matrix

            >>> L = _TrilinosMatrixFromShape(rows=3, cols=3)
            >>> L.putDiagonal((3., 10., numerix.pi))
            >>> print(L)
             3.000000      ---        ---    
                ---    10.000000      ---    
                ---        ---     3.141593  
            >>> L.putDiagonal((10., 3.))
            >>> print(L)
            10.000000      ---        ---    
                ---     3.000000      ---    
                ---        ---     3.141593  
        """

        if type(vector) in [type(1), type(1.)]:
            ids = numerix.arange(self.matrix.NumGlobalRows())
            tmp = numerix.zeros((self.matrix.NumGlobalRows), 'd')
            tmp[:] = vector
            if ids.dtype.name == 'int64':
                ids = ids.astype('int32')
            self.put(tmp, ids, ids)
        else:
            ids = numerix.arange(len(vector))
            if ids.dtype.name == 'int64':
                ids = ids.astype('int32')
            self.put(vector, ids, ids)

    def take(self, id1, id2):
        import warnings
        warnings.warn("Trying to take from a Trilinos Matrix. That doesn't work.",
                      UserWarning, stacklevel=2)
        raise TypeError

    def takeDiagonal(self):
        self.fillComplete()

        result = Epetra.Vector(self.rangeMap)
        self.matrix.ExtractDiagonalCopy(result)

        return result

    def addAt(self, vector, id1, id2):
        """
        Add elements of `vector` to the positions in the matrix corresponding to (`id1`,`id2`)

            >>> L = _TrilinosMatrixFromShape(rows=3, cols=3)
            >>> L.addAt((3., 10., numerix.pi, 2.5), (0, 0, 1, 2), (2, 1, 1, 0))
            >>> L.addAt((1.73, 2.2, 8.4, 3.9, 1.23), (1, 2, 0, 0, 1), (2, 2, 0, 0, 2))
            >>> print(L)
            12.300000  10.000000   3.000000  
                ---     3.141593   2.960000  
             2.500000      ---     2.200000  
        """

        ## This was added as it seems that trilinos does not like int64 arrays
        if hasattr(id1, 'astype') and id1.dtype.name == 'int64':
            id1 = id1.astype('int32')
        if hasattr(id2, 'astype') and id2.dtype.name == 'int64':
            id2 = id2.astype('int32')

        if not self.matrix.Filled():
            err = self.matrix.InsertGlobalValues(id1, id2, vector)
            if err < 0:
                raise RuntimeError("Processor %d, error code %d" \
                  % (self.comm.MyPID(), err))
        else:
            if self.matrix.SumIntoGlobalValues(id1, id2, vector) != 0:
                import warnings
                warnings.warn("Summing into unfilled matrix returned error code",
                              UserWarning, stacklevel=2)
                # Possible change to this part of the code to do the following:
                #
                # Make a new matrix,
                # Use addAt to put the values in it
                # Add the old one into the new one
                # Replace the old one.
                #
                # Would incur performance costs, and since FiPy does not use
                # this function in such a way as would generate these errors,
                # I have not implemented the change.


    def addAtDiagonal(self, vector):
        if isinstance(vector, (int, float)):

            if hasattr(self.matrix, 'GetMyRows'):
                Nrows = self.matrix.GetMyRows()
            else:
                Nrows = self.matrix.NumMyRows()

            ids = numerix.arange(Nrows)
            tmp = numerix.zeros((Nrows,), 'd')
            tmp[:] = vector
            self.addAt(tmp, ids, ids)
        else:
            ids = numerix.arange(len(vector))
            self.addAt(vector, ids, ids)

    def exportMmf(self, filename):
        """
        Exports the matrix to a Matrix Market file of the given `filename`.
        """
        self.fillComplete()
        EpetraExt.RowMatrixToMatrixMarketFile(text_to_native_str(filename), self.matrix)

    @property
    def numpyArray(self):
        import tempfile
        import os
        from scipy.io import mmio
        from fipy.tools import parallelComm

        if parallelComm.procID == 0:
            (_, mtxName) = tempfile.mkstemp(suffix='.mtx')
        else:
            mtxName = None

        mtxName = parallelComm.bcast(mtxName)

        self.exportMmf(mtxName)

        parallelComm.Barrier()
        mtx = mmio.mmread(mtxName)
        parallelComm.Barrier()

        if parallelComm.procID == 0:
            os.remove(mtxName)

        coo = mtx.tocoo()
        trilinosMatrix = self.matrix
        numpyArray = numerix.zeros((trilinosMatrix.NumGlobalRows(),
                                    trilinosMatrix.NumGlobalCols()), 'd')
        numpyArray[coo.row, coo.col] = coo.data
        return numpyArray

    def _getDistributedMatrix(self):
        """
        Returns an equivalent Trilinos matrix, but redistributed evenly over
        all processors.
        """
        if self.comm.NumProc() == 1:
            return self.matrix
            # No redistribution necessary in serial mode
        else:
##            self._matrix.GlobalAssemble()
            totalElements = self.matrix.NumGlobalRows()

            # Epetra.Map(numGlobalElements, indexBase, comm)
            # implicit number of elements per processor
            DistributedMap = Epetra.Map(totalElements, 0, self.comm)
            RootToDist = Epetra.Import(DistributedMap, self.rangeMap)

            DistMatrix = Epetra.CrsMatrix(Epetra.Copy, DistributedMap, (self.bandwidth*3) // 2)

            DistMatrix.Import(self.matrix, RootToDist, Epetra.Insert)

            return DistMatrix

    def fillComplete(self):
        if not self.matrix.Filled():
            self.matrix.FillComplete(self.domainMap, self.rangeMap)

    def finalize(self):
        self.fillComplete()
        self.matrix.OptimizeStorage()

    @property
    def CSR(self):
        """The Compact Sparse Row description of the local matrix

        Returns
        -------
        ptrs : array_like of int
            Locations in `cols` and `data` vectors that start a row,
            terminated with len(data) + 1
        cols : array_like of int
            Sequence of non-sparse column indices.
        data : array_like of float
            Sequence of non-sparse values.

        Examples
        --------

        >>> L = _TrilinosMatrixFromShape(rows=3, cols=3, bandwidth=3)
        >>> L.put([3.,10.,numerix.pi,2.5], [0,0,1,2], [2,1,1,0])
        >>> L.addAt([1.73,2.2,8.4,3.9,1.23], [1,2,0,0,1], [2,2,0,0,2])
        >>> ptrs, cols, data = L.CSR
        >>> print(numerix.asarray(ptrs))
        [0 3 5 7]
        >>> print(numerix.asarray(cols))
        [0 1 2 1 2 0 2]
        >>> print(numerix.asarray(data))
        [ 12.3         10.           3.           3.14159265   2.96
           2.5          2.2       ]
        """
        rows, lildata = self.LIL

        ptrs = [0] + [len(row) for row in rows if len(row)]
        ptrs = list(numerix.cumsum(ptrs))
        cols = [col for row in rows for col in row]
        data = [datum for row in lildata for datum in row]

        return ptrs, cols, data

    @property
    def LIL(self):
        """The List of Lists description of the local matrix

        Returns
        -------
        rows : list of sequence of int
            List of non-sparse column indices on each row.
        data : list of sequence of float
            List of non-sparse values on each row.

        Examples
        --------

        >>> L = _TrilinosMatrixFromShape(rows=3, cols=3, bandwidth=3)
        >>> L.put([3.,10.,numerix.pi,2.5], [0,0,1,2], [2,1,1,0])
        >>> L.addAt([1.73,2.2,8.4,3.9,1.23], [1,2,0,0,1], [2,2,0,0,2])
        >>> rows, data = L.LIL
        >>> from scipy.stats.mstats import argstoarray # doctest: +SCIPY
        >>> print(argstoarray(*rows)) # doctest: +SCIPY
        [[0.0 1.0 2.0]
         [1.0 2.0 --]
         [0.0 2.0 --]]
        >>> print(argstoarray(*data)) # doctest: +SCIPY
        [[12.3 10.0 3.0]
         [3.141592653589793 2.96 --]
         [2.5 2.2 --]]
        """
        self.finalize()

        rows = []
        data = []
        for row in self.matrix.RowMap().MyGlobalElements():
            rowdata, rowcols = self.matrix.ExtractGlobalRowCopy(row)
            rows.append(rowcols)
            data.append(rowdata)

        return rows, data

    @property
    def T(self):
        """Transpose matrix

        After https://github.com/trilinos/Trilinos_tutorial/wiki/Tpetra_Exercises_Advanced_CrsMatrix_ExplicitTranspose#how-to-compute-the-transpose-of-a-crsmatrix

        Sadly, `EpetraExt.RowMatrix_Transpose` isn't exposed in PyTrilinos.

        Returns
        -------
        ~fipy.matrices.trilinosMatrix._TrilinosMatrix

        Examples
        --------

        >>> import fipy as fp

        >>> mesh = fp.Grid1D(nx=10)
        >>> ids = fp.CellVariable(mesh=mesh, value=mesh._globalOverlappingCellIDs)

        >>> mat = _TrilinosColMeshMatrix(mesh=mesh, rows=1)
        >>> mat.put(vector=ids.value,
        ...         id1=[fp.parallelComm.procID] * mesh.numberOfCells,
        ...         id2=mesh._localOverlappingCellIDs,
        ...         overlapping=True)

        >>> print(mat.T.numpyArray) # doctest: +SERIAL
        [[ 0.]
         [ 1.]
         [ 2.]
         [ 3.]
         [ 4.]
         [ 5.]
         [ 6.]
         [ 7.]
         [ 8.]
         [ 9.]]
        >>> print(mat.T.numpyArray) # doctest: +PARALLEL_2
        [[ 0.  0.]
         [ 1.  0.]
         [ 2.  0.]
         [ 3.  3.]
         [ 4.  4.]
         [ 5.  5.]
         [ 6.  6.]
         [ 0.  7.]
         [ 0.  8.]
         [ 0.  9.]]
        """
        self.finalize()

        # transpose the maps
        rowMap = self.matrix.ColMap()
        colMap = self.matrix.RowMap()
        domainMap = self.matrix.RangeMap()
        rangeMap = self.matrix.DomainMap()

        # 2. Create a new CrsMatrix AT, with A's column Map as AT's row Map.
        numEntriesPerRow = self.matrix.NumGlobalNonzeros() // self.matrix.NumGlobalCols()
        A_T = Epetra.CrsMatrix(Epetra.Copy,
                               rowMap,
                               # RuntimeError: InsertMyValues cannot be
                               # called on Epetra_CrsMatrix that does not
                               # have a column map
                               colMap,
                               numEntriesPerRow)

        # 1. On each process, extract and transpose the local data.
        for irow in range(self.matrix.NumMyRows()):
            # "Returns a two-tuple of numpy arrays of the same size; the first is
            # an array of integers that represent the nonzero columns on the
            # matrix"
            #
            # No, it's not. Returns values *then* indices
            val, jcol = self.matrix.ExtractMyRowCopy(irow)

            # 3. Insert the transposed local data into AT.
            # Note that the column Map owns all of AT's local data.  This
            # means that Step 3 can use local indices.
            A_T.InsertMyValues(jcol, [irow] * len(jcol), val)

        # 4.  Fill complete AT (you'll need to supply the domain and range
        # Map, because the row Map of AT is not one to one in general).
        A_T.FillComplete(domainMap, rangeMap)

        # 5.  If desired, redistribute AT (using an Export) to have a
        # one-to-one row Map.
        exporter = Epetra.Export(rowMap, rangeMap)
        A_T_bis = Epetra.CrsMatrix(Epetra.Copy,
                                   rangeMap,
                                   numEntriesPerRow)
        A_T_bis.Export(A_T, exporter, Epetra.Insert)
        A_T_bis.FillComplete(domainMap, rangeMap)

        return _TrilinosMatrix(matrix=A_T_bis)

class _TrilinosMatrixFromShape(_TrilinosMatrix):
    def __init__(self, rows, cols, bandwidth=1, sizeHint=None, matrix=None):
        """Instantiates and wraps an `Epetra.CrsMatrix`

        Parameters
        ----------
        rows : int
            The number of matrix rows
        cols : int
            The number of matrix columns
        bandwidth : int
            The proposed band width of the matrix.
        sizeHint : int
            Estimate of the number of non-zeros
        matrix : ~PyTrilinos.Epetra.CrsMatrix
            Pre-assembled Trilinos matrix to use for storage.
        """
        self.rows = rows
        self.cols = cols

        size = max(rows, cols)
        if sizeHint is not None and bandwidth == 0:
            bandwidth = (sizeHint + size - 1) // (size or 1)
        else:
            bandwidth = bandwidth

        if matrix is None:
            # On Mar 26, 2011, at 11:04 AM, Williams, Alan B <william@sandia.gov> wrote:
            #
            # In almost every case, it is advisable to construct epetra
            # matrices with only a row-map and allow the column-map to be
            # generated internally.  Especially if you have a rectangular
            # matrix, which is even trickier to get correct.
            matrix = Epetra.CrsMatrix(Epetra.Copy, self.rowMap, (bandwidth*3)//2)

        # Leave extra bandwidth, to handle multiple insertions into the
        # same spot. It's memory-inefficient, but it'll get cleaned up when
        # FillComplete is called, and according to the Trilinos devs the
        # performance boost will be worth it.

        super(_TrilinosMatrixFromShape, self).__init__(matrix=matrix,
                                                       bandwidth=bandwidth)

    @property
    def rowMap(self):
        comm = Epetra.SerialComm()
        # Matrix building gets done on all processors
        # Epetra.Map(numGlobalElements, numMyElements, indexBase, comm)
        # If NumGlobalElements = NumMyElements (and not equal to zero) the map is
        # defined to be a local replicated map
        return Epetra.Map(self.rows, self.rows, 0, comm)

    @property
    def colMap(self):
        comm = Epetra.SerialComm()
        # Matrix building gets done on all processors
        # Epetra.Map(numGlobalElements, numMyElements, indexBase, comm)
        # If NumGlobalElements = NumMyElements (and not equal to zero) the map is
        # defined to be a local replicated map
        return Epetra.Map(self.cols, self.cols, 0, comm)

class _TrilinosBaseMeshMatrix(_TrilinosMatrixFromShape):
    def __init__(self, mesh, rows, cols, m2m, bandwidth=0, sizeHint=None, matrix=None):
        """Creates a `_TrilinosMatrixFromShape` associated with a `Mesh`.

        Parameters
        ----------
        mesh : ~fipy.meshes.mesh.Mesh
            The `Mesh` to assemble the matrix for.
        rows : int
            The number of local matrix rows.
        cols : int
            The number of local matrix columns.
        m2m : ~fipy.matrices.sparseMatrix._Mesh2Matrix
            Object to convert between mesh coordinates and matrix coordinates.
        bandwidth : int
            The proposed band width of the matrix.
        sizeHint : int
            Estimate of the number of non-zeros.
        matrix : ~PyTrilinos.Epetra.CrsMatrix
            Pre-assembled Trilinos matrix to use for storage.
        """
        self.mesh = mesh
        self._m2m = m2m

        super(_TrilinosBaseMeshMatrix, self).__init__(rows=rows,
                                                      cols=cols,
                                                      bandwidth=bandwidth,
                                                      sizeHint=sizeHint,
                                                      matrix=matrix)

    @property
    def rowMap(self):
        comm = self.mesh.communicator.epetra_comm
        # Epetra.Map(numGlobalElements, indexBase, comm)
        # If NumGlobalElements = -1 and NumMyElements is passed in then
        # NumGlobalElements will be computed as the sum of NumMyElements across
        # all processors.
        return Epetra.Map(-1, self.rows, 0, comm)

    @property
    def colMap(self):
        comm = self.mesh.communicator.epetra_comm
        # Epetra.Map(numGlobalElements, indexBase, comm)
        # If NumGlobalElements = -1 and NumMyElements is passed in then
        # NumGlobalElements will be computed as the sum of NumMyElements across
        # all processors.
        return Epetra.Map(-1, self.cols, 0, comm)

    def copy(self):
        tmp = super(_TrilinosBaseMeshMatrix, self).copy()
        copy = self.__class__(mesh=self.mesh, bandwidth=self.bandwidth)
        copy.matrix = tmp.matrix
        return copy

    @staticmethod
    def _ArrayAsEpetraVector(arr, emap):
        """Convert numpy `ndarray` to `Epetra.Vector`

        Parameters
        ----------
        arr : array_like
            Input data, in any form that can be converted to an array.
        emap : Epetra.Map
            The map to distribute across.

        Returns
        -------
        vec : Epetra.Vector
            Epetra interpretation of `arr`.
        dtype : data-type
            The original data type of `arr`, as :class:`Epetra.Vector` can
            only handle doubles.
        """
        return (Epetra.Vector(emap, numerix.asarray(arr, dtype=float)),
                arr.dtype)

    @staticmethod
    def _EpetraVectorAsArray(vec, dtype=None, shape=None):
        """Convert `Epetra.Vector` to numpy `ndarray`

        Parameters
        ----------
        vec : Epetra.Vector
            Input data.
        dtype : data-type, optional
            By default, the data-type is float, as :class:`Epetra.Vector`
            can only hold doubles.
        shape : tuple
            Desired shape of result (default: None leaves shape along)

        Returns
        -------
        out : ndarray
            Array interpretation of `vec`.
        """
        arr = numerix.asarray(vec, dtype=dtype)
        if shape is not None:
            arr = numerix.reshape(arr, shape=shape)
        return arr

    def _getGhostedValues(self, var):
        """Obtain current ghost values from across processes

        Returns
        -------
        ndarray
            Ghosted values
        """
        s = (var.mesh._localNonOverlappingCellIDs,)

        ## The following conditional is required because empty indexing is
        ## not altogether functional.  This numpy.empty((0,))[[]] and this
        ## numpy.empty((0,))[...,[]] both work, but this numpy.empty((3,
        ## 0))[...,[]] is broken.
        if var.shape[-1] != 0:
            s = (Ellipsis,) + s

        (nonOverlappingVector,
         original_dtype) = self._ArrayAsEpetraVector(arr=var[s].ravel(),
                                                     emap=self.domainMap)

        overlappingVector = Epetra.Vector(self.colMap)
        overlappingVector.Import(nonOverlappingVector,
                                 Epetra.Import(self.colMap, self.domainMap),
                                 Epetra.Insert)

        return self._EpetraVectorAsArray(overlappingVector,
                                         dtype=original_dtype,
                                         shape=var.shape)

    def put(self, vector, id1, id2, overlapping=False):
        """Insert local overlapping values and coordinates into global

        Parameters
        ----------
        vector : array_like
            The overlapping values to insert.
        id1 : array_like
            The local overlapping row indices.
        id2 : array_like
            The local overlapping column indices.
        overlapping : bool
            Whether to insert ghosted values or not (default False)
        """
        vector, id1, id2 = self._m2m.globalVectorAndIDs(vector, id1, id2, overlapping)
        super(_TrilinosBaseMeshMatrix, self).put(vector=vector, id1=id1, id2=id2)

    def addAt(self, vector, id1, id2, overlapping=False):
        """Accumulate local overlapping values and coordinates into global

        Parameters
        ----------
        vector : array_like
            The overlapping values to insert.
        id1 : array_like
            The local overlapping row indices.
        id2 : array_like
            The local overlapping column indices.
        overlapping : bool
            Whether to add ghosted values or not (default False)
        """
        vector, id1, id2 = self._m2m.globalVectorAndIDs(vector, id1, id2, overlapping)
        super(_TrilinosBaseMeshMatrix, self).addAt(vector=vector, id1=id1, id2=id2)

class _TrilinosRowMeshMatrix(_TrilinosBaseMeshMatrix):
    def __init__(self, mesh, cols, numberOfEquations=1, bandwidth=0,
                 sizeHint=None, matrix=None, m2m=None):
        """Creates a `_TrilinosMatrixFromShape` with rows associated with equations.

        Parameters
        ----------
        mesh : ~fipy.meshes.mesh.Mesh
            The `Mesh` to assemble the matrix for.
        cols : int
            The number of matrix columns
        numberOfEquations : int
            The local rows of the matrix are determined by
            `numberOfEquations * mesh._localNonOverlappingCellIDs`.
        bandwidth : int
            The proposed band width of the matrix.
        sizeHint : int
            Estimate of the number of non-zeros.
        matrix : ~PyTrilinos.Epetra.CrsMatrix
            Pre-assembled Trilinos matrix to use for storage.
        m2m : ~fipy.matrices.sparseMatrix._RowMesh2Matrix
            Object to convert between mesh coordinates and matrix coordinates.
        """
        if m2m is None:
            m2m = _RowMesh2Matrix(mesh=mesh, matrix=self,
                                  numberOfEquations=numberOfEquations)

        rows = numberOfEquations * mesh.globalNumberOfCells
        super(_TrilinosRowMeshMatrix, self).__init__(mesh=mesh,
                                                     rows=rows,
                                                     cols=cols,
                                                     m2m=m2m,
                                                     bandwidth=bandwidth,
                                                     sizeHint=sizeHint,
                                                     matrix=matrix)

    @property
    def rowMap(self):
        comm = self.mesh.communicator.epetra_comm
        # Epetra.Map(numGlobalElements, myGlobalElements, indexBase, comm)
        # Specify -1 to have the constructor compute the number of global elements.
        return Epetra.Map(-1, list(self._m2m.globalNonOverlappingRowIDs), 0, comm)

    @property
    def domainMap(self):
        """Distribution of input vector `x` in `y = A*x`

        The domain Map must be one-to-one: that is, each global index in
        the Map must be uniquely owned by one and only one process.
        """
        return self.rowMap

class _TrilinosColMeshMatrix(_TrilinosBaseMeshMatrix):
    def __init__(self, mesh, rows, numberOfVariables=1, bandwidth=0,
                 sizeHint=None, matrix=None):
        """Creates a `_TrilinosMatrixFromShape` with columns associated with solution variables.

        Parameters
        ----------
        mesh : ~fipy.meshes.mesh.Mesh
            The `Mesh` to assemble the matrix for.
        rows : int
            The number of matrix rows.
        numberOfVariables : int
            The local columns of the matrix are determined by
            `numberOfVariables * mesh.globalNumberOfCells`.
        bandwidth : int
            The proposed band width of the matrix.
        sizeHint : int
            Estimate of the number of non-zeros.
        matrix : ~PyTrilinos.Epetra.CrsMatrix
            Pre-assembled Trilinos matrix to use for storage.
        """
        m2m = _ColMesh2Matrix(mesh=mesh, matrix=self,
                              numberOfVariables=numberOfVariables)

        cols = numberOfVariables * mesh.globalNumberOfCells
        super(_TrilinosColMeshMatrix, self).__init__(mesh=mesh,
                                                     rows=rows,
                                                     cols=cols,
                                                     m2m=m2m,
                                                     bandwidth=bandwidth,
                                                     sizeHint=sizeHint,
                                                     matrix=matrix)

    @property
    def colMap(self):
        comm = self.mesh.communicator.epetra_comm
        # Epetra.Map(numGlobalElements, myGlobalElements, indexBase, comm)
        # Specify -1 to have the constructor compute the number of global elements.
        return Epetra.Map(-1, list(self._m2m.globalOverlappingColIDs), 0, comm)

    @property
    def domainMap(self):
        """Distribution of input vector `x` in `y = A*x`

        The domain Map must be one-to-one: that is, each global index in
        the Map must be uniquely owned by one and only one process.
        """
        comm = self.mesh.communicator.epetra_comm
        # Epetra.Map(numGlobalElements, myGlobalElements, indexBase, comm)
        # Specify -1 to have the constructor compute the number of global elements.
        return Epetra.Map(-1, list(self._m2m.globalNonOverlappingColIDs), 0, comm)

class _TrilinosMeshMatrix(_TrilinosRowMeshMatrix):
    def __init__(self, mesh, numberOfVariables=1, numberOfEquations=1,
                 bandwidth=0, sizeHint=None, matrix=None):
        """Creates a `_TrilinosRowMeshMatrix` associated with equations and variables.

        Parameters
        ----------
        mesh : ~fipy.meshes.mesh.Mesh
            The `Mesh` to assemble the matrix for.
        numberOfVariables : int
            The local columns of the matrix are determined by
            `numberOfVariables * len(mesh._localNonOverlappingCellIDs)`.
        numberOfEquations : int
            The local rows of the matrix are determined by
            `numberOfEquations * len(mesh._localNonOverlappingCellIDs)`.
        bandwidth : int
            The proposed band width of the matrix.
        sizeHint : int
            Estimate of the number of non-zeros
        matrix : ~PyTrilinos.Epetra.CrsMatrix
            Pre-assembled Trilinos matrix to use for storage.
        """
        m2m = _RowColMesh2Matrix(mesh=mesh, matrix=self,
                                 numberOfVariables=numberOfVariables,
                                 numberOfEquations=numberOfEquations)

        cols = numberOfVariables * mesh.globalNumberOfCells
        super(_TrilinosMeshMatrix, self).__init__(mesh=mesh,
                                                  cols=cols,
                                                  numberOfEquations=numberOfEquations,
                                                  bandwidth=bandwidth,
                                                  sizeHint=sizeHint,
                                                  matrix=matrix,
                                                  m2m=m2m)

    @property
    def colMap(self):
        comm = self.mesh.communicator.epetra_comm
        # Epetra.Map(numGlobalElements, myGlobalElements, indexBase, comm)
        # Specify -1 to have the constructor compute the number of global elements.
        return Epetra.Map(-1, list(self._m2m.globalOverlappingColIDs), 0, comm)

    def asTrilinosMeshMatrix(self):
        self.finalize()
        return self

    def flush(self):
        pass

    def _getMatrixProperty(self):
        if not hasattr(self, '_matrix'):
            self._matrix = _TrilinosMeshMatrix(self.mesh,
                                               bandwidth=self.bandwidth,
                                               numberOfVariables=self._m2m.numberOfVariables,
                                               numberOfEquations=self._m2m.numberOfEquations).matrix
        return super(_TrilinosMeshMatrix, self).matrix

    matrix = property(_getMatrixProperty, _TrilinosMatrixFromShape._setMatrix)

    def takeDiagonal(self):
        nonoverlapping_result = _TrilinosMatrixFromShape.takeDiagonal(self)

        overlapping_result = Epetra.Vector(self.colMap)
        overlapping_result.Import(nonoverlapping_result,
                                  Epetra.Import(self.colMap,
                                                self.domainMap),
                                  Epetra.Insert)

        return overlapping_result

    def __mul__(self, other):
        """
        Multiply a sparse matrix by another sparse matrix.

            >>> L1 = _TrilinosMatrixFromShape(rows=3, cols=3)
            >>> L1.addAt((3, 10, numerix.pi, 2.5), (0, 0, 1, 2), (2, 1, 1, 0))
            >>> L2 = _TrilinosIdentityMatrix(size=3)
            >>> L2.addAt((4.38, 12357.2, 1.1), (2, 1, 0), (1, 0, 2))

            >>> tmp = numerix.array(((1.23572000e+05, 2.31400000e+01, 3.00000000e+00),
            ...                      (3.88212887e+04, 3.14159265e+00, 0.00000000e+00),
            ...                      (2.50000000e+00, 0.00000000e+00, 2.75000000e+00)))

            >>> L = (L1 * L2).numpyArray

            >>> print(numerix.allclose(tmp, L))
            True

        or a sparse matrix by a vector

            >>> tmp = numerix.array((29., 6.28318531, 2.5))
            >>> print(numerix.allclose(L1 * numerix.array((1, 2, 3), 'd'), tmp)) # doctest: +SERIAL
            True

        or a vector by a sparse matrix

            >>> tmp = numerix.array((7.5, 16.28318531,  3.))
            >>> numerix.allclose(numerix.array((1, 2, 3), 'd') * L1, tmp) # doctest: +SERIAL
            True

        Should be able to multiply an overlapping value obtained from a
        `CellVariable`. This is required to make the `--no-pysparse` flag
        work correctly.

            >>> from fipy import *
            >>> m = Grid1D(nx=6)
            >>> v0 = CellVariable(mesh=m, value=numerix.arange(m.globalNumberOfCells, dtype=float))
            >>> v1 = CellVariable(mesh=m, value=_TrilinosIdentityMeshMatrix(mesh=m) * v0.value)
            >>> print(numerix.allclose(v0, v1))
            True

        """
        self.fillComplete()

        N = self.matrix.NumMyCols()

        if isinstance(other, _TrilinosMatrix):
            return super(_TrilinosMeshMatrix, self).__mul__(other=other)
        else:
            shape = numerix.shape(other)


            if shape == ():
                result = self.copy()
                result.matrix.Scale(other)
                return result
            else:

                if isinstance(other, Epetra.Vector):
                    other_map = other.Map()
                else:
                    other_map = self.colMap

                if other_map.SameAs(self.colMap):
                    localNonOverlappingColIDs = self._m2m.localNonOverlappingColIDs

                    other = Epetra.Vector(self.domainMap,
                                          other[localNonOverlappingColIDs])

                if other.Map().SameAs(self.matrix.DomainMap()):
                    nonoverlapping_result = Epetra.Vector(self.rangeMap)
                    self.matrix.Multiply(False, other, nonoverlapping_result)

                    if other_map.SameAs(self.colMap):
                        overlapping_result = Epetra.Vector(self.colMap)
                        overlapping_result.Import(nonoverlapping_result,
                                                  Epetra.Import(self.colMap,
                                                                self.domainMap),
                                                  Epetra.Insert)

                        return overlapping_result
                    else:
                        return nonoverlapping_result

                else:
                    raise TypeError("%s: %s != (%d,)" % (self.__class__, str(shape), N))

class _TrilinosIdentityMatrix(_TrilinosMatrixFromShape):
    """
    Represents a sparse identity matrix for Trilinos.
    """
    def __init__(self, size):
        """
        Create a sparse matrix with `1` in the diagonal

            >>> print(_TrilinosIdentityMatrix(size=3))
             1.000000      ---        ---    
                ---     1.000000      ---    
                ---        ---     1.000000  
        """
        _TrilinosMatrixFromShape.__init__(self, rows=size, cols=size, bandwidth=1)
        ids = numerix.arange(size)
        self.addAt(numerix.ones(size, 'l'), ids, ids)

class _TrilinosIdentityMeshMatrix(_TrilinosMeshMatrix):
    def __init__(self, mesh):
        """
        Create a sparse matrix associated with a `Mesh` with `1` in the diagonal

            >>> from fipy import Grid1D
            >>> mesh = Grid1D(nx=3)
            >>> print(_TrilinosIdentityMeshMatrix(mesh=mesh))
             1.000000      ---        ---    
                ---     1.000000      ---    
                ---        ---     1.000000  
        """
        _TrilinosMeshMatrix.__init__(self, mesh=mesh, bandwidth=1)
        size = mesh.numberOfCells
        ids = numerix.arange(size)
        self.addAt(numerix.ones(size, 'l'), ids, ids)

class _TrilinosMeshMatrixKeepStencil(_TrilinosMeshMatrix):

    def _getStencil(self, id1, id2):
        if not hasattr(self, 'stencil'):
            self.stencil = super(_TrilinosMeshMatrixKeepStencil, self)._getStencil(self, id1, id2)

        return self.stencil

    def flush(self, cacheStencil=False):
        """Deletes the matrix but maintains the stencil used
        `_globalNonOverlapping()` in as it can be expensive to construct.

        Parameters
        ----------
        cacheStencil : bool
            Whether to determine whether to keep the stencil (tuple of IDs
            and a mask) even after deleting the matrix.

        """

        del self._matrix
        if not cacheStencil:
            del self.stencil

def _test():
    import fipy.tests.doctestPlus
    return fipy.tests.doctestPlus.testmod()

if __name__ == "__main__":
    _test()
