""" Test functions for the sparse.linalg._isolve module
"""

from contextlib import nullcontext
import itertools
import platform
import pytest
import types

import numpy as np
from numpy.testing import assert_allclose
from numpy import zeros, arange, array, ones, eye, iscomplexobj

from scipy._lib._array_api import (
    is_numpy, make_xp_pytest_param, make_xp_pytest_marks,
    xp_assert_close, xp_assert_less, xp_vector_norm, xp_assert_equal,
    xp_assert_less_equal,
)
from scipy._external import array_api_extra as xpx
from scipy._lib._sparse import issparse
from scipy.sparse import dia_array, csr_array, kronsum

from scipy.sparse.linalg import LinearOperator, aslinearoperator
from scipy.sparse.linalg._isolve import (bicg, bicgstab, cg, cgs,
                                         gcrotmk, gmres, lgmres,
                                         minres, qmr, tfqmr)
from scipy.sparse.linalg._interface import IdentityOperator

# TODO check that method preserve shape and type
# TODO test both preconditioner methods


# list of all solvers under test
_SOLVERS = [bicg, bicgstab, cg, cgs, gcrotmk, gmres, lgmres,
            minres, qmr, tfqmr]

CB_TYPE_FILTER = ".*called without specifying `callback_type`.*"


def _assert_success(*, A, x, b, xp, rtol=1.0, atol=0.0, less_equal=False):
    Ax = xp.squeeze(A @ x[..., xp.newaxis], axis=-1)
    residual = Ax - b
    err = xp_vector_norm(residual, axis=-1)
    assertion = xp_assert_less_equal if less_equal else xp_assert_less
    limit = atol + rtol * xp_vector_norm(b, axis=-1)
    assertion(err, limit)


# create parametrized fixture for easy reuse in tests
@pytest.fixture(
    scope="session",
    params=[make_xp_pytest_param(func) for func in _SOLVERS],
)
def solver(request):
    """
    Fixture for all solvers in scipy.sparse.linalg._isolve
    """
    return request.param


class Case:
    def __init__(self, name, A, b=None, skip=None, nonconvergence=None):
        self.name = name
        self.A = A
        if b is None:
            self.b = arange(A.shape[0], dtype=A.dtype)
        else:
            self.b = b
        if skip is None:
            self.skip = []
        else:
            self.skip = skip
        if nonconvergence is None:
            self.nonconvergence = []
        else:
            self.nonconvergence = nonconvergence


class SingleTest:
    # case with a specific solver
    def __init__(self, A, b, solver, casename, convergence=True):
        self.A = A
        self.b = b
        self.solver = solver
        self.casename = casename
        self.name = casename + '-' + solver.__name__
        self.convergence = convergence

    def __repr__(self):
        return f"<{self.name}>"


def xp_case(case, xp, batch_A, batch_b, rng=None):
    sparse = issparse(case.A)

    if (not_np := not is_numpy(xp)) and sparse: # TODO: pydata/sparse for sparse tests?
        pytest.skip("sparse tests run only with `np` backend")

    A = xp.asarray(case.A) if not_np else case.A
    b = xp.asarray(case.b) if not_np else case.b

    rng = np.random.default_rng(rng)
    if sparse and (batch_A or batch_b):
        A = A.tocoo()
    if batch_A:
        # apply some random scaling to each system in batch,
        # preserving symmetry, (non-)positive definiteness, Hermiticity
        scaling = np.abs(rng.standard_normal(batch_A + (1, 1)))
        batch_array = xp.asarray(scaling, dtype=A.dtype)
        A = batch_array * A
    if batch_b:
        # apply some random scaling to each RHS in batch
        b = b * xp.asarray(rng.random(batch_b + (1,)), dtype=b.dtype)

    return types.SimpleNamespace(
        name=case.name,
        A=A,
        b=b,
        convergence=case.convergence,
        solver=case.solver,
        casename=case.casename,
    )


