import torch


class Jacobi:

    def __init__(self, A: torch.Tensor):
        if A.is_sparse or A.layout == torch.sparse_csc:
            A_coo = A.to_sparse_coo().coalesce()
            indices = A_coo.indices()
            values = A_coo.values()
            diag_mask = indices[0] == indices[1]
            self.D = torch.zeros(A.shape[0], dtype=A.dtype, device=A.device)
            self.D.scatter_add_(0, indices[0, diag_mask], values[diag_mask])
        else:
            self.D = torch.diagonal(A).clone()

        zero_mask = self.D.abs() < 1e-15
        if zero_mask.any():
            raise ValueError(f"Matrix has {zero_mask.sum().item()} zero diagonal entries, Jacobi not applicable")

    def apply(self, r: torch.Tensor) -> torch.Tensor:
        return r / self.D.unsqueeze(-1) if r.dim() > 1 else r / self.D
