from __future__ import annotations

import operator
import warnings

import numpy

try:
    import scipy.sparse
    _scipy_available = True
except ImportError:
    _scipy_available = False

import cupy
from cupy.cuda import runtime
from cupyx.scipy.sparse import _base
from cupyx.scipy.sparse import _compressed
from cupyx.scipy.sparse import SparseEfficiencyWarning
from cupyx.scipy.sparse import _sputils
from cupyx.scipy.sparse import _util


class _csr_base(_compressed._compressed_sparse_matrix):
    """CSR format base (shared by csr_matrix and csr_array)."""

    format = 'csr'
    # Index of the major axis in shape: rows for CSR.
    _major_axis = 0

    def get(self, stream=None):
        """Returns a copy of the array on host memory.

        Args:
            stream (cupy.cuda.Stream): CUDA stream object. If it is given, the
                copy runs asynchronously. Otherwise, the copy is synchronous.

        Returns:
            scipy.sparse.csr_array or scipy.sparse.csr_matrix:
                Copy of the array on host memory.

        """
        if not _scipy_available:
            raise RuntimeError('scipy is not available')
        data = self.data.get(stream)
        indices = self.indices.get(stream)
        indptr = self.indptr.get(stream)
        if isinstance(self, _base.sparray):
            sp_cls = scipy.sparse.csr_array
        else:
            sp_cls = scipy.sparse.csr_matrix
        return sp_cls((data, indices, indptr), shape=self._shape)

    def _convert_dense(self, x):
        m = dense2csr(x)
        return m.data, m.indices, m.indptr

    def _swap(self, x, y):
        return (x, y)

    def _add_sparse(self, other, alpha, beta):
        from cupyx import cusparse

        self.sum_duplicates()
        other = other.tocsr()
        other.sum_duplicates()
        if cusparse.check_availability('csrgeam2'):
            csrgeam = cusparse.csrgeam2
        elif cusparse.check_availability('csrgeam'):
            csrgeam = cusparse.csrgeam
        else:
            raise NotImplementedError
        return self._as_csr_type(csrgeam(self, other, alpha, beta))

    def _comparison(self, other, op, op_name):
        if self.ndim != 2:
            # 1-D: compare on the (1, N) backing (2-D result kept, 1-D
            # squeezed -- see _run_1d_backing_op).
            return self._run_1d_backing_op(
                other, lambda a, o: a._comparison(o, op, op_name))
        cls = type(self)
        if _util.isscalarlike(other):
            data = cupy.asarray(other, dtype=self.dtype).reshape(1)
            if numpy.isnan(data[0]):
                if op_name == '_ne_':
                    return cls(cupy.ones(self.shape, dtype=numpy.bool_))
                else:
                    return cls(self.shape, dtype=numpy.bool_)
            scalar = data[0]
            # Fast path: when op(0, scalar) is False, unstored
            # entries (implicit zeros) all produce False.  The result
            # is sparse -- only stored entries can contribute True.
            # This avoids the O(m*n) expansion in binopt_csr.
            try:
                zero_cmp = bool(op(self.dtype.type(0), scalar))
            except TypeError:
                # Complex dtypes don't support < > in Python
                zero_cmp = True  # fall through to binopt
            if not zero_cmp:
                self.sum_duplicates()
                new_data = op(self.data, scalar)
                # Keep only True entries (filter explicit False).  The
                # post-``sum_duplicates`` structure is canonical
                # (sorted, no duplicates); filtering preserves both
                # invariants, so propagate ``has_canonical_format=True``.
                mask = new_data
                if mask.all():  # synchronize!
                    return cls._from_parts(
                        new_data, self.indices.copy(),
                        self.indptr.copy(), self.shape,
                        has_canonical_format=True)
                from cupyx.cusparse import (
                    _indptr_to_coo, _build_indptr)
                idx_dtype = self.indices.dtype
                rows = _indptr_to_coo(self.indptr, nnz=self.nnz)
                rows = rows[mask]
                cols = self.indices[mask]
                data = new_data[mask]
                M = self._swap(*self.shape)[0]
                indptr = _build_indptr(rows, M, idx_dtype)
                return cls._from_parts(
                    data, cols, indptr, self.shape,
                    has_canonical_format=True)
            # Slow path: op(0, scalar) is True, so unstored entries
            # contribute.  Fall through to binopt_csr which expands
            # to O(m*n).  This is unavoidable when the result is
            # dense.  Match scipy by emitting a warning.
            _inv = {
                '_eq_': ('==', '!='), '_ne_': ('!=', '=='),
                '_lt_': ('<', '>='), '_gt_': ('>', '<='),
                '_le_': ('<=', '>'), '_ge_': ('>=', '<'),
            }
            sym, alt = _inv.get(op_name, (op_name, '!='))
            warnings.warn(
                f'Comparing a sparse matrix with {scalar} using '
                f'{sym} is inefficient. Try using {alt} instead.',
                _base.SparseEfficiencyWarning)
            idx_dtype = self.indices.dtype
            other = cls._from_parts(
                data,
                cupy.zeros((1,), dtype=idx_dtype),
                cupy.arange(2, dtype=idx_dtype),
                (1, 1))
            return binopt_csr(self, other, op_name)
        elif _util.isdense(other):
            return op(self.todense(), other)
        elif _is_csr(other):
            self.sum_duplicates()
            other.sum_duplicates()
            if op_name in ('_ne_', '_lt_', '_gt_'):
                return binopt_csr(self, other, op_name)

            warnings.warn(
                "Comparing sparse matrices using ==, <=, and >= is "
                "inefficient, try using !=, <, or > instead.",
                SparseEfficiencyWarning)
            if op_name == '_eq_':
                opposite_op_name = '_ne_'
            elif op_name == '_le_':
                opposite_op_name = '_gt_'
            elif op_name == '_ge_':
                opposite_op_name = '_lt_'
            res = binopt_csr(self, other, opposite_op_name)
            out = cupy.logical_not(res.toarray())
            return cls(out)
        elif _base.issparse(other):
            # Convert non-CSR sparse operands to CSR (matches scipy).
            return self._comparison(other.tocsr(), op, op_name)
        raise NotImplementedError

    def __eq__(self, other):
        return self._comparison(other, operator.eq, '_eq_')

    def __ne__(self, other):
        return self._comparison(other, operator.ne, '_ne_')

    def __lt__(self, other):
        return self._comparison(other, operator.lt, '_lt_')

    def __gt__(self, other):
        return self._comparison(other, operator.gt, '_gt_')

    def __le__(self, other):
        return self._comparison(other, operator.le, '_le_')

    def __ge__(self, other):
        return self._comparison(other, operator.ge, '_ge_')

    def _matmul_dispatch(self, other):
        from cupyx import cusparse

        if cupy.isscalar(other):
            self.sum_duplicates()
            return self._with_data(self.data * other)
        if self.ndim != 2 or (_base.issparse(other) and other.ndim != 2):
            # A 1-D operand is involved: run on the (1, N)/(N, 1) backing
            # with numpy-style 1-D matmul semantics.
            return self._matmul_1d(other)
        if _is_csr(other):
            self.sum_duplicates()
            other.sum_duplicates()
            # int64: always route through cusparse.spgemm, which has a
            # pure-CuPy fallback when cuSPARSE lacks int64 SpGEMM.
            is_int64 = (self.indices.dtype == cupy.int64
                        or other.indices.dtype == cupy.int64)
            if is_int64 or cusparse.check_availability('spgemm'):
                return self._as_csr_type(cusparse.spgemm(self, other))
            elif cusparse.check_availability('csrgemm2'):
                return self._as_csr_type(cusparse.csrgemm2(self, other))
            elif cusparse.check_availability('csrgemm'):
                return self._as_csr_type(cusparse.csrgemm(self, other))
            else:
                raise RuntimeError('no cuSPARSE spgemm backend available')
        elif _is_csc(other):
            self.sum_duplicates()
            other.sum_duplicates()
            is_int64 = (self.indices.dtype == cupy.int64
                        or other.indices.dtype == cupy.int64)
            if is_int64:
                # csrgemm/csrgemm2 are int32-only; route via spgemm,
                # which has a pure-CuPy fallback for older cuSPARSE.
                b = other.tocsr()
                b.sum_duplicates()
                return self._as_csr_type(cusparse.spgemm(self, b))
            if cusparse.check_availability('csrgemm') and not runtime.is_hip:
                # trans=True is still buggy as of ROCm 4.2.0
                return self._as_csr_type(
                    cusparse.csrgemm(self, other.T, transb=True))
            elif cusparse.check_availability('spgemm'):
                b = other.tocsr()
                b.sum_duplicates()
                return self._as_csr_type(cusparse.spgemm(self, b))
            elif cusparse.check_availability('csrgemm2'):
                b = other.tocsr()
                b.sum_duplicates()
                return self._as_csr_type(cusparse.csrgemm2(self, b))
            else:
                raise RuntimeError('no cuSPARSE spgemm backend available')
        elif _base.issparse(other):
            return self._matmul_dispatch(other.tocsr())
        elif _base.isdense(other):
            if other.ndim == 0:
                self.sum_duplicates()
                return self._with_data(self.data * other)
            elif other.ndim == 1:
                self.sum_duplicates()
                other = cupy.asfortranarray(other)
                if cusparse.check_availability('csrmv'):
                    csrmv = cusparse.csrmv
                elif cusparse.check_availability('spmv'):
                    csrmv = cusparse.spmv
                else:
                    raise RuntimeError('no cuSPARSE spmv backend available')
                return csrmv(self, other)
            elif other.ndim == 2:
                self.sum_duplicates()
                if cusparse.check_availability('csrmm2'):
                    csrmm = cusparse.csrmm2
                elif cusparse.check_availability('spmm'):
                    csrmm = cusparse.spmm
                else:
                    raise RuntimeError('no cuSPARSE spmm backend available')
                return csrmm(self, cupy.asfortranarray(other))
            else:
                raise ValueError('could not interpret dimensions')
        else:
            return NotImplemented

    def __div__(self, other):
        raise NotImplementedError

    def __rdiv__(self, other):
        raise NotImplementedError

    def __truediv__(self, other):
        """Point-wise division by another matrix, vector or scalar"""
        if _util.isscalarlike(other):
            # Pick the result dtype.  scipy upcasts any dtype that's
            # castable to float64 (bool/int*/float16/float32/...) before
            # dividing; mirror that behaviour, and additionally fall
            # back to float64 for any non-cuSPARSE-supported dtype so
            # the result remains usable with cuSPARSE-backed ops.
            # Real-vs-complex follows numpy's natural promotion
            # (``result_type(float64, complex64) -> complex128``).
            dtype = self.dtype
            if dtype == numpy.float32:
                dtype = numpy.float64
            dtype = cupy.result_type(dtype, other)
            if not _sputils.is_sparse_data_dtype(dtype):
                dtype = numpy.dtype(numpy.float64)
            d = cupy.reciprocal(other, dtype=dtype)
            return multiply_by_scalar(self, d)
        if self.ndim != 2:
            # 1-D: divide on the (1, N) backing (2-D result kept, 1-D
            # squeezed -- see _run_1d_backing_op).
            return self._run_1d_backing_op(
                other, lambda a, o: a.__truediv__(o))
        if _util.isdense(other):
            other = cupy.atleast_2d(other)
            other = cupy.broadcast_to(other, self.shape)
            check_shape_for_pointwise_op(self.shape, other.shape)
            ret = self.tocoo()
            ret.data = _cupy_divide_by_dense()(
                ret.data, ret.row, ret.col, ret.shape[1], other)
            return ret
        elif _base.issparse(other):
            # Note: If broadcasting is needed, an exception is raised here for
            # compatibility with SciPy, as SciPy does not support broadcasting
            # in the "sparse / sparse" case.
            check_shape_for_pointwise_op(self.shape, other.shape,
                                         allow_broadcasting=False)
            dtype = numpy.promote_types(self.dtype, other.dtype)
            if dtype.char not in 'FD':
                dtype = numpy.promote_types(numpy.float64, dtype)
            # Note: The following implementation converts two sparse matrices
            # into dense matrices and then performs a point-wise division,
            # which can use lots of memory.
            self_dense = self.todense().astype(dtype, copy=False)
            return self_dense / other.todense()
        return NotImplemented

    def __rtruediv__(self, other):
        return NotImplemented

    # TODO(unno): Implement check_format

    def diagonal(self, k=0):
        self._require_2d('diagonal')
        rows, cols = self.shape
        ylen = min(rows + min(k, 0), cols - max(k, 0))
        if ylen <= 0:
            return cupy.empty(0, dtype=self.dtype)
        self.sum_duplicates()
        y = cupy.empty(ylen, dtype=self.dtype)
        idx_dtype = self.indptr.dtype
        _cupy_csr_diagonal()(idx_dtype.type(k),
                             idx_dtype.type(rows), idx_dtype.type(cols),
                             self.data, self.indptr, self.indices, y)
        return y

    def eliminate_zeros(self):
        """Removes zero entries in place."""
        from cupyx import cusparse

        if self.ndim == 1:
            # Route through the (1, N) backing (rebuilds indptr as
            # [0, nnz]); the direct paths below assume a 2-D shape.
            return self._apply_2d_inplace(lambda m: m.eliminate_zeros())

        if self.indices.dtype == cupy.int64:
            # TODO(eriknw): cuSPARSE--csr2csr_compress doesn't support int64
            mask = self.data != 0
            if mask.all():  # synchronize!
                return
            new_data = self.data[mask]
            new_indices = self.indices[mask]
            nrows = self.shape[0]
            idx_dtype = self.indptr.dtype
            if len(new_data) == 0:
                self.data = new_data
                self.indices = new_indices
                self.indptr = cupy.zeros(nrows + 1, dtype=idx_dtype)
                return
            row_of_each = cusparse._indptr_to_coo(
                self.indptr, nnz=self.nnz)
            kept_rows = row_of_each[mask]
            new_indptr = cusparse._build_indptr(
                kept_rows, nrows, idx_dtype)
            self.data = new_data
            self.indices = new_indices
            self.indptr = new_indptr
            return

        compress = cusparse.csr2csr_compress(self, 0)
        self.data = compress.data
        self.indices = compress.indices
        self.indptr = compress.indptr

    def _maximum_minimum(self, other, cupy_op, op_name, dense_check):
        cls = type(self)
        if _util.isscalarlike(other):
            dtype = cupy.result_type(self.dtype, other)
            other = cupy.asarray(other)
            if dense_check(other):
                other = other.astype(dtype, copy=False)
                new_array = cupy_op(self.todense(), other)
                return cls(new_array)
            else:
                self.sum_duplicates()
                new_data = cupy_op(self.data, other)
                new_data = new_data.astype(dtype, copy=False)
                return cls._from_parts(
                    new_data, self.indices, self.indptr,
                    self.shape,
                    has_canonical_format=getattr(
                        self, '_has_canonical_format', None),
                    has_sorted_indices=getattr(
                        self, '_has_sorted_indices', None))
        if self.ndim != 2:
            # 1-D: run on the (1, N) backing (2-D result kept, 1-D
            # squeezed -- see _run_1d_backing_op).
            return self._run_1d_backing_op(
                other,
                lambda a, o: a._maximum_minimum(
                    o, cupy_op, op_name, dense_check))
        if _util.isdense(other):
            self.sum_duplicates()
            other = cupy.atleast_2d(other)
            return cupy_op(self.todense(), other)
        elif _is_csr(other):
            self.sum_duplicates()
            other.sum_duplicates()
            return binopt_csr(self, other, op_name)
        elif _base.issparse(other):
            # Convert non-CSR sparse operands to CSR (matches scipy).
            return self._maximum_minimum(
                other.tocsr(), cupy_op, op_name, dense_check)
        raise NotImplementedError

    def maximum(self, other):
        return self._maximum_minimum(other, cupy.maximum, '_maximum_',
                                     lambda x: x > 0)

    def minimum(self, other):
        return self._maximum_minimum(other, cupy.minimum, '_minimum_',
                                     lambda x: x < 0)

    def multiply(self, other):
        """Point-wise multiplication by another matrix, vector or scalar"""
        if cupy.isscalar(other):
            return multiply_by_scalar(self, other)
        if self.ndim != 2:
            # 1-D: multiply on the (1, N) backing (2-D result kept, 1-D
            # squeezed -- see _run_1d_backing_op).
            return self._run_1d_backing_op(
                other, lambda a, o: a.multiply(o))
        if _util.isdense(other):
            self.sum_duplicates()
            other = cupy.atleast_2d(other)
            return multiply_by_dense(self, other)
        elif _is_csr(other):
            self.sum_duplicates()
            other.sum_duplicates()
            return multiply_by_csr(self, other, out_cls=type(self))
        elif _base.issparse(other):
            return self.multiply(other.tocsr())
        else:
            msg = 'expected scalar, dense matrix/vector or csr matrix'
            raise TypeError(msg)

    def setdiag(self, values, k=0):
        """Set diagonal or off-diagonal elements of the array.

        Args:
            values: New values of the diagonal elements.  Accepts a
                scalar, list, or 1-D array; any non-cupy input is
                coerced via :func:`cupy.asarray`.
            k (int): Which diagonal to set.  Default 0 (main diagonal).
        """
        self._require_2d('setdiag')
        rows, cols = self.shape
        row_st, col_st = max(0, -k), max(0, k)
        x_len = min(rows - row_st, cols - col_st)
        if x_len <= 0:
            raise ValueError('k exceeds matrix dimensions')
        # Coerce list/scalar/numpy input; matches scipy's
        # ``np.asarray(values)`` in ``_spbase.setdiag``.
        values = cupy.asarray(values, dtype=self.dtype)
        if values.ndim > 1:
            raise ValueError('values must be 0-d or 1-d')
        if values.ndim == 0:
            # broadcast
            x_data = cupy.full((x_len,), values, dtype=self.dtype)
        else:
            x_len = min(x_len, values.size)
            x_data = values[:x_len]
        idx_dtype = self.indices.dtype
        x_indices = cupy.arange(col_st, col_st + x_len, dtype=idx_dtype)
        x_indptr = cupy.zeros((rows + 1,), dtype=idx_dtype)
        x_indptr[row_st:row_st+x_len+1] = cupy.arange(
            x_len+1, dtype=idx_dtype)
        x_indptr[row_st+x_len+1:] = x_len
        # Out-of-place subtraction: ``x_data`` may alias ``values``
        # (slice view) when ``values`` was already a cupy array of the
        # right dtype, so an in-place ``-=`` would mutate the caller.
        x_data = x_data - self.diagonal(k=k)[:x_len]
        y = self + type(self)._from_parts(
            x_data, x_indices, x_indptr, self.shape)
        self.data = y.data
        self.indices = y.indices
        self.indptr = y.indptr

    def sort_indices(self):
        """Sorts the indices of this matrix *in place*.

        .. warning::
            Calling this function might synchronize the device.

        """
        from cupyx import cusparse

        if self.ndim == 1:
            # csrsort needs a 2-D shape; sort on the (1, N) backing.
            return self._apply_2d_inplace(lambda m: m.sort_indices())

        if not self.has_sorted_indices:
            cusparse.csrsort(self)
            self.has_sorted_indices = True

    def toarray(self, order=None, out=None):
        """Returns a dense matrix representing the same value.

        Args:
            order ({'C', 'F', None}): Whether to store data in C (row-major)
                order or F (column-major) order. Default is C-order.
            out: Not supported.

        Returns:
            cupy.ndarray: Dense array representing the same matrix.

        .. seealso:: :meth:`scipy.sparse.csr_matrix.toarray`

        """
        if self.ndim == 1:
            # Densify through the (1, N) backing, then drop the length-1
            # leading axis.
            return self._as_2d().toarray(order=order).reshape(self.shape)

        from cupyx import cusparse

        order = 'C' if order is None else order.upper()
        if self.nnz == 0:
            return cupy.zeros(shape=self.shape, dtype=self.dtype, order=order)

        if self.dtype.char not in 'fdFD':
            return csr2dense(self, order)

        x = self.copy()
        x.has_canonical_format = False  # need to enforce sum_duplicates
        x.sum_duplicates()
        if (cusparse.check_availability('sparseToDense')
                and (not runtime.is_hip or (x.nnz > 0))):
            # On HIP, nnz=0 is problematic as of ROCm 4.2.0
            y = cusparse.sparseToDense(x)
            if order == 'F':
                return y
            elif order == 'C':
                return cupy.ascontiguousarray(y)
            else:
                raise ValueError('order not understood')
        else:
            # csr2dense returns F-contiguous array.
            if order == 'C':
                # To return C-contiguous array, it uses transpose.
                return cusparse.csc2dense(x.T).T
            elif order == 'F':
                return cusparse.csr2dense(x)
            else:
                raise ValueError('order not understood')

    def tobsr(self, blocksize=None, copy=False):
        # TODO(unno): Implement tobsr
        raise NotImplementedError

    def tocoo(self, copy=False):
        """Converts the matrix to COOrdinate format.

        Args:
            copy (bool): If ``False``, it shares data arrays as much as
                possible.

        Returns:
            cupyx.scipy.sparse.coo_matrix: Converted matrix.

        """
        if self.ndim == 1:
            # Build the (1, N) coo backing, then drop the length-1 axis.
            return self._as_2d().tocoo(copy=copy).reshape(self.shape)

        from cupyx import cusparse

        if copy:
            data = self.data.copy()
            indices = self.indices.copy()
        else:
            data = self.data
            indices = self.indices

        result = cusparse.csr2coo(self, data, indices)
        if not isinstance(result, self._coo_container):
            result = self._coo_container(result)
        return result

    def tocsc(self, copy=False):
        """Converts the matrix to Compressed Sparse Column format.

        Args:
            copy (bool): If ``False``, it shares data arrays as much as
                possible. Actually this option is ignored because all
                arrays in a matrix cannot be shared in csr to csc conversion.

        Returns:
            cupyx.scipy.sparse.csc_matrix: Converted matrix.

        """
        if self.ndim == 1:
            raise ValueError(
                'Cannot convert a 1-D sparse array to csc format')

        from cupyx import cusparse

        # copy is ignored
        if cusparse.check_availability('csr2csc'):
            csr2csc = cusparse.csr2csc
        elif cusparse.check_availability('csr2cscEx2'):
            csr2csc = cusparse.csr2cscEx2
        else:
            raise NotImplementedError
        # don't touch has_sorted_indices, as cuSPARSE made no guarantee
        result = csr2csc(self)
        if not isinstance(result, self._csc_container):
            result = self._csc_container(result)
        return result

    def tocsr(self, copy=False):
        """Converts the matrix to Compressed Sparse Row format.

        Args:
            copy (bool): If ``False``, the method returns itself.
                Otherwise it makes a copy of the matrix.

        Returns:
            cupyx.scipy.sparse.csr_matrix: Converted matrix.

        """
        if copy:
            return self.copy()
        else:
            return self

    def _tocsx(self):
        """Inverts the format.
        """
        return self.tocsc()

    def todia(self, copy=False):
        # TODO(unno): Implement todia
        raise NotImplementedError

    def todok(self, copy=False):
        # TODO(unno): Implement todok
        raise NotImplementedError

    def tolil(self, copy=False):
        # TODO(unno): Implement tolil
        raise NotImplementedError

    def transpose(self, axes=None, copy=False):
        """Returns the transpose of this object.

        Args:
            axes: This option is not supported.
            copy (bool): If ``True``, a returned matrix shares no data.
                Otherwise, it shared data arrays as much as possible.

        Returns:
            csc_array if ``self`` is a sparse array, else csc_matrix,
            with the dimensions reversed.

        """
        if axes is not None:
            raise ValueError(
                'Sparse matrices do not support an \'axes\' parameter because '
                'swapping dimensions is the only logical permutation.')

        if self.ndim == 1:
            # Transpose of a 1-D array is a no-op (matches scipy).
            return self.copy() if copy else self

        shape = self.shape[1], self.shape[0]
        if copy:
            data = self.data.copy()
            indices = self.indices.copy()
            indptr = self.indptr.copy()
        else:
            data, indices, indptr = self.data, self.indices, self.indptr
        return self._csc_container._from_parts(
            data, indices, indptr, shape,
            has_canonical_format=getattr(
                self, '_has_canonical_format', None),
            has_sorted_indices=getattr(
                self, '_has_sorted_indices', None))

    def _getrow(self, i):
        """Return a copy of row i as a (1 x n) CSR row vector."""
        return self._major_slice(slice(i, i + 1), copy=True)

    def _getcol(self, i):
        """Return a copy of column i as an (m x 1) CSR column vector."""
        return self._minor_slice(slice(i, i + 1), copy=True)

    def _get_intXarray(self, row, col):
        row = slice(row, row + 1)
        return self._major_slice(row)._minor_index_fancy(col)

    def _get_intXslice(self, row, col):
        row = slice(row, row + 1)
        return self._major_slice(row)._minor_slice(col, copy=True)

    def _get_sliceXint(self, row, col):
        col = slice(col, col + 1)
        copy = row.step in (1, None)
        return self._major_slice(row)._minor_slice(col, copy=copy)

    def _get_sliceXarray(self, row, col):
        return self._major_slice(row)._minor_index_fancy(col)

    def _get_arrayXint(self, row, col):
        col = slice(col, col + 1)
        return self._major_index_fancy(row)._minor_slice(col)

    def _get_arrayXslice(self, row, col):
        if col.step not in (1, None):
            start, stop, step = col.indices(self.shape[1])
            cols = cupy.arange(start, stop, step, self.indices.dtype)
            return self._get_arrayXarray(row, cols)
        return self._major_index_fancy(row)._minor_slice(col)