class IterativeParams:
    def __init__(self):
        sym_solvers = [minres, cg]
        posdef_solvers = [cg]
        real_solvers = [minres]

        # list of Cases
        self.cases = []

        # Symmetric and Positive Definite
        N = 40
        data = ones((3, N))
        data[0, :] = 2
        data[1, :] = -1
        data[2, :] = -1
        Poisson1D = dia_array((data, [0, -1, 1]), shape=(N, N)).tocsr()
        self.cases.append(Case("poisson1d", Poisson1D))
        # note: minres fails for single precision
        self.cases.append(Case("poisson1d-F", Poisson1D.astype(np.float32),
                               skip=[minres]))

        # Symmetric and Negative Definite
        self.cases.append(Case("neg-poisson1d", -Poisson1D,
                               skip=posdef_solvers))
        # note: minres fails for single precision
        self.cases.append(Case("neg-poisson1d-F", (-Poisson1D).astype('f'),
                               skip=posdef_solvers + [minres]))

        # 2-dimensional Poisson equations
        Poisson2D = kronsum(Poisson1D, Poisson1D)
        # note: minres fails for 2-d poisson problem,
        # it will be fixed in the future PR
        self.cases.append(Case("poisson2d", Poisson2D, skip=[minres]))
        # note: minres fails for single precision
        self.cases.append(Case("poisson2d-F", Poisson2D.astype('f'),
                               skip=[minres]))

        # Symmetric and Indefinite
        data = array([[6, -5, 2, 7, -1, 10, 4, -3, -8, 9]], dtype='d')
        RandDiag = dia_array((data, [0]), shape=(10, 10)).tocsr()
        self.cases.append(Case("rand-diag", RandDiag, skip=posdef_solvers))
        self.cases.append(Case("rand-diag-F", RandDiag.astype('f'),
                               skip=posdef_solvers))

        # Random real-valued
        rng = np.random.RandomState(1234)
        data = rng.rand(4, 4)
        self.cases.append(Case("rand", data,
                               skip=posdef_solvers + sym_solvers))
        self.cases.append(Case("rand-F", data.astype('f'),
                               skip=posdef_solvers + sym_solvers))

        # Random symmetric real-valued
        rng = np.random.RandomState(1234)
        data = rng.rand(4, 4)
        data = data + data.T
        self.cases.append(Case("rand-sym", data, skip=posdef_solvers))
        self.cases.append(Case("rand-sym-F", data.astype('f'),
                               skip=posdef_solvers))

        # Random pos-def symmetric real
        np.random.seed(1234)
        data = np.random.rand(9, 9)
        data = np.dot(data.conj(), data.T)
        self.cases.append(Case("rand-sym-pd", data))
        # note: minres fails for single precision
        self.cases.append(Case("rand-sym-pd-F", data.astype('f'),
                               skip=[minres]))

        # Random complex-valued
        rng = np.random.RandomState(1234)
        data = rng.rand(4, 4) + 1j * rng.rand(4, 4)
        skip_cmplx = posdef_solvers + sym_solvers + real_solvers
        self.cases.append(Case("rand-cmplx", data, skip=skip_cmplx))
        self.cases.append(Case("rand-cmplx-F", data.astype('F'),
                               skip=skip_cmplx))

        # Random hermitian complex-valued
        rng = np.random.RandomState(1234)
        data = rng.rand(4, 4) + 1j * rng.rand(4, 4)
        data = data + data.T.conj()
        self.cases.append(Case("rand-cmplx-herm", data,
                               skip=posdef_solvers + real_solvers))
        self.cases.append(Case("rand-cmplx-herm-F", data.astype('F'),
                               skip=posdef_solvers + real_solvers))

        # Random pos-def hermitian complex-valued
        rng = np.random.RandomState(1234)
        data = rng.rand(9, 9) + 1j * rng.rand(9, 9)
        data = np.dot(data.conj(), data.T)
        self.cases.append(Case(
            "rand-cmplx-sym-pd", data, skip=real_solvers)
        )
        self.cases.append(Case("rand-cmplx-sym-pd-F", data.astype('F'),
                               skip=real_solvers))

        # Non-symmetric and Positive Definite
        #
        # cgs, qmr, bicg and tfqmr fail to converge on this one
        #   -- algorithmic limitation apparently
        data = ones((2, 10))
        data[0, :] = 2
        data[1, :] = -1
        A = dia_array((data, [0, -1]), shape=(10, 10)).tocsr()
        self.cases.append(Case("nonsymposdef", A,
                               skip=sym_solvers + [cgs, qmr, bicg, tfqmr]))
        self.cases.append(Case("nonsymposdef-F", A.astype('F'),
                               skip=sym_solvers + [cgs, qmr, bicg, tfqmr]))

        # Symmetric, non-pd, hitting cgs/bicg/bicgstab/qmr/tfqmr breakdown
        A = np.array([[0, 0, 0, 0, 0, 1, -1, -0, -0, -0, -0],
                      [0, 0, 0, 0, 0, 2, -0, -1, -0, -0, -0],
                      [0, 0, 0, 0, 0, 2, -0, -0, -1, -0, -0],
                      [0, 0, 0, 0, 0, 2, -0, -0, -0, -1, -0],
                      [0, 0, 0, 0, 0, 1, -0, -0, -0, -0, -1],
                      [1, 2, 2, 2, 1, 0, -0, -0, -0, -0, -0],
                      [-1, 0, 0, 0, 0, 0, -1, -0, -0, -0, -0],
                      [0, -1, 0, 0, 0, 0, -0, -1, -0, -0, -0],
                      [0, 0, -1, 0, 0, 0, -0, -0, -1, -0, -0],
                      [0, 0, 0, -1, 0, 0, -0, -0, -0, -1, -0],
                      [0, 0, 0, 0, -1, 0, -0, -0, -0, -0, -1]], dtype=float)
        b = np.array([0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0], dtype=float)
        assert (A == A.T).all()
        self.cases.append(Case("sym-nonpd", A, b,
                               skip=posdef_solvers,
                               nonconvergence=[cgs, bicg, bicgstab, qmr, tfqmr]
                               )
                          )

    def generate_tests(self):
        # generate test cases with skips applied
        tests = []
        for case in self.cases:
            for solver in _SOLVERS:
                if (solver in case.skip):
                    continue
                if solver in case.nonconvergence:
                    tests += [SingleTest(case.A, case.b, solver, case.name,
                                         convergence=False)]
                else:
                    tests += [
                        SingleTest(case.A, case.b, solver, case.name)
                    ]
        return tests


