#!/usr/bin/env python
"""Tests for the linalg._isolve.gcrotmk module
"""
import threading
import warnings

from numpy.testing import (assert_, assert_allclose, assert_equal)

import numpy as np
from numpy import zeros, array, allclose
from scipy.linalg import norm
from scipy.sparse import csr_array, eye_array, random_array

from scipy.sparse.linalg._interface import LinearOperator
from scipy.sparse.linalg import splu
from scipy.sparse.linalg._isolve import gcrotmk, gmres


Am = csr_array(array([[-2,1,0,0,0,9],
                       [1,-2,1,0,5,0],
                       [0,1,-2,1,0,0],
                       [0,0,1,-2,1,0],
                       [0,3,0,1,-2,1],
                       [1,0,0,0,1,-2]]))
b = array([1,2,3,4,5,6])
count = threading.local()  # [0]
niter = threading.local()  # [0]


def matvec(v):
    if not hasattr(count, 'c'):
        count.c = [0]
    count.c[0] += 1
    return Am@v


def cb(v):
    if not hasattr(niter, 'n'):
        niter.n = [0]
    niter.n[0] += 1


A = LinearOperator(matvec=matvec, shape=Am.shape, dtype=Am.dtype)


def do_solve(**kw):
    if not hasattr(niter, 'n'):
        niter.n = [0]

    if not hasattr(count, 'c'):
        count.c = [0]

    count.c[0] = 0
    with warnings.catch_warnings():
        warnings.filterwarnings(
            "ignore", ".*called without specifying.*", DeprecationWarning)
        x0, flag = gcrotmk(A, b, x0=zeros(A.shape[0]), rtol=1e-14, **kw)
    count_0 = count.c[0]
    assert_(allclose(A@x0, b, rtol=1e-12, atol=1e-12), norm(A@x0-b))
    return x0, count_0


class TestGCROTMK:
    def test_preconditioner(self):
        # Check that preconditioning works
        pc = splu(Am.tocsc())
        M = LinearOperator(matvec=pc.solve, shape=A.shape, dtype=A.dtype)

        x0, count_0 = do_solve()
        niter.n[0] = 0
        x1, count_1 = do_solve(M=M, callback=cb)

        assert_equal(count_1, 3)
        assert count_1 < count_0/2
        assert allclose(x1, x0, rtol=1e-14)
        assert niter.n[0] < 3

    def test_arnoldi(self):
        rng = np.random.default_rng(1)

        A = eye_array(2000) + random_array((2000, 2000), density=5e-4, rng=rng)
        b = rng.random(2000)

        # The inner arnoldi should be equivalent to gmres
        with warnings.catch_warnings():
            warnings.filterwarnings(
                "ignore", ".*called without specifying.*", DeprecationWarning)
            x0, flag0 = gcrotmk(A, b, x0=zeros(A.shape[0]), m=10, k=0, maxiter=1)
            x1, flag1 = gmres(A, b, x0=zeros(A.shape[0]), restart=10, maxiter=1)

        assert_equal(flag0, 1)
        assert_equal(flag1, 1)
        assert np.linalg.norm(A.dot(x0) - b) > 1e-4

        assert_allclose(x0, x1)

    def test_cornercase(self):
        np.random.seed(1234)

        # Rounding error may prevent convergence with tol=0 --- ensure
        # that the return values in this case are correct, and no
        # exceptions are raised

        for n in [3, 5, 10, 100]:
            A = 2*eye_array(n)

            with warnings.catch_warnings():
                warnings.filterwarnings(
                    "ignore", ".*called without specifying.*", DeprecationWarning)
                b = np.ones(n)
                x, info = gcrotmk(A, b, maxiter=10)
                assert_equal(info, 0)
                assert_allclose(A.dot(x) - b, 0, atol=1e-14)

                x, info = gcrotmk(A, b, rtol=0, maxiter=10)
                if info == 0:
                    assert_allclose(A.dot(x) - b, 0, atol=1e-14)

                b = np.random.rand(n)
                x, info = gcrotmk(A, b, maxiter=10)
                assert_equal(info, 0)
                assert_allclose(A.dot(x) - b, 0, atol=1e-14)

                x, info = gcrotmk(A, b, rtol=0, maxiter=10)
                if info == 0:
                    assert_allclose(A.dot(x) - b, 0, atol=1e-14)

    def test_nans(self):
        A = eye_array(3, format='lil')
        A[1,1] = np.nan
        b = np.ones(3)

        with warnings.catch_warnings():
            warnings.filterwarnings(
                "ignore", ".*called without specifying.*", DeprecationWarning)
            x, info = gcrotmk(A, b, rtol=0, maxiter=10)
            assert_equal(info, 1)

    def test_truncate(self):
        np.random.seed(1234)
        A = np.random.rand(30, 30) + np.eye(30)
        b = np.random.rand(30)

        for truncate in ['oldest', 'smallest']:
            with warnings.catch_warnings():
                warnings.filterwarnings(
                    "ignore", ".*called without specifying.*", DeprecationWarning)
                x, info = gcrotmk(A, b, m=10, k=10, truncate=truncate,
                                  rtol=1e-4, maxiter=200)
            assert_equal(info, 0)
            assert_allclose(A.dot(x) - b, 0, atol=1e-3)

    def test_CU(self):
        for discard_C in (True, False):
            # Check that C,U behave as expected
            CU = []
            x0, count_0 = do_solve(CU=CU, discard_C=discard_C)
            assert_(len(CU) > 0)
            assert_(len(CU) <= 6)

            if discard_C:
                for c, u in CU:
                    assert_(c is None)

            # should converge immediately
            x1, count_1 = do_solve(CU=CU, discard_C=discard_C)
            if discard_C:
                assert_equal(count_1, 2 + len(CU))
            else:
                assert_equal(count_1, 3)
            assert_(count_1 <= count_0/2)
            assert_allclose(x1, x0, atol=1e-14)

    def test_denormals(self):
        # Check that no warnings are emitted if the matrix contains
        # numbers for which 1/x has no float representation, and that
        # the solver behaves properly.
        A = np.array([[1, 2], [3, 4]], dtype=float)
        A *= 100 * np.nextafter(0, 1)

        b = np.array([1, 1])

        with warnings.catch_warnings():
            warnings.filterwarnings(
                "ignore", ".*called without specifying.*", DeprecationWarning)
            xp, info = gcrotmk(A, b)

        if info == 0:
            assert_allclose(A.dot(xp), b)