import torch
from scipy.sparse import csc_matrix
from scipy.sparse.linalg import splu, spsolve_triangular

from .sparse_convert import torch_csc_to_scipy_csc, to_numpy, from_numpy


def _extract_block_diagonal(A_scipy, block_size: int) -> csc_matrix:
    A_coo = A_scipy.tocoo()
    block_mask = (A_coo.row // block_size) == (A_coo.col // block_size)
    return csc_matrix(
        (A_coo.data[block_mask], (A_coo.row[block_mask], A_coo.col[block_mask])),
        shape=A_scipy.shape,
    )


class BlockJacobi:

    def __init__(self, A: torch.Tensor, block_size: int = 32):
        A_scipy = torch_csc_to_scipy_csc(A)
        D = _extract_block_diagonal(A_scipy, block_size)
        lu = splu(D)
        self.L = lu.L.tocsr()
        self.U = lu.U.tocsr()
        self.perm_c = lu.perm_c
        self.perm_r = lu.perm_r
        self.device = A.device
        self.dtype = A.dtype

    def apply(self, r: torch.Tensor) -> torch.Tensor:
        r_np = to_numpy(r)
        b = r_np.copy()
        b[self.perm_r] = r_np
        c = spsolve_triangular(self.L, b, lower=True, unit_diagonal=True)
        x = spsolve_triangular(self.U, c, lower=False)
        x = x[self.perm_c]
        return from_numpy(x, self.dtype, self.device)