cases = IterativeParams().generate_tests()


@pytest.fixture(
    params=[pytest.param(c, marks=make_xp_pytest_marks(c.solver)) for c in cases],
    ids=[x.name for x in cases],
    scope="module"
)
def case(request):
    """
    Fixture for all cases in IterativeParams
    """
    return request.param


@pytest.mark.parametrize("batch_A", [()])
@pytest.mark.parametrize("batch_b", [()])
def test_maxiter(case, xp, batch_A, batch_b):
    case = xp_case(case, xp, batch_A, batch_b, rng=38)
    if not case.convergence:
        pytest.skip("Solver - Breakdown case, see gh-8829")
    A = case.A
    rtol = 1e-12

    b = case.b
    x0 = 0 * b

    residuals = []

    def callback(x):
        if x.ndim == 0:
            residuals.append(xp_vector_norm(b - case.A * x))
        else:
            Ax = xp.squeeze(case.A @ x[..., xp.newaxis], axis=-1)
            residuals.append(xp_vector_norm(b - Ax, axis=-1))

    if case.solver == gmres:
        with pytest.warns(DeprecationWarning, match=CB_TYPE_FILTER):
            x, info = case.solver(A, b, x0=x0, rtol=rtol, maxiter=1, callback=callback)
    else:
        x, info = case.solver(A, b, x0=x0, rtol=rtol, maxiter=1, callback=callback)

    assert len(residuals) == 1
    assert info == 1


@pytest.mark.parametrize("batch_A", [()])
@pytest.mark.parametrize("batch_b", [()])
def test_convergence(case, xp, batch_A, batch_b):
    if (case.solver is tfqmr) and ("poisson2d-F" in case.name):
        pytest.skip("Struggles to converge with single precision on some platforms")
    case = xp_case(case, xp, batch_A, batch_b, rng=38)
    A = case.A

    if A.dtype in (xp.float64, xp.complex128):
        rtol = 1e-8
    else:
        rtol = 1e-2

    b = case.b
    x0 = 0 * b

    x, info = case.solver(A, b, x0=x0, rtol=rtol)
    xp_assert_equal(x0, 0 * b)  # ensure that x0 is not overwritten

    if case.convergence:
        assert info == 0
        _assert_success(A=A, x=x, b=b, xp=xp, rtol=rtol)
    else:
        assert info != 0
        _assert_success(A=A, x=x, b=b, xp=xp, rtol=1, less_equal=True)


