import torch
from pyamg.blackbox import solver as amg_solver, solver_configuration
from pyamg.classical import air_solver

from .sparse_convert import torch_csc_to_scipy_csr, to_numpy, from_numpy


class AMG:

    def __init__(self, A: torch.Tensor):
        A_csr = torch_csc_to_scipy_csr(A)
        config = solver_configuration(A_csr, verb=False)
        self.M = amg_solver(A_csr, config).aspreconditioner()
        self.device = A.device
        self.dtype = A.dtype

    def apply(self, r: torch.Tensor) -> torch.Tensor:
        z_np = self.M * to_numpy(r)
        return from_numpy(z_np, self.dtype, self.device)


class AMG_AIR:

    def __init__(self, A: torch.Tensor):
        A_csr = torch_csc_to_scipy_csr(A)
        config = solver_configuration(A_csr, verb=False)

        if config["symmetry"] == "hermitian":
            mg = amg_solver(A_csr, config)
        else:
            mg = air_solver(
                A_csr,
                presmoother=config["presmoother"],
                postsmoother=config["postsmoother"],
            )

        self.M = mg.aspreconditioner()
        self.device = A.device
        self.dtype = A.dtype

    def apply(self, r: torch.Tensor) -> torch.Tensor:
        z_np = self.M * to_numpy(r)
        return from_numpy(z_np, self.dtype, self.device)
