import torch
from ..solver.fgmres import FGMRES


class GMRESPreconditioner:

    def __init__(self, A: torch.Tensor, inner_iters: int = 10, inner_rtol: float = 1e-6):
        self.A = A
        self.solver = FGMRES(restart=inner_iters, max_iters=inner_iters, rtol=inner_rtol)

    def apply(self, r: torch.Tensor) -> torch.Tensor:
        result = self.solver.solve(self.A, r, M=None, progress_bar=False)
        return result.solution