@pytest.mark.parametrize("batch_A", [()])
@pytest.mark.parametrize("batch_b", [()])
def test_precond_dummy(case, xp, batch_A, batch_b):
    dtype = case.A.dtype
    case = xp_case(case, xp, batch_A, batch_b)
    if (case.solver is cgs) and ("pd-F" in case.name):
        pytest.skip("Struggles to converge with single precision")
    if (case.solver is tfqmr) and ("poisson2d-F" in case.name):
        pytest.skip("Hits divide-by-zero with single precision")
    if not case.convergence:
        pytest.skip("Solver - Breakdown case, see gh-8829")

    rtol = 1e-8 if np.finfo(dtype).eps < 1e-8 else 1.2e-3

    A = case.A
    
    # NOTE: the following was previously uncommented as dead code --
    # was the intention to set `A = dia_array(...)`?

    # _, M, N = A.shape
    # # Ensure the diagonal elements of A are non-zero before calculating
    # # 1.0 / xp.linalg.diagonal(A)
    # diagOfA = xp.linalg.diagonal(A) if not is_numpy(xp) else A.diagonal()
    # if xp.count_nonzero(diagOfA) == diagOfA.shape[0]:
    #     dia_array(([1.0 / diagOfA], [0]), shape=(M, N))

    b = case.b
    x0 = 0 * b

    precond = IdentityOperator(shape=A.shape)

    if case.solver is qmr:
        x, info = case.solver(A, b, M1=precond, M2=precond, x0=x0, rtol=rtol)
    else:
        x, info = case.solver(A, b, M=precond, x0=x0, rtol=rtol)

    assert info == 0
    _assert_success(A=A, x=x, b=b, xp=np, rtol=rtol)

    A = aslinearoperator(A)

    x, info = case.solver(A, b, x0=x0, rtol=rtol)
    assert info == 0
    _assert_success(A=A, x=x, b=b, xp=np, rtol=rtol)


@pytest.mark.parametrize("batch_A", [()])
@pytest.mark.parametrize("batch_b", [()])
@pytest.mark.fail_slow(10)
def test_precond_inverse(case, xp, batch_A, batch_b):
    if case.casename not in ('poisson1d', 'poisson2d'):
        pytest.skip("specific to poisson1d and poisson2d cases")
    if case.solver is qmr:
        pytest.skip("skipped for qmr")
    
    case = xp_case(case, xp, batch_A, batch_b, rng=38)
    rtol = 1e-8

    def inverse(b, which=None):
        """inverse preconditioner"""
        A = case.A
        if is_numpy(xp) and not isinstance(A, np.ndarray):
            A = A.toarray()
        return xp.linalg.solve(A, b[..., np.newaxis])

    def rinverse(b, which=None):
        """inverse preconditioner"""
        A = case.A
        if is_numpy(xp) and not isinstance(A, np.ndarray):
            A = A.toarray()
        return xp.linalg.solve(A.T, b[..., np.newaxis])

    matvec_count = [0]

    def matvec(b):
        matvec_count[0] += 1
        return case.A @ b[..., np.newaxis]

    def rmatvec(b):
        matvec_count[0] += 1
        return case.A.T @ b[..., np.newaxis]

    b = case.b
    x0 = 0 * b

    A = LinearOperator(case.A.shape, matvec, rmatvec=rmatvec)
    precond = LinearOperator(case.A.shape, inverse, rmatvec=rinverse)

    # Solve with preconditioner
    matvec_count = [0]
    x, info = case.solver(A, b, M=precond, x0=x0, rtol=rtol)

    assert info == 0
    _assert_success(A=case.A, x=x, b=b, xp=xp, rtol=rtol)

    # Solution should be nearly instant
    assert matvec_count[0] <= 3


@pytest.mark.skip_xp_backends("dask.array", reason="https://github.com/dask/dask/issues/11711")
@pytest.mark.parametrize("batch_A", [()])
@pytest.mark.parametrize("batch_b", [()])
def test_atol(solver, xp, batch_A, batch_b):
    # TODO: minres / tfqmr. It didn't historically use absolute tolerances, so
    # fixing it is less urgent.
    if solver in (minres, tfqmr):
        pytest.skip("TODO: Add atol to minres/tfqmr")

    # Historically this is tested as below, all pass but for some reason
    # gcrotmk is over-sensitive to difference between random.seed/rng.random
    # Hence tol lower bound is changed from -10 to -9
    # np.random.seed(1234)
    # A = np.random.rand(10, 10)
    # A = A @ A.T + 10 * np.eye(10)
    # b = 1e3*np.random.rand(10)

    rng = np.random.default_rng(168441431005389)
    A = rng.uniform(size=(*batch_A, 10, 10))
    A = A @ A.mT + 10*np.eye(10)
    b = 1e3 * rng.uniform(size=(*batch_b, 10))

    dtype = xpx.default_dtype(xp)
    A = xp.asarray(A, dtype=dtype) 
    b = xp.asarray(b, dtype=dtype)

    tols = np.r_[0, np.logspace(-9, 2, 7), np.inf]

    # Check effect of badly scaled preconditioners
    M0 = rng.standard_normal(size=(*batch_A, 10, 10))
    M0 = xp.asarray(M0)
    M0 = M0 @ M0.mT
    Ms = [None, 1e-6 * M0, 1e6 * M0]

    for M, rtol, atol in itertools.product(Ms, tols, tols):
        if rtol == 0 and atol == 0:
            continue

        if solver is qmr:
            if M is not None:
                M = aslinearoperator(M)
                M2 = IdentityOperator(shape=(*batch_A, 10, 10))
            else:
                M2 = None
            x, info = solver(A, b, M1=M, M2=M2, rtol=rtol, atol=atol)
        else:
            x, info = solver(A, b, M=M, rtol=rtol, atol=atol)

        assert info == 0
        _assert_success(A=A, x=x, b=b, xp=xp, rtol=rtol, atol=atol)


