"""Extra functionality for [sparse](https://sparse.pydata.org) and [scipy.sparse](https://docs.scipy.org/doc/scipy/reference/sparse.html)"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/api/utilities/sparse.ipynb.

# %% auto 0
__all__ = ['sparse2d_rand', 'spilu', 'sparse2d', 'sparse', 'sparse2d_identity', 'sparse_identity', 'sparse2d_diag', 'sparse_diag',
           'dense_toeplitz', 'sparse_toeplitz', 'sparse2d_kronecker_matrix', 'sparse_kronecker_matrix', 'kron',
           'restrict_bandwidth']

# %% ../../nbs/api/utilities/sparse.ipynb
import itertools as it

from fastcore.meta import delegates

import numpy as np
from numpy import array, ndarray
import scipy.sparse as sps
from scipy.sparse import sparray
import scipy.linalg as spl
import scipy.sparse.linalg as spa
from sparse import COO, GCXS, SparseArray
import sparse as ss

# %% ../../nbs/api/utilities/sparse.ipynb
def sparse2d(a):
    """Create a 2D sparse array in specified format."""
    if isinstance(a, SparseArray): a = a.to_scipy_sparse()
    return sps.csr_array(a)

# %% ../../nbs/api/utilities/sparse.ipynb
@delegates(COO, but='data')
def sparse(a, format='coo', shape=None, **kwargs)->SparseArray:
    """Create an N-D sparse array in specified format."""
    if format == 'gcxs': cls = GCXS
    elif format == 'coo': cls = COO
    if isinstance(a, ndarray): return cls.from_numpy(a)
    if sps.issparse(a): return cls.from_scipy_sparse(a)
    if isinstance(a, SparseArray): return cls(a)
    # Assume a is coords and data in COO format
    if shape is None: shape = [max(c) - min(c) + 1 for c in a[0]]
    return cls(COO(*a, shape=shape, **kwargs))

# %% ../../nbs/api/utilities/sparse.ipynb
def sparse2d_identity(n):
    return sps.eye_array(n)

# %% ../../nbs/api/utilities/sparse.ipynb
def sparse_identity(n, format='coo')->SparseArray:
    return ss.eye(n, format=format)

# %% ../../nbs/api/utilities/sparse.ipynb
def sparse2d_diag(diags, offsets=0):
    return sps.diags_array(diags, offsets=offsets)

# %% ../../nbs/api/utilities/sparse.ipynb
def sparse_diag(diag, format='coo')->SparseArray:
    return sparse(sparse2d_diag(diag), format=format)

# %% ../../nbs/api/utilities/sparse.ipynb
def dense_toeplitz(
    diags:ndarray, # Value on each diagonal, starting at lower left.
):
    """Return a toeplitz array given the value on each diagonal (starting at lower left)."""
    n = (len(diags) - 1) // 2
    return spl.toeplitz(diags[:n+1][::-1], diags[n:])

# %% ../../nbs/api/utilities/sparse.ipynb
def sparse_toeplitz(
    diags:ndarray, # Value on each diagonal, starting at lower left.
    format='coo'
)->SparseArray:
    """Return a toeplitz sparse array in specified format given the value on each diagonal (starting at lower left)."""
    return sparse(dense_toeplitz(diags), format=format)

# %% ../../nbs/api/utilities/sparse.ipynb
def sparse2d_kronecker_matrix(n, k=0):
    """The nxn sparse kronecker matrix delta_(i,j-k) in specified format."""
    return sparse2d_diag(np.ones(n - abs(k)), offsets=k)

# %% ../../nbs/api/utilities/sparse.ipynb
def sparse_kronecker_matrix(n, k=0, format='coo')->SparseArray:
    """The nxn sparse kronecker matrix delta_(i,j-k) in specified format."""
    return sparse(sparse2d_kronecker_matrix(n, k), format=format)

# %% ../../nbs/api/utilities/sparse.ipynb
def kron(a, b, *rest):
    if rest: return kron(kron(a, b), *rest)
    if sps.issparse(a) or sps.issparse(b): 
        return sps.kron(a, b, format='csr')
    return np.kron(a, b)

# %% ../../nbs/api/utilities/sparse.ipynb
sparse2d_rand = sps.random_array

# %% ../../nbs/api/utilities/sparse.ipynb
spilu = spa.spilu

# %% ../../nbs/api/utilities/sparse.ipynb
def restrict_bandwidth(a_sparray:sparray, width):
    """Make a scipy sparse array banded by setting all elements outside the bandwidth to zero."""
    result = a_sparray.copy()
    i, j = result.nonzero()
    result.data[np.abs(i - j) > width] = 0.
    result.eliminate_zeros()
    return result
