"""Test functions for the sparse.linalg._krylov_funm module."""
from functools import partial

import numpy as np
import pytest
from numpy.testing import (assert_allclose)
import scipy.sparse
import scipy.linalg
from scipy.sparse.linalg import aslinearoperator
from scipy.linalg import (expm, cosm, coshm, sinm, sinhm)

from scipy.sparse.linalg._funm_multiply_krylov import funm_multiply_krylov

REAL_DTYPES = (np.float32, np.float64)
COMPLEX_DTYPES = (np.complex64, np.complex128)
DTYPES = REAL_DTYPES + COMPLEX_DTYPES

# This the phi_1 function from exponential integrators: phi_1 (z) = (e^z - 1) / z
def custom(X):
    return scipy.linalg.solve(X, expm(X) - np.eye(X.shape[0]))

FUNCS = (expm, cosm, coshm, sinm, sinhm, custom)

class TestKrylovFunmv:

    def test_krylov_funm_zero_vector(self):
        n = 20
        A = np.zeros((n, n))
        b = np.zeros(n)
        observed = funm_multiply_krylov(expm, A, b)
        expected = np.zeros(n)
        assert_allclose(observed, expected)

    @pytest.mark.parametrize("f", FUNCS)
    def test_funm_multiply_krylov_nonhermitian_dense(self, f):
        rng = np.random.default_rng(1738151906092735)
        n = 60
        nsamples = 10

        for i in range(nsamples):
            A = rng.standard_normal((n, n))
            b = rng.standard_normal(n)

            fA = f(A)
            expected = fA @ b
            observed = funm_multiply_krylov(f, A, b)
            assert_allclose(observed, expected, rtol = 1E-6, atol = 1E-8)
            observed = funm_multiply_krylov(f, aslinearoperator(A), b)
            assert_allclose(observed, expected, rtol = 1E-6, atol = 1E-8)

    @pytest.mark.parametrize("f", FUNCS)
    def test_funm_multiply_krylov_nonhermitian_sparse(self, f, num_parallel_threads):

        rng = np.random.default_rng(1738151906092735)
        n = 100
        nsamples = 1 + 9 // num_parallel_threads  # Very slow otherwise

        for i in range(nsamples):
            D = scipy.sparse.diags_array(rng.standard_normal(n))
            A = scipy.sparse.random_array((n, n), density = 0.01, rng = rng) + D
            denseA = A.todense()
            b = rng.standard_normal(n)

            fA = f(denseA)
            expected = fA @ b
            observed = funm_multiply_krylov(f, A, b)
            assert_allclose(observed, expected, rtol = 1E-6, atol = 1E-8)
            observed = funm_multiply_krylov(f, aslinearoperator(A), b)
            assert_allclose(observed, expected, rtol = 1E-6, atol = 1E-8)

    @pytest.mark.parametrize("f", FUNCS)
    def test_funm_multiply_krylov_hermitian_dense(self, f):

        rng = np.random.default_rng(1738151906092735)
        n = 60
        nsamples = 10

        for i in range(nsamples):
            R = np.triu(rng.standard_normal((n, n)))
            A = R.T + R
            b = rng.standard_normal(n)

            fA = f(A)
            expected = fA @ b
            observed = funm_multiply_krylov(f, A, b, assume_a = 'her')
            assert_allclose(observed, expected, rtol = 1E-6, atol = 1E-8)
            observed = funm_multiply_krylov(f, aslinearoperator(A), b, assume_a = 'her')
            assert_allclose(observed, expected, rtol = 1E-6, atol = 1E-8)

    @pytest.mark.parametrize("f", FUNCS)
    def test_funm_multiply_krylov_hermitian_sparse(self, f, num_parallel_threads):

        rng = np.random.default_rng(1738151906092735)
        n = 100
        nsamples = 1 + 9 // num_parallel_threads  # Very slow otherwise

        for i in range(nsamples):
            D = scipy.sparse.diags_array(rng.standard_normal(n))
            A = scipy.sparse.random_array((n, n), density = 0.01, rng = rng)
            R = scipy.sparse.triu(A)
            A = R + R.T + D
            denseA = A.todense()
            b = rng.standard_normal(n)

            fA = f(denseA)
            expected = fA @ b
            observed = funm_multiply_krylov(f, A, b, assume_a = 'her')
            assert_allclose(observed, expected, rtol = 1E-6, atol = 1E-8)

            observed = funm_multiply_krylov(f, aslinearoperator(A), b, assume_a = 'her')
            assert_allclose(observed, expected, rtol = 1E-6, atol = 1E-8)

    def test_funm_multiply_krylov_breakdown(self):
        rng = np.random.default_rng(1738151906092735)

        # From test_iterative
        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 = rng.standard_normal(A.shape[0])

        fA = expm(A)
        expected = fA @ b
        observed = funm_multiply_krylov(expm, A, b, restart_every_m = 40)
        assert_allclose(observed, expected)

    def test_funm_multiply_krylov_invalid_input(self):
            A = np.array([[1, 2], [3, 4]])  # Non-hermitian matrix
            b = np.array([1.0, 2.0])  # Ensure 'b' is a 1D array of floats

            # Test for invalid 'b' (not 1D)
            b_invalid = np.array([[1.0], [2.0]])  # 2D array
            with pytest.raises(ValueError, 
                    match="argument 'b' must be a 1D array."): 
                funm_multiply_krylov(np.exp, A, b_invalid)

            # Test for invalid restart parameter
            with pytest.raises(ValueError, 
                    match="argument 'restart_every_m' must be positive."): 
                funm_multiply_krylov(np.exp, A, b, restart_every_m=0)

            # Test for invalid max_restarts
            with pytest.raises(ValueError, 
                    match="argument 'max_restarts' must be positive."): 
                funm_multiply_krylov(np.exp, A, b, max_restarts=0)

            # Test for invalid 'assume_a' string
            with pytest.raises(ValueError, 
                    match="is not a recognized matrix structure"): 
                funm_multiply_krylov(np.exp, A, b, assume_a='invalid')


@pytest.mark.parametrize("dtype_a", DTYPES)
@pytest.mark.parametrize("dtype_b", DTYPES)
def test_funm_multiply_krylov_types(dtype_a, dtype_b):
    assert_allclose_ = (partial(assert_allclose, rtol = 1.8e-3, atol = 1e-5)
                        if {dtype_a, dtype_b} else assert_allclose)

    rng = np.random.default_rng(1738151906092735)
    n = 50

    if dtype_a in REAL_DTYPES:
        A = rng.random([n, n]).astype(dtype_a)
    else:
        A = (rng.random([n, n]) + 1j * rng.random([n, n])).astype(dtype_a)

    if dtype_b in REAL_DTYPES:
        b = (2 * rng.random(n)).astype(dtype_b)
    else:
        b = (rng.random(n) + 1j * rng.random(n)).astype(dtype_b)

        expA = expm(A)
        expected = expA @ b
        observed = funm_multiply_krylov(expm, A, b)
        assert_allclose_(observed, expected)
        observed = funm_multiply_krylov(expm, aslinearoperator(A), b)
        assert_allclose_(observed, expected)