@pytest.mark.skip_xp_backends("dask.array", reason="https://github.com/dask/dask/issues/11711")
@pytest.mark.parametrize("batch_A", [()])
@pytest.mark.parametrize("batch_b", [()])
def test_zero_rhs(solver, xp, batch_A, batch_b):
    rng = np.random.default_rng(1684414984100503)
    dtype = xpx.default_dtype(xp) 
    A = xp.asarray(rng.random(size=(*batch_A, 10, 10)), dtype=dtype)
    A = A @ A.mT + 10 * xp.eye(10)

    b = xp.zeros((*batch_b, 10), dtype=dtype)
    tols = np.r_[np.logspace(-10, 2, 7)]

    expected = xp.broadcast_to(b, (*np.broadcast_shapes(batch_A, batch_b), 10))
    for tol in tols:
        tol = float(tol)
        x, info = solver(A, b, rtol=tol)
        assert info == 0
        xp_assert_close(x, expected, atol=1e-15)

        x, info = solver(A, b, rtol=tol, x0=xp.ones((*batch_b, 10)))
        assert info == 0
        xp_assert_close(x, expected, atol=tol)

        if solver is not minres:
            x, info = solver(A, b, rtol=tol, atol=0.0, x0=xp.ones((*batch_b, 10)))
            if info == 0:
                xp_assert_close(x, expected)

            x, info = solver(A, b, rtol=tol, atol=tol)
            assert info == 0
            xp_assert_close(x, expected, atol=1e-300)

            x, info = solver(A, b, rtol=tol, atol=0.0)
            assert info == 0
            xp_assert_close(x, expected, atol=1e-300)


# unbatched test
@pytest.mark.xfail(reason="see gh-18697")
def test_maxiter_worsening(solver, xp):
    if solver not in (gmres, lgmres, qmr):
        # these were skipped from the very beginning, see gh-9201; gh-14160
        pytest.skip("Solver breakdown case")
    # Check error does not grow (boundlessly) with increasing maxiter.
    # This can occur due to the solvers hitting close to breakdown,
    # which they should detect and halt as necessary.
    # cf. gh-9100
    if (solver is lgmres and
            platform.machine() not in ['x86_64' 'x86', 'aarch64', 'arm64']):
        # see gh-17839
        pytest.xfail(reason="fails on at least ppc64le, ppc64 and riscv64")

    # Singular matrix, rhs numerically not in range
    A = np.array([[-0.1112795288033378, 0, 0, 0.16127952880333685],
                  [0, -0.13627952880333782 + 6.283185307179586j, 0, 0],
                  [0, 0, -0.13627952880333782 - 6.283185307179586j, 0],
                  [0.1112795288033368, 0j, 0j, -0.16127952880333785]])
    v = np.ones(4)
    dtype = xpx.default_dtype(xp)
    A, v = (xp.asarray(arr, dtype=dtype) for arr in [A, v])
    best_error = np.inf

    # Unable to match the Fortran code tolerance levels with this example
    # Original tolerance values

    # slack_tol = 7 if platform.machine() == 'aarch64' else 5
    slack_tol = 9

    rtol = 1e-8
    for maxiter in range(1, 20):
        x, info = solver(A, v, maxiter=maxiter, rtol=rtol, atol=0.0)
        if info == 0:
            _assert_success(A=A, x=x, b=v, xp=xp, rtol=rtol)

        # Check with slack
        error = xp_vector_norm(A @ x - v)
        best_error = xp.min(best_error, error)
        assert error <= slack_tol * best_error