class csr_matrix(_base.spmatrix, _csr_base):
    """Compressed Sparse Row matrix.

    This can be instantiated in several ways.

    ``csr_matrix(D)``
        ``D`` is a rank-2 :class:`cupy.ndarray`.
    ``csr_matrix(S)``
        ``S`` is another sparse matrix. It is equivalent to ``S.tocsr()``.
    ``csr_matrix((M, N), [dtype])``
        It constructs an empty matrix whose shape is ``(M, N)``. Default dtype
        is float64.
    ``csr_matrix((data, (row, col)))``
        All ``data``, ``row`` and ``col`` are one-dimensional
        :class:`cupy.ndarray`.
    ``csr_matrix((data, indices, indptr))``
        All ``data``, ``indices`` and ``indptr`` are one-dimensional
        :class:`cupy.ndarray`.

    Args:
        arg1: Arguments for the initializer.
        shape (tuple): Shape of a matrix. Its length must be two.
        dtype: Data type. It must be an argument of :class:`numpy.dtype`.
        copy (bool): If ``True``, copies of given arrays are always used.

    .. seealso::
        - :class:`scipy.sparse.csr_matrix`
        - :class:`~cupyx.scipy.sparse.csr_array`

    """
    pass


class csr_array(_csr_base, _base.sparray):
    """Compressed Sparse Row array.

    Same constructor interface as :class:`csr_matrix`.

    Unlike ``csr_matrix``, the ``*`` operator performs element-wise
    multiplication (not matrix multiplication).  Use ``@`` for matmul.

    Supports 1-D shapes (``ndim == 1``) in addition to 2-D, matching
    :class:`scipy.sparse.csr_array`.  A 1-D array is stored as a single
    CSR row (``indptr == [0, nnz]``) presenting ``shape == (N,)``, and
    supports element-wise arithmetic, comparisons, matmul, reductions,
    indexing, and assignment.  ``diagonal``/``setdiag`` and conversion
    to CSC require two dimensions, as in scipy.

    .. seealso::
       - :class:`scipy.sparse.csr_array`
       - :class:`~cupyx.scipy.sparse.csr_matrix`
    """

    _allow_nd = (1, 2)


