import torch
from scipy.sparse.linalg import spilu, spsolve_triangular

from .sparse_convert import torch_csc_to_scipy_csc, to_numpy, from_numpy


class ILU:

    def __init__(self, A: torch.Tensor, fill_factor: float | None = None):
        A_scipy = torch_csc_to_scipy_csc(A)

        kwargs = {}
        if fill_factor is not None:
            kwargs["fill_factor"] = fill_factor

        ilu = spilu(A_scipy, **kwargs)
        self.L = ilu.L.tocsr()
        self.U = ilu.U.tocsr()
        self.perm_c = ilu.perm_c
        self.perm_r = ilu.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)