def _setup_random_system(xp, batch_A, batch_b):
    rng = np.random.default_rng(1685363802304750)
    n = 10
    A = rng.random(size=(*batch_A, n, n))
    A = A @ A.mT
    b = rng.random((*batch_b, n))
    x0 = rng.random((*batch_b, n))
    dtype = xpx.default_dtype(xp)
    A, b, x0 = (xp.asarray(arr, dtype=dtype) for arr in [A, b, x0])
    return A, b, x0


@pytest.mark.parametrize("batch_A", [()])
@pytest.mark.parametrize("batch_b", [()])
def test_x0_working(solver, xp, batch_A, batch_b):
    # Easy problem
    A, b, x0 = _setup_random_system(xp, batch_A, batch_b)

    if solver is minres:
        kw = dict(rtol=1e-6)
    else:
        kw = dict(atol=0.0, rtol=1e-6)

    x, info = solver(A, b, **kw)
    assert info == 0

    _assert_success(A=A, x=x, b=b, xp=xp, rtol=1e-6)

    x, info = solver(A, b, x0=x0, **kw)
    assert info == 0
    _assert_success(A=A, x=x, b=b, xp=xp, rtol=1e-5)


@pytest.mark.parametrize("batch_A", [()])
@pytest.mark.parametrize("batch_b", [()])
def test_x0_equals_Mb(case, xp, batch_A, batch_b):
    dtype = case.A.dtype
    case = xp_case(case, xp, batch_A, batch_b, rng=38)
    if (case.solver is cgs) and ("pd-F" in case.name):
        pytest.skip("Struggles to converge with single precision")
    if (case.solver is bicgstab) and (case.name == 'nonsymposdef-bicgstab'):
        pytest.skip("Solver fails due to numerical noise "
                    "on some architectures (see gh-15533).")
    if case.solver is tfqmr:
        pytest.skip("Solver does not support x0='Mb'")

    A = case.A
    b = case.b
    x0 = 'Mb'
    rtol = 1e-8 if np.finfo(dtype).eps < 1e-8 else 1.5e-3
    x, info = case.solver(A, b, x0=x0, rtol=rtol)

    assert x0 == 'Mb'  # ensure that x0 is not overwritten
    assert info == 0
    _assert_success(A=A, x=x, b=b, xp=xp, rtol=rtol)


@pytest.mark.parametrize("batch_A", [()])
@pytest.mark.parametrize("batch_b", [()])
def test_x0_solves_problem_exactly(solver, xp, batch_A, batch_b):
    # See gh-19948
    dtype = xpx.default_dtype(xp)
    mat = xp.tile(xp.eye(2, dtype=dtype), (*batch_A, 1, 1))
    rhs = xp.tile(xp.asarray([-1., -1.], dtype=dtype), (*batch_b, 1))

    sol, info = solver(mat, rhs, x0=rhs)
    expected = xp.broadcast_to(rhs, (*np.broadcast_shapes(batch_A, batch_b), 2))
    xp_assert_close(sol, expected)
    assert info == 0


@pytest.mark.parametrize("batch_A", [()])
@pytest.mark.parametrize("batch_b", [()])
def test_show(case, capsys, xp, batch_A, batch_b):
    if case.solver != tfqmr:
        pytest.skip("tfqmr specific test")
    if "poisson2d-F" in case.name:
        pytest.skip("Hits divide-by-zero with single precision")
    if "pd-F" in case.name:
        pytest.skip("Struggles to converge with single precision on some platforms")
    case = xp_case(case, xp, batch_A, batch_b)
    def cb(x):
        pass

    ctx = (
        np.errstate(all='ignore')
        if case.casename == "nonsymposdef"
        else nullcontext()
    )
    with ctx:
        x, info = tfqmr(case.A, case.b, callback=cb, show=True)

    out, err = capsys.readouterr()

    if case.casename == "sym-nonpd":
        # no logs for some reason
        exp = ""
    elif case.casename in ("nonsymposdef", "nonsymposdef-F"):
        # Asymmetric and Positive Definite
        exp = "TFQMR: Linear solve not converged due to reach MAXIT iterations"
    else:  # all other cases
        exp = "TFQMR: Linear solve converged due to reach TOL iterations"

    assert out.startswith(exp)
    assert err == ""


@pytest.mark.parametrize("batch_A", [()])
@pytest.mark.parametrize("batch_b", [()])
def test_positional_error(solver, xp, batch_A, batch_b):
    A, b, x0 = _setup_random_system(xp, batch_A, batch_b)
    with pytest.raises(TypeError):
        solver(A, b, x0, 1e-5)


