import numpy as np
from lacuna.sparse.csr import CSR

# A = [[1, 0, 2],
#      [0, 3, 0]]
indptr = np.array([0, 2, 3], dtype=np.int64)
indices = np.array([0, 2, 1], dtype=np.int64)
data = np.array([1.0, 2.0, 3.0], dtype=np.float64)
A = CSR(indptr, indices, data, (2, 3), check=False)

# SpMV
x = np.array([10.0, 20.0, 30.0], dtype=np.float64)
y = A @ x  # -> array([1*10 + 2*30, 3*20])

# SpMM
B = np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], dtype=np.float64)
Y = A @ B  # shape (2, 2)

# Reductions
total = A.sum()            # 6.0
row_sums = A.sum(axis=1)   # [3.0, 3.0]
col_sums = A.sum(axis=0)   # [1.0, 3.0, 2.0]

# Transpose
AT = A.T  # (3, 2)

# Arithmetic
C = A + A
Z = A - A
H = A.multiply(A)  # Hadamard (elementwise)

# Cleanup
Az = A.eliminate_zeros()
Ap = A.prune(1e-6)

# Indexing
a00 = A[0, 0]      # 1.0
row0 = A[0, :]     # dense row (numpy)
col1 = A[:, 1]     # dense col (numpy)

import numpy as np
from lacuna.sparse import CSC, COO

# CSC
indptr = np.array([0, 1, 2, 3], dtype=np.int64)
indices = np.array([0, 1, 1], dtype=np.int64)
data = np.array([1.0, 2.0, 3.0], dtype=np.float64)
C = CSC(indptr, indices, data, (2, 3), check=False)
_ = C.sum()            # total sum
_ = C.sum(axis=0)      # column sums
_ = C.sum(axis=1)      # row sums
CT = C.T               # transpose (still CSC)

# COO
row = np.array([0, 1, 1], dtype=np.int64)
col = np.array([0, 0, 2], dtype=np.int64)
val = np.array([1.0, 2.0, 3.0], dtype=np.float64)
O = COO(row, col, val, (2, 3), check=False)
_ = O.sum()            # total sum
OT = O.T               # transpose (COO)

import numpy as np
from lacuna.sparse import COOND

shape = (2, 3, 4)
# indices can be flattened (nnz * ndim)
indices = np.array([
    0, 1, 2,
    1, 2, 3,
], dtype=np.int64)
data = np.array([1.0, 3.0], dtype=np.float64)
A = COOND(shape, indices, data, check=False)

total = A.sum()                # 4.0
avg = A.mean()                 # mean over full dense shape
B = A.permute_axes([2, 1, 0])  # shape -> (4, 3, 2)
R = A.reduce_sum_axes([2])     # sum over last axis -> shape (2, 3)
M = A.reduce_mean_axes([0, 2]) # mean over axes 0 and 2 -> shape (3,)
S = A.reshape((3, 2, 4))
H = A.hadamard_broadcast(S)    # broadcasting elementwise product

# Unfold to 2D sparse matrices
CSR0 = A.mode_unfold_to_csr(axis=0)  # rows = shape[0], cols = prod(shape[1:])
CSC1 = A.mode_unfold_to_csc(axis=1)
CSRrg = A.axes_unfold_to_csr([0, 2]) # group axes as rows
CSCrg = A.axes_unfold_to_csc([1])    # group axes as rows (CSC)