def _is_csr(x):
    """True if x is any CSR sparse (array or matrix)."""
    return isinstance(x, _csr_base)


def _is_csc(x):
    """True if x is any CSC sparse (array or matrix)."""
    # Uses format string instead of isinstance(_csc_base) to avoid a
    # circular import between _csr and _csc.
    return _base.issparse(x) and x.format == 'csc'


def isspmatrix_csr(x):
    """Checks if a given matrix is of CSR format.

    Returns:
        bool: Returns if ``x`` is :class:`cupyx.scipy.sparse.csr_matrix`.

    """
    return isinstance(x, csr_matrix)


def check_shape_for_pointwise_op(a_shape, b_shape, allow_broadcasting=True):
    if allow_broadcasting:
        a_m, a_n = a_shape
        b_m, b_n = b_shape
        if not (a_m == b_m or a_m == 1 or b_m == 1):
            raise ValueError('inconsistent shape')
        if not (a_n == b_n or a_n == 1 or b_n == 1):
            raise ValueError('inconsistent shape')
    else:
        if a_shape != b_shape:
            raise ValueError('inconsistent shape')


def multiply_by_scalar(sp, a):
    data = sp.data * a
    return type(sp)._from_parts(
        data, sp.indices.copy(), sp.indptr.copy(), sp.shape,
        has_canonical_format=getattr(
            sp, '_has_canonical_format', None),
        has_sorted_indices=getattr(
            sp, '_has_sorted_indices', None))