@pytest.mark.parametrize("batch_A", [()])
@pytest.mark.parametrize("batch_b", [()])
@pytest.mark.parametrize("atol", ["legacy", None, -1])
def test_invalid_atol(solver, atol, xp, batch_A, batch_b):
    if solver == minres:
        pytest.skip("minres has no `atol` argument")
    A, b, x0 = _setup_random_system(xp, batch_A, batch_b)
    with pytest.raises(ValueError):
        solver(A, b, x0, atol=atol)


class TestQMR:
    @pytest.mark.filterwarnings('ignore::scipy.sparse.SparseEfficiencyWarning')
    def test_leftright_precond(self):
        """Check that QMR works with left and right preconditioners"""

        from scipy.sparse.linalg._dsolve import splu
        from scipy.sparse.linalg._interface import LinearOperator

        n = 100

        dat = ones(n)
        A = dia_array(([-2 * dat, 4 * dat, -dat], [-1, 0, 1]), shape=(n, n))
        b = arange(n, dtype='d')

        L = dia_array(([-dat / 2, dat], [-1, 0]), shape=(n, n))
        U = dia_array(([4 * dat, -dat], [0, 1]), shape=(n, n))
        L_solver = splu(L)
        U_solver = splu(U)

        def L_solve(b):
            return L_solver.solve(b)

        def U_solve(b):
            return U_solver.solve(b)

        def LT_solve(b):
            return L_solver.solve(b, 'T')

        def UT_solve(b):
            return U_solver.solve(b, 'T')

        M1 = LinearOperator((n, n), matvec=L_solve, rmatvec=LT_solve)
        M2 = LinearOperator((n, n), matvec=U_solve, rmatvec=UT_solve)

        rtol = 1e-8
        x, info = qmr(A, b, rtol=rtol, maxiter=15, M1=M1, M2=M2)

        assert info == 0
        assert xp_vector_norm(A @ x - b) <= rtol * xp_vector_norm(b)


class TestGMRES:
    def test_basic(self):
        A = np.vander(np.arange(10) + 1)[:, ::-1]
        b = np.zeros(10)
        b[0] = 1

        x_gm, err = gmres(A, b, restart=5, maxiter=1)

        assert_allclose(x_gm[0], 0.359, rtol=1e-2)

    @pytest.mark.filterwarnings(f"ignore:{CB_TYPE_FILTER}:DeprecationWarning")
    def test_callback(self):

        def store_residual(r, rvec):
            rvec[rvec.nonzero()[0].max() + 1] = r

        # Define, A,b
        A = csr_array(array([[-2, 1, 0, 0, 0, 0],
                             [1, -2, 1, 0, 0, 0],
                             [0, 1, -2, 1, 0, 0],
                             [0, 0, 1, -2, 1, 0],
                             [0, 0, 0, 1, -2, 1],
                             [0, 0, 0, 0, 1, -2]]))
        b = ones((A.shape[0],))
        maxiter = 1
        rvec = zeros(maxiter + 1)
        rvec[0] = 1.0

        def callback(r):
            return store_residual(r, rvec)

        x, flag = gmres(A, b, x0=zeros(A.shape[0]), rtol=1e-16,
                        maxiter=maxiter, callback=callback)

        # Expected output from SciPy 1.0.0
        assert_allclose(rvec, array([1.0, 0.81649658092772603]), rtol=1e-10)

        # Test preconditioned callback
        M = 1e-3 * np.eye(A.shape[0])
        rvec = zeros(maxiter + 1)
        rvec[0] = 1.0
        x, flag = gmres(A, b, M=M, rtol=1e-16, maxiter=maxiter,
                        callback=callback)

        # Expected output from SciPy 1.0.0
        # (callback has preconditioned residual!)
        assert_allclose(rvec, array([1.0, 1e-3 * 0.81649658092772603]),
                        rtol=1e-10)

    def test_abi(self):
        # Check we don't segfault on gmres with complex argument
        A = eye(2)
        b = ones(2)
        r_x, r_info = gmres(A, b)
        r_x = r_x.astype(complex)
        x, info = gmres(A.astype(complex), b.astype(complex))

        assert iscomplexobj(x)
        assert_allclose(r_x, x)
        assert r_info == info

    @pytest.mark.fail_slow(10)
    def test_atol_legacy(self):

        A = eye(2)
        b = ones(2)
        x, info = gmres(A, b, rtol=1e-5)
        assert xp_vector_norm(A @ x - b) <= 1e-5 * xp_vector_norm(b)
        assert_allclose(x, b, atol=0.0, rtol=1e-8)

        rndm = np.random.RandomState(12345)
        A = rndm.rand(30, 30)
        b = 1e-6 * ones(30)
        x, info = gmres(A, b, rtol=1e-7, restart=20)
        assert xp_vector_norm(A @ x - b) > 1e-7

        A = eye(2)
        b = 1e-10 * ones(2)
        x, info = gmres(A, b, rtol=1e-8, atol=0.0)
        assert xp_vector_norm(A @ x - b) <= 1e-8 * xp_vector_norm(b)

    def test_defective_precond_breakdown(self):
        # Breakdown due to defective preconditioner
        M = np.eye(3)
        M[2, 2] = 0

        b = np.array([0, 1, 1])
        x = np.array([1, 0, 0])
        A = np.diag([2, 3, 4])

        x, info = gmres(A, b, x0=x, M=M, rtol=1e-15, atol=0.0)

        # Should not return nans, nor terminate with false success
        assert not np.isnan(x).any()
        if info == 0:
            assert xp_vector_norm(A @ x - b) <= 1e-15 * xp_vector_norm(b)

        # The solution should be OK outside null space of M
        assert_allclose(M @ (A @ x), M @ b)

    def test_defective_matrix_breakdown(self):
        # Breakdown due to defective matrix
        A = np.array([[0, 1, 0], [1, 0, 0], [0, 0, 0]])
        b = np.array([1, 0, 1])
        rtol = 1e-8
        x, info = gmres(A, b, rtol=rtol, atol=0)

        # Should not return nans, nor terminate with false success
        assert not np.isnan(x).any()
        if info == 0:
            assert xp_vector_norm(A @ x - b) <= rtol * xp_vector_norm(b)

        # The solution should be OK outside null space of A
        assert_allclose(A @ (A @ x), A @ b)

    @pytest.mark.filterwarnings(f"ignore:{CB_TYPE_FILTER}:DeprecationWarning")
    def test_callback_type(self):
        # The legacy callback type changes meaning of 'maxiter'
        np.random.seed(1)
        A = np.random.rand(20, 20)
        b = np.random.rand(20)

        cb_count = [0]

        def pr_norm_cb(r):
            cb_count[0] += 1
            assert isinstance(r, float)

        def x_cb(x):
            cb_count[0] += 1
            assert isinstance(x, np.ndarray)

        # 2 iterations is not enough to solve the problem
        cb_count = [0]
        x, info = gmres(A, b, rtol=1e-6, atol=0.0, callback=pr_norm_cb,
                        maxiter=2, restart=50)
        assert info == 2
        assert cb_count[0] == 2

        # With `callback_type` specified, no warning should be raised
        cb_count = [0]
        x, info = gmres(A, b, rtol=1e-6, atol=0.0, callback=pr_norm_cb,
                        maxiter=2, restart=50, callback_type='legacy')
        assert info == 2
        assert cb_count[0] == 2

        # 2 restart cycles is enough to solve the problem
        cb_count = [0]
        x, info = gmres(A, b, rtol=1e-6, atol=0.0, callback=pr_norm_cb,
                        maxiter=2, restart=50, callback_type='pr_norm')
        assert info == 0
        assert cb_count[0] > 2

        # 2 restart cycles is enough to solve the problem
        cb_count = [0]
        x, info = gmres(A, b, rtol=1e-6, atol=0.0, callback=x_cb, maxiter=2,
                        restart=50, callback_type='x')
        assert info == 0
        assert cb_count[0] == 1

    def test_callback_x_monotonic(self):
        # Check that callback_type='x' gives monotonic norm decrease
        rng = np.random.RandomState(1)
        A = rng.rand(20, 20) + np.eye(20)
        b = rng.rand(20)

        prev_r = [np.inf]
        count = [0]

        def x_cb(x):
            r = xp_vector_norm(A @ x - b)
            assert r <= prev_r[0]
            prev_r[0] = r
            count[0] += 1

        x, info = gmres(A, b, rtol=1e-6, atol=0.0, callback=x_cb, maxiter=20,
                        restart=10, callback_type='x')
        assert info == 20
        assert count[0] == 20


def test_nD(solver, xp):
    """Check that >2-D operators are rejected cleanly."""
    def id(x):
        return x
    A = LinearOperator(shape=(2, 2, 2), matvec=id, dtype=xp.float64, xp=xp)
    b = xp.ones((2, 2))
    with pytest.raises(ValueError, match="expected 2-D"):
        solver(A, b)