def multiply_by_dense(sp, dn):
    check_shape_for_pointwise_op(sp.shape, dn.shape)
    sp_m, sp_n = sp.shape
    dn_m, dn_n = dn.shape
    m, n = max(sp_m, dn_m), max(sp_n, dn_n)
    nnz = sp.nnz * (m // sp_m) * (n // sp_n)
    dtype = numpy.promote_types(sp.dtype, dn.dtype)
    data = cupy.empty(nnz, dtype=dtype)
    indices = cupy.empty(nnz, dtype=sp.indices.dtype)
    if m > sp_m:
        if n > sp_n:
            indptr = cupy.arange(0, nnz+1, n, dtype=sp.indptr.dtype)
        else:
            indptr = cupy.arange(0, nnz+1, sp.nnz, dtype=sp.indptr.dtype)
    else:
        indptr = sp.indptr.copy()
        if n > sp_n:
            indptr *= n

    # out = sp * dn
    idx_dtype = sp.indptr.dtype
    it = idx_dtype.type
    cupy_multiply_by_dense()(
        sp.data, sp.indptr, sp.indices, it(sp_m), it(sp_n),
        dn, it(dn_m), it(dn_n), indptr, it(m), it(n),
        data, indices)

    return type(sp)._from_parts(
        data, indices, indptr, shape=(m, n))


_GET_ROW_ID_ = '''
template<typename I>
__device__ inline I get_row_id(I i, I min, I max, const I *indptr) {
    I row = (min + max) / 2;
    while (min < max) {
        if (i < indptr[row]) {
            max = row - 1;
        } else if (i >= indptr[row + 1]) {
            min = row + 1;
        } else {
            break;
        }
        row = (min + max) / 2;
    }
    return row;
}
'''

_FIND_INDEX_HOLDING_COL_IN_ROW_ = '''
template<typename I>
__device__ inline I find_index_holding_col_in_row(
        I row, I col, const I *indptr, const I *indices) {
    I j_min = indptr[row];
    I j_max = indptr[row+1] - 1;
    while (j_min <= j_max) {
        I j = (j_min + j_max) / 2;
        I j_col = indices[j];
        if (j_col == col) {
            return j;
        } else if (j_col < col) {
            j_min = j + 1;
        } else {
            j_max = j - 1;
        }
    }
    return (I)(-1);
}
'''


@cupy._util.memoize(for_each_device=True)
def cupy_multiply_by_dense():
    return cupy.ElementwiseKernel(
        '''
        raw S SP_DATA, raw I SP_INDPTR, raw I SP_INDICES,
        I SP_M, I SP_N,
        raw D DN_DATA, I DN_M, I DN_N,
        raw I OUT_INDPTR, I OUT_M, I OUT_N
        ''',
        'O OUT_DATA, I OUT_INDICES',
        '''
        I i_out = (I)i;
        I m_out = get_row_id(i_out, (I)0, OUT_M - 1, &(OUT_INDPTR[0]));
        I i_sp = i_out;
        if (OUT_M > SP_M && SP_M == 1) {
            i_sp -= OUT_INDPTR[m_out];
        }
        if (OUT_N > SP_N && SP_N == 1) {
            i_sp /= OUT_N;
        }
        I n_out = SP_INDICES[i_sp];
        if (OUT_N > SP_N && SP_N == 1) {
            n_out = i_out - OUT_INDPTR[m_out];
        }
        I m_dn = m_out;
        if (OUT_M > DN_M && DN_M == 1) {
            m_dn = (I)0;
        }
        I n_dn = n_out;
        if (OUT_N > DN_N && DN_N == 1) {
            n_dn = (I)0;
        }
        OUT_DATA = (O)(SP_DATA[i_sp] * DN_DATA[n_dn + (DN_N * m_dn)]);
        OUT_INDICES = n_out;
        ''',
        'cupyx_scipy_sparse_csr_multiply_by_dense',
        preamble=_GET_ROW_ID_
    )


@cupy._util.memoize(for_each_device=True)
def _cupy_divide_by_dense():
    return cupy.ElementwiseKernel(
        'T data, I row, I col, I width, raw T other',
        'T res',
        '''
        res = data / other[row * width + col]
        ''',
        'cupyx_scipy_sparse_coo_divide_dense',
    )


def multiply_by_csr(a, b, *, out_cls=None):
    # ``out_cls`` pins the result type to the caller (e.g. ``type(self)``
    # from ``_csr_base.multiply``) so the swap below for performance
    # doesn't leak the other operand's type into the result.
    if out_cls is None:
        out_cls = type(a)
    check_shape_for_pointwise_op(a.shape, b.shape)
    a_m, a_n = a.shape
    b_m, b_n = b.shape
    m, n = max(a_m, b_m), max(a_n, b_n)
    a_nnz = a.nnz * (m // a_m) * (n // a_n)
    b_nnz = b.nnz * (m // b_m) * (n // b_n)
    if a_nnz > b_nnz:
        return multiply_by_csr(b, a, out_cls=out_cls)
    # Harmonise index dtypes so the kernel's templated ``I`` parameter
    # is consistent across both operands (mirrors ``binopt_csr``).
    # Without this, mixed int32/int64 inputs trip ``TypeError: Type is
    # mismatched. B_INDPTR int32 int64 I``.
    idx_dtype = numpy.promote_types(a.indices.dtype, b.indices.dtype)
    a_indices = a.indices.astype(idx_dtype, copy=False)
    a_indptr = a.indptr.astype(idx_dtype, copy=False)
    b_indices = b.indices.astype(idx_dtype, copy=False)
    b_indptr = b.indptr.astype(idx_dtype, copy=False)
    c_nnz = a_nnz
    dtype = numpy.promote_types(a.dtype, b.dtype)
    c_data = cupy.empty(c_nnz, dtype=dtype)
    c_indices = cupy.empty(c_nnz, dtype=idx_dtype)
    if m > a_m:
        if n > a_n:
            c_indptr = cupy.arange(0, c_nnz+1, n, dtype=idx_dtype)
        else:
            c_indptr = cupy.arange(0, c_nnz+1, a.nnz, dtype=idx_dtype)
    else:
        c_indptr = a_indptr.copy()
        if n > a_n:
            c_indptr *= n
    flags = cupy.zeros(c_nnz+1, dtype=idx_dtype)
    nnz_each_row = cupy.zeros(m+1, dtype=idx_dtype)

    # compute c = a * b where necessary and get sparsity pattern of matrix d
    it = idx_dtype.type
    cupy_multiply_by_csr_step1()(
        a.data, a_indptr, a_indices, it(a_m), it(a_n),
        b.data, b_indptr, b_indices, it(b_m), it(b_n),
        c_indptr, it(m), it(n), c_data, c_indices, flags, nnz_each_row)

    flags = cupy.cumsum(flags, dtype=idx_dtype)
    d_indptr = cupy.cumsum(nnz_each_row, dtype=idx_dtype)
    d_nnz = int(d_indptr[-1])  # synchronize!
    d_data = cupy.empty(d_nnz, dtype=dtype)
    d_indices = cupy.empty(d_nnz, dtype=idx_dtype)

    # remove zero elements in matrix c
    cupy_multiply_by_csr_step2()(c_data, c_indices, flags, d_data, d_indices)

    return out_cls._from_parts(
        d_data, d_indices, d_indptr, shape=(m, n))


@cupy._util.memoize(for_each_device=True)
def cupy_multiply_by_csr_step1():
    return cupy.ElementwiseKernel(
        '''
        raw A A_DATA, raw I A_INDPTR, raw I A_INDICES, I A_M, I A_N,
        raw B B_DATA, raw I B_INDPTR, raw I B_INDICES, I B_M, I B_N,
        raw I C_INDPTR, I C_M, I C_N
        ''',
        'C C_DATA, I C_INDICES, raw I FLAGS, raw I NNZ_EACH_ROW',
        '''
        I i_c = (I)i;
        I m_c = get_row_id(i_c, (I)0, C_M - 1, &(C_INDPTR[0]));

        I i_a = (I)i;
        if (C_M > A_M && A_M == 1) {
            i_a -= C_INDPTR[m_c];
        }
        if (C_N > A_N && A_N == 1) {
            i_a /= C_N;
        }
        I n_c = A_INDICES[i_a];
        if (C_N > A_N && A_N == 1) {
            n_c = (I)i % C_N;
        }
        I m_b = m_c;
        if (C_M > B_M && B_M == 1) {
            m_b = (I)0;
        }
        I n_b = n_c;
        if (C_N > B_N && B_N == 1) {
            n_b = (I)0;
        }
        I i_b = find_index_holding_col_in_row(m_b, n_b,
            &(B_INDPTR[0]), &(B_INDICES[0]));
        if (i_b >= (I)0) {
            atomicAdd(&NNZ_EACH_ROW[m_c+1], (I)1);
            FLAGS[i+1] = 1;
            C_DATA = (C)(A_DATA[i_a] * B_DATA[i_b]);
            C_INDICES = n_c;
        }
        ''',
        'cupyx_scipy_sparse_csr_multiply_by_csr_step1',
        preamble=_GET_ROW_ID_ + _FIND_INDEX_HOLDING_COL_IN_ROW_
    )


@cupy._util.memoize(for_each_device=True)
def cupy_multiply_by_csr_step2():
    return cupy.ElementwiseKernel(
        'T C_DATA, I C_INDICES, raw I FLAGS',
        'raw D D_DATA, raw I D_INDICES',
        '''
        I j = FLAGS[i];
        if (j < FLAGS[i+1]) {
            D_DATA[j] = (D)(C_DATA);
            D_INDICES[j] = C_INDICES;
        }
        ''',
        'cupyx_scipy_sparse_csr_multiply_by_csr_step2'
    )


_BINOPT_MAX_ = '''
__device__ inline O binopt(T in1, T in2) {
    return max(in1, in2);
}
'''
_BINOPT_MIN_ = '''
__device__ inline O binopt(T in1, T in2) {
    return min(in1, in2);
}
'''
_BINOPT_EQ_ = '''
__device__ inline O binopt(T in1, T in2) {
    return (in1 == in2);
}
'''
_BINOPT_NE_ = '''
__device__ inline O binopt(T in1, T in2) {
    return (in1 != in2);
}
'''
_BINOPT_LT_ = '''
__device__ inline O binopt(T in1, T in2) {
    return (in1 < in2);
}
'''
_BINOPT_GT_ = '''
__device__ inline O binopt(T in1, T in2) {
    return (in1 > in2);
}
'''
_BINOPT_LE_ = '''
__device__ inline O binopt(T in1, T in2) {
    return (in1 <= in2);
}
'''
_BINOPT_GE_ = '''
__device__ inline O binopt(T in1, T in2) {
    return (in1 >= in2);
}
'''


def binopt_csr(a, b, op_name):
    check_shape_for_pointwise_op(a.shape, b.shape)
    a_m, a_n = a.shape
    b_m, b_n = b.shape
    m, n = max(a_m, b_m), max(a_n, b_n)
    a_nnz = a.nnz * (m // a_m) * (n // a_n)
    b_nnz = b.nnz * (m // b_m) * (n // b_n)

    idx_dtype = numpy.promote_types(a.indices.dtype, b.indices.dtype)
    a_indptr = a.indptr.astype(idx_dtype, copy=False)
    a_indices = a.indices.astype(idx_dtype, copy=False)
    b_indptr = b.indptr.astype(idx_dtype, copy=False)
    b_indices = b.indices.astype(idx_dtype, copy=False)
    a_info = cupy.zeros(a_nnz + 1, dtype=idx_dtype)
    b_info = cupy.zeros(b_nnz + 1, dtype=idx_dtype)
    a_valid = cupy.zeros(a_nnz, dtype=numpy.int8)
    b_valid = cupy.zeros(b_nnz, dtype=numpy.int8)
    c_indptr = cupy.zeros(m + 1, dtype=idx_dtype)
    in_dtype = numpy.promote_types(a.dtype, b.dtype)
    a_data = a.data.astype(in_dtype, copy=False)
    b_data = b.data.astype(in_dtype, copy=False)
    funcs = _GET_ROW_ID_
    if op_name == '_maximum_':
        funcs += _BINOPT_MAX_
        out_dtype = in_dtype
    elif op_name == '_minimum_':
        funcs += _BINOPT_MIN_
        out_dtype = in_dtype
    elif op_name == '_eq_':
        funcs += _BINOPT_EQ_
        out_dtype = numpy.bool_
    elif op_name == '_ne_':
        funcs += _BINOPT_NE_
        out_dtype = numpy.bool_
    elif op_name == '_lt_':
        funcs += _BINOPT_LT_
        out_dtype = numpy.bool_
    elif op_name == '_gt_':
        funcs += _BINOPT_GT_
        out_dtype = numpy.bool_
    elif op_name == '_le_':
        funcs += _BINOPT_LE_
        out_dtype = numpy.bool_
    elif op_name == '_ge_':
        funcs += _BINOPT_GE_
        out_dtype = numpy.bool_
    else:
        raise ValueError('invalid op_name: {}'.format(op_name))
    a_tmp_data = cupy.empty(a_nnz, dtype=out_dtype)
    b_tmp_data = cupy.empty(b_nnz, dtype=out_dtype)
    a_tmp_indices = cupy.empty(a_nnz, dtype=idx_dtype)
    b_tmp_indices = cupy.empty(b_nnz, dtype=idx_dtype)
    _size = a_nnz + b_nnz
    it = idx_dtype.type
    cupy_binopt_csr_step1(op_name, preamble=funcs)(
        it(m), it(n),
        a_indptr, a_indices, a_data, it(a_m), it(a_n), it(a.nnz), it(a_nnz),
        b_indptr, b_indices, b_data, it(b_m), it(b_n), it(b.nnz), it(b_nnz),
        a_info, a_valid, a_tmp_indices, a_tmp_data,
        b_info, b_valid, b_tmp_indices, b_tmp_data,
        c_indptr, size=_size)
    a_info = cupy.cumsum(a_info, dtype=a_info.dtype)
    b_info = cupy.cumsum(b_info, dtype=b_info.dtype)
    c_indptr = cupy.cumsum(c_indptr, dtype=c_indptr.dtype)
    c_nnz = int(c_indptr[-1])  # synchronize!
    c_indices = cupy.empty(c_nnz, dtype=idx_dtype)
    c_data = cupy.empty(c_nnz, dtype=out_dtype)
    cupy_binopt_csr_step2(op_name)(
        a_info, a_valid, a_tmp_indices, a_tmp_data, it(a_nnz),
        b_info, b_valid, b_tmp_indices, b_tmp_data, it(b_nnz),
        c_indices, c_data, size=_size)
    return type(a)._from_parts(
        c_data, c_indices, c_indptr, shape=(m, n))


@cupy._util.memoize(for_each_device=True)
def cupy_binopt_csr_step1(op_name, preamble=''):
    name = 'cupyx_scipy_sparse_csr_binopt_' + op_name + 'step1'
    return cupy.ElementwiseKernel(
        '''
        I M, I N,
        raw I A_INDPTR, raw I A_INDICES, raw T A_DATA,
        I A_M, I A_N, I A_NNZ_ACT, I A_NNZ,
        raw I B_INDPTR, raw I B_INDICES, raw T B_DATA,
        I B_M, I B_N, I B_NNZ_ACT, I B_NNZ
        ''',
        '''
        raw I A_INFO, raw B A_VALID, raw I A_TMP_INDICES, raw O A_TMP_DATA,
        raw I B_INFO, raw B B_VALID, raw I B_TMP_INDICES, raw O B_TMP_DATA,
        raw I C_INFO
        ''',
        '''
        if (i >= A_NNZ + B_NNZ) return;

        const I *MY_INDPTR, *MY_INDICES;  I *MY_INFO;  const T *MY_DATA;
        const I *OP_INDPTR, *OP_INDICES;  I *OP_INFO;  const T *OP_DATA;
        I MY_M, MY_N, MY_NNZ_ACT, MY_NNZ;
        I OP_M, OP_N, OP_NNZ_ACT, OP_NNZ;
        signed char *MY_VALID;  I *MY_TMP_INDICES;  O *MY_TMP_DATA;

        I my_j;
        if (i < A_NNZ) {
            // in charge of one of non-zero element of sparse matrix A
            my_j = (I)i;
            MY_INDPTR  = &(A_INDPTR[0]);   OP_INDPTR  = &(B_INDPTR[0]);
            MY_INDICES = &(A_INDICES[0]);  OP_INDICES = &(B_INDICES[0]);
            MY_INFO    = &(A_INFO[0]);     OP_INFO    = &(B_INFO[0]);
            MY_DATA    = &(A_DATA[0]);     OP_DATA    = &(B_DATA[0]);
            MY_M       = A_M;              OP_M       = B_M;
            MY_N       = A_N;              OP_N       = B_N;
            MY_NNZ_ACT = A_NNZ_ACT;        OP_NNZ_ACT = B_NNZ_ACT;
            MY_NNZ     = A_NNZ;            OP_NNZ     = B_NNZ;
            MY_VALID   = &(A_VALID[0]);
            MY_TMP_DATA= &(A_TMP_DATA[0]);
            MY_TMP_INDICES = &(A_TMP_INDICES[0]);
        } else {
            // in charge of one of non-zero element of sparse matrix B
            my_j = (I)i - A_NNZ;
            MY_INDPTR  = &(B_INDPTR[0]);   OP_INDPTR  = &(A_INDPTR[0]);
            MY_INDICES = &(B_INDICES[0]);  OP_INDICES = &(A_INDICES[0]);
            MY_INFO    = &(B_INFO[0]);     OP_INFO    = &(A_INFO[0]);
            MY_DATA    = &(B_DATA[0]);     OP_DATA    = &(A_DATA[0]);
            MY_M       = B_M;              OP_M       = A_M;
            MY_N       = B_N;              OP_N       = A_N;
            MY_NNZ_ACT = B_NNZ_ACT;        OP_NNZ_ACT = A_NNZ_ACT;
            MY_NNZ     = B_NNZ;            OP_NNZ     = A_NNZ;
            MY_VALID   = &(B_VALID[0]);
            MY_TMP_DATA= &(B_TMP_DATA[0]);
            MY_TMP_INDICES = &(B_TMP_INDICES[0]);
        }
        I _min, _max, _mid;

        // get column location
        I my_col;
        I my_j_act = my_j;
        if (MY_M == 1 && MY_M < M) {
            if (MY_N == 1 && MY_N < N) my_j_act = 0;
            else                       my_j_act = my_j % MY_NNZ_ACT;
        } else {
            if (MY_N == 1 && MY_N < N) my_j_act = my_j / N;
        }
        my_col = MY_INDICES[my_j_act];
        if (MY_N == 1 && MY_N < N) {
            my_col = my_j % N;
        }

        // get row location
        I my_row = get_row_id(my_j_act, (I)0, MY_M - 1, &(MY_INDPTR[0]));
        if (MY_M == 1 && MY_M < M) {
            if (MY_N == 1 && MY_N < N) my_row = my_j / N;
            else                       my_row = my_j / MY_NNZ_ACT;
        }

        I op_row = my_row;
        I op_row_act = op_row;
        if (OP_M == 1 && OP_M < M) {
            op_row_act = 0;
        }

        I op_col = 0;
        _min = OP_INDPTR[op_row_act];
        _max = OP_INDPTR[op_row_act + 1] - 1;
        I op_j_act = _min;
        bool op_nz = false;
        if (_min <= _max) {
            if (OP_N == 1 && OP_N < N) {
                op_col = my_col;
                op_nz = true;
            }
            else {
                _mid = (_min + _max) / 2;
                op_col = OP_INDICES[_mid];
                while (_min < _max) {
                    if (op_col < my_col) {
                        _min = _mid + 1;
                    } else if (op_col > my_col) {
                        _max = _mid;
                    } else {
                        break;
                    }
                    _mid = (_min + _max) / 2;
                    op_col = OP_INDICES[_mid];
                }
                op_j_act = _mid;
                if (op_col == my_col) {
                    op_nz = true;
                } else if (op_col < my_col) {
                    op_col = N;
                    op_j_act += 1;
                }
            }
        }

        I op_j = op_j_act;
        if (OP_M == 1 && OP_M < M) {
            if (OP_N == 1 && OP_N < N) {
                op_j = (op_col + N * op_row) * OP_NNZ_ACT;
            } else {
                op_j = op_j_act + OP_NNZ_ACT * op_row;
            }
        } else {
            if (OP_N == 1 && OP_N < N) {
                op_j = op_col + N * op_j_act;
            }
        }

        if (i < A_NNZ || !op_nz) {
            T my_data = MY_DATA[my_j_act];
            T op_data = 0;
            if (op_nz) op_data = OP_DATA[op_j_act];
            O out;
            if (i < A_NNZ) out = binopt(my_data, op_data);
            else           out = binopt(op_data, my_data);
            if (out != static_cast<O>(0)) {
                MY_VALID[my_j] = 1;
                MY_TMP_DATA[my_j] = out;
                MY_TMP_INDICES[my_j] = my_col;
                atomicAdd(&C_INFO[my_row + 1], (I)1);
                atomicAdd(&MY_INFO[my_j + 1], (I)1);
                atomicAdd(&OP_INFO[op_j], (I)1);
            }
        }
        ''',
        name, preamble=preamble,
    )


@cupy._util.memoize(for_each_device=True)
def cupy_binopt_csr_step2(op_name):
    name = 'cupyx_scipy_sparse_csr_binopt' + op_name + 'step2'
    return cupy.ElementwiseKernel(
        '''
        raw I A_INFO, raw B A_VALID, raw I A_TMP_INDICES, raw O A_TMP_DATA,
        I A_NNZ,
        raw I B_INFO, raw B B_VALID, raw I B_TMP_INDICES, raw O B_TMP_DATA,
        I B_NNZ
        ''',
        'raw I C_INDICES, raw O C_DATA',
        '''
        if (i < A_NNZ) {
            I j = (I)i;
            if (A_VALID[j]) {
                C_INDICES[A_INFO[j]] = A_TMP_INDICES[j];
                C_DATA[A_INFO[j]]    = A_TMP_DATA[j];
            }
        } else if (i < A_NNZ + B_NNZ) {
            I j = (I)i - A_NNZ;
            if (B_VALID[j]) {
                C_INDICES[B_INFO[j]] = B_TMP_INDICES[j];
                C_DATA[B_INFO[j]]    = B_TMP_DATA[j];
            }
        }
        ''',
        name,
    )


def csr2dense(a, order):
    out = cupy.zeros(a.shape, dtype=a.dtype, order=order)
    m, n = a.shape
    idx_dtype = a.indptr.dtype
    kern = _cupy_csr2dense(a.dtype)
    kern(idx_dtype.type(m), idx_dtype.type(n),
         a.indptr, a.indices, a.data, out)
    return out


@cupy._util.memoize(for_each_device=True)
def _cupy_csr2dense(dtype):
    if dtype == '?':
        op = "if (DATA) OUT[index] = true;"
    else:
        op = "atomicAdd(&OUT[index], DATA);"

    return cupy.ElementwiseKernel(
        'I M, I N, raw I INDPTR, I INDICES, T DATA',
        'raw T OUT',
        '''
        I row = get_row_id((I)i, (I)0, M - 1, &(INDPTR[0]));
        I col = INDICES;
        // Logical (C-order) position: the raw-array indexer applies the
        // output's actual strides, so this addresses both C- and
        // F-order allocations correctly.  Compute in 64-bit: with int32
        // indices the flat position ``col + N * row`` can exceed INT32_MAX
        // (M*N > 2^31) even though M, N, and each index fit int32, which
        // would otherwise wrap negative and write out of bounds.
        long long index = (long long)col + (long long)N * (long long)row;
        ''' + op,
        'cupyx_scipy_sparse_csr2dense',
        preamble=_GET_ROW_ID_
    )


def dense2csr(a):
    from cupyx import cusparse

    if a.dtype.char in 'fdFD':
        if cusparse.check_availability('denseToSparse'):
            return cusparse.denseToSparse(a, format='csr')
        else:
            return cusparse.dense2csr(a)
    m, n = a.shape
    mn = m * n
    # The prefix-sum scratch values run up to mn, hence maxval=mn.
    idx_dtype = _sputils.get_index_dtype(maxval=mn)
    a = cupy.ascontiguousarray(a)
    indptr = cupy.zeros(m + 1, dtype=idx_dtype)
    info = cupy.zeros(mn + 1, dtype=idx_dtype)
    it = numpy.dtype(idx_dtype).type
    cupy_dense2csr_step1()(it(m), it(n), a, indptr, info)
    cupy.cumsum(indptr, out=indptr)
    cupy.cumsum(info, out=info)
    nnz = int(indptr[-1])  # synchronize!
    indices = cupy.empty(nnz, dtype=idx_dtype)
    data = cupy.empty(nnz, dtype=a.dtype)
    cupy_dense2csr_step2()(it(m), it(n), a, info, indices, data)
    return csr_matrix._from_parts(
        data, indices, indptr, (m, n))


@cupy._util.memoize(for_each_device=True)
def cupy_dense2csr_step1():
    return cupy.ElementwiseKernel(
        'I M, I N, T A',
        'raw I INDPTR, raw I INFO',
        '''
        I row = (I)i / N;
        I col = (I)i % N;
        if (A != static_cast<T>(0)) {
            atomicAdd(&INDPTR[row + 1], (I)1);
            INFO[(I)i + 1] = 1;
        }
        ''',
        'cupyx_scipy_sparse_dense2csr_step1')


@cupy._util.memoize(for_each_device=True)
def cupy_dense2csr_step2():
    return cupy.ElementwiseKernel(
        'I M, I N, T A, raw I INFO',
        'raw I INDICES, raw T DATA',
        '''
        I row = (I)i / N;
        I col = (I)i % N;
        if (A != static_cast<T>(0)) {
            I idx = INFO[(I)i];
            INDICES[idx] = col;
            DATA[idx] = A;
        }
        ''',
        'cupyx_scipy_sparse_dense2csr_step2')


@cupy._util.memoize(for_each_device=True)
def _cupy_csr_diagonal():
    return cupy.ElementwiseKernel(
        'I k, I rows, I cols, '
        'raw T data, raw I indptr, raw I indices',
        'T y',
        '''
        I row = (I)i;
        I col = (I)i;
        if (k < 0) row -= (I)k;
        if (k > 0) col += (I)k;
        if (row >= rows || col >= cols) return;
        I j = find_index_holding_col_in_row(row, col,
            &(indptr[0]), &(indices[0]));
        if (j >= (I)0) {
            y = data[j];
        } else {
            y = static_cast<T>(0);
        }
        ''',
        'cupyx_scipy_sparse_csr_diagonal',
        preamble=_FIND_INDEX_HOLDING_COL_IN_ROW_
    )
