#
# Created by: Pearu Peterson, March 2002
#
""" Test functions for scipy.linalg.matfuncs module

"""
from __future__ import division, print_function, absolute_import

import math

import numpy as np
from numpy import array, eye, exp, random
from numpy.linalg import matrix_power
from numpy.testing import (
        assert_allclose, assert_, assert_array_almost_equal, assert_equal,
        assert_array_almost_equal_nulp)
from scipy._lib._numpy_compat import suppress_warnings

from scipy.sparse import csc_matrix, SparseEfficiencyWarning
from scipy.sparse.construct import eye as speye
from scipy.sparse.linalg.matfuncs import (expm, _expm,
        ProductOperator, MatrixPowerOperator,
        _onenorm_matrix_power_nnm)
from scipy.sparse.sputils import matrix
from scipy.linalg import logm
from scipy.special import factorial, binom
import scipy.sparse
import scipy.sparse.linalg


def _burkardt_13_power(n, p):
    """
    A helper function for testing matrix functions.

    Parameters
    ----------
    n : integer greater than 1
        Order of the square matrix to be returned.
    p : non-negative integer
        Power of the matrix.

    Returns
    -------
    out : ndarray representing a square matrix
        A Forsythe matrix of order n, raised to the power p.

    """
    # Input validation.
    if n != int(n) or n < 2:
        raise ValueError('n must be an integer greater than 1')
    n = int(n)
    if p != int(p) or p < 0:
        raise ValueError('p must be a non-negative integer')
    p = int(p)

    # Construct the matrix explicitly.
    a, b = divmod(p, n)
    large = np.power(10.0, -n*a)
    small = large * np.power(10.0, -n)
    return np.diag([large]*(n-b), b) + np.diag([small]*b, b-n)


def test_onenorm_matrix_power_nnm():
    np.random.seed(1234)
    for n in range(1, 5):
        for p in range(5):
            M = np.random.random((n, n))
            Mp = np.linalg.matrix_power(M, p)
            observed = _onenorm_matrix_power_nnm(M, p)
            expected = np.linalg.norm(Mp, 1)
            assert_allclose(observed, expected)


class TestExpM(object):
    def test_zero_ndarray(self):
        a = array([[0.,0],[0,0]])
        assert_array_almost_equal(expm(a),[[1,0],[0,1]])

    def test_zero_sparse(self):
        a = csc_matrix([[0.,0],[0,0]])
        assert_array_almost_equal(expm(a).toarray(),[[1,0],[0,1]])

    def test_zero_matrix(self):
        a = matrix([[0.,0],[0,0]])
        assert_array_almost_equal(expm(a),[[1,0],[0,1]])

    def test_misc_types(self):
        A = expm(np.array([[1]]))
        assert_allclose(expm(((1,),)), A)
        assert_allclose(expm([[1]]), A)
        assert_allclose(expm(matrix([[1]])), A)
        assert_allclose(expm(np.array([[1]])), A)
        assert_allclose(expm(csc_matrix([[1]])).A, A)
        B = expm(np.array([[1j]]))
        assert_allclose(expm(((1j,),)), B)
        assert_allclose(expm([[1j]]), B)
        assert_allclose(expm(matrix([[1j]])), B)
        assert_allclose(expm(csc_matrix([[1j]])).A, B)

    def test_bidiagonal_sparse(self):
        A = csc_matrix([
            [1, 3, 0],
            [0, 1, 5],
            [0, 0, 2]], dtype=float)
        e1 = math.exp(1)
        e2 = math.exp(2)
        expected = np.array([
            [e1, 3*e1, 15*(e2 - 2*e1)],
            [0, e1, 5*(e2 - e1)],
            [0, 0, e2]], dtype=float)
        observed = expm(A).toarray()
        assert_array_almost_equal(observed, expected)

    def test_padecases_dtype_float(self):
        for dtype in [np.float32, np.float64]:
            for scale in [1e-2, 1e-1, 5e-1, 1, 10]:
                A = scale * eye(3, dtype=dtype)
                observed = expm(A)
                expected = exp(scale) * eye(3, dtype=dtype)
                assert_array_almost_equal_nulp(observed, expected, nulp=100)

    def test_padecases_dtype_complex(self):
        for dtype in [np.complex64, np.complex128]:
            for scale in [1e-2, 1e-1, 5e-1, 1, 10]:
                A = scale * eye(3, dtype=dtype)
                observed = expm(A)
                expected = exp(scale) * eye(3, dtype=dtype)
                assert_array_almost_equal_nulp(observed, expected, nulp=100)

    def test_padecases_dtype_sparse_float(self):
        # float32 and complex64 lead to errors in spsolve/UMFpack
        dtype = np.float64
        for scale in [1e-2, 1e-1, 5e-1, 1, 10]:
            a = scale * speye(3, 3, dtype=dtype, format='csc')
            e = exp(scale) * eye(3, dtype=dtype)
            with suppress_warnings() as sup:
                sup.filter(SparseEfficiencyWarning,
                           "Changing the sparsity structure of a csc_matrix is expensive.")
                exact_onenorm = _expm(a, use_exact_onenorm=True).toarray()
                inexact_onenorm = _expm(a, use_exact_onenorm=False).toarray()
            assert_array_almost_equal_nulp(exact_onenorm, e, nulp=100)
            assert_array_almost_equal_nulp(inexact_onenorm, e, nulp=100)

    def test_padecases_dtype_sparse_complex(self):
        # float32 and complex64 lead to errors in spsolve/UMFpack
        dtype = np.complex128
        for scale in [1e-2, 1e-1, 5e-1, 1, 10]:
            a = scale * speye(3, 3, dtype=dtype, format='csc')
            e = exp(scale) * eye(3, dtype=dtype)
            with suppress_warnings() as sup:
                sup.filter(SparseEfficiencyWarning,
                           "Changing the sparsity structure of a csc_matrix is expensive.")
                assert_array_almost_equal_nulp(expm(a).toarray(), e, nulp=100)

    def test_logm_consistency(self):
        random.seed(1234)
        for dtype in [np.float64, np.complex128]:
            for n in range(1, 10):
                for scale in [1e-4, 1e-3, 1e-2, 1e-1, 1, 1e1, 1e2]:
                    # make logm(A) be of a given scale
                    A = (eye(n) + random.rand(n, n) * scale).astype(dtype)
                    if np.iscomplexobj(A):
                        A = A + 1j * random.rand(n, n) * scale
                    assert_array_almost_equal(expm(logm(A)), A)

    def test_integer_matrix(self):
        Q = np.array([
            [-3, 1, 1, 1],
            [1, -3, 1, 1],
            [1, 1, -3, 1],
            [1, 1, 1, -3]])
        assert_allclose(expm(Q), expm(1.0 * Q))

    def test_integer_matrix_2(self):
        # Check for integer overflows
        Q = np.array([[-500, 500, 0, 0],
                      [0, -550, 360, 190],
                      [0, 630, -630, 0],
                      [0, 0, 0, 0]], dtype=np.int16)
        assert_allclose(expm(Q), expm(1.0 * Q))

        Q = csc_matrix(Q)
        assert_allclose(expm(Q).A, expm(1.0 * Q).A)

    def test_triangularity_perturbation(self):
        # Experiment (1) of
        # Awad H. Al-Mohy and Nicholas J. Higham (2012)
        # Improved Inverse Scaling and Squaring Algorithms
        # for the Matrix Logarithm.
        A = np.array([
            [3.2346e-1, 3e4, 3e4, 3e4],
            [0, 3.0089e-1, 3e4, 3e4],
            [0, 0, 3.221e-1, 3e4],
            [0, 0, 0, 3.0744e-1]],
            dtype=float)
        A_logm = np.array([
            [-1.12867982029050462e+00, 9.61418377142025565e+04,
             -4.52485573953179264e+09, 2.92496941103871812e+14],
            [0.00000000000000000e+00, -1.20101052953082288e+00,
             9.63469687211303099e+04, -4.68104828911105442e+09],
            [0.00000000000000000e+00, 0.00000000000000000e+00,
             -1.13289322264498393e+00, 9.53249183094775653e+04],
            [0.00000000000000000e+00, 0.00000000000000000e+00,
             0.00000000000000000e+00, -1.17947533272554850e+00]],
            dtype=float)
        assert_allclose(expm(A_logm), A, rtol=1e-4)

        # Perturb the upper triangular matrix by tiny amounts,
        # so that it becomes technically not upper triangular.
        random.seed(1234)
        tiny = 1e-17
        A_logm_perturbed = A_logm.copy()
        A_logm_perturbed[1, 0] = tiny
        with suppress_warnings() as sup:
            sup.filter(RuntimeWarning, "Ill-conditioned.*")
            A_expm_logm_perturbed = expm(A_logm_perturbed)
        rtol = 1e-4
        atol = 100 * tiny
        assert_(not np.allclose(A_expm_logm_perturbed, A, rtol=rtol, atol=atol))

    def test_burkardt_1(self):
        # This matrix is diagonal.
        # The calculation of the matrix exponential is simple.
        #
        # This is the first of a series of matrix exponential tests
        # collected by John Burkardt from the following sources.
        #
        # Alan Laub,
        # Review of "Linear System Theory" by Joao Hespanha,
        # SIAM Review,
        # Volume 52, Number 4, December 2010, pages 779--781.
        #
        # Cleve Moler and Charles Van Loan,
        # Nineteen Dubious Ways to Compute the Exponential of a Matrix,
        # Twenty-Five Years Later,
        # SIAM Review,
        # Volume 45, Number 1, March 2003, pages 3--49.
        #
        # Cleve Moler,
        # Cleve's Corner: A Balancing Act for the Matrix Exponential,
        # 23 July 2012.
        #
        # Robert Ward,
        # Numerical computation of the matrix exponential
        # with accuracy estimate,
        # SIAM Journal on Numerical Analysis,
        # Volume 14, Number 4, September 1977, pages 600--610.
        exp1 = np.exp(1)
        exp2 = np.exp(2)
        A = np.array([
            [1, 0],
            [0, 2],
            ], dtype=float)
        desired = np.array([
            [exp1, 0],
            [0, exp2],
            ], dtype=float)
        actual = expm(A)
        assert_allclose(actual, desired)

    def test_burkardt_2(self):
        # This matrix is symmetric.
        # The calculation of the matrix exponential is straightforward.
        A = np.array([
            [1, 3],
            [3, 2],
            ], dtype=float)
        desired = np.array([
            [39.322809708033859, 46.166301438885753],
            [46.166301438885768, 54.711576854329110],
            ], dtype=float)
        actual = expm(A)
        assert_allclose(actual, desired)

    def test_burkardt_3(self):
        # This example is due to Laub.
        # This matrix is ill-suited for the Taylor series approach.
        # As powers of A are computed, the entries blow up too quickly.
        exp1 = np.exp(1)
        exp39 = np.exp(39)
        A = np.array([
            [0, 1],
            [-39, -40],
            ], dtype=float)
        desired = np.array([
            [
                39/(38*exp1) - 1/(38*exp39),
                -np.expm1(-38) / (38*exp1)],
            [
                39*np.expm1(-38) / (38*exp1),
                -1/(38*exp1) + 39/(38*exp39)],
            ], dtype=float)
        actual = expm(A)
        assert_allclose(actual, desired)

    def test_burkardt_4(self):
        # This example is due to Moler and Van Loan.
        # The example will cause problems for the series summation approach,
        # as well as for diagonal Pade approximations.
        A = np.array([
            [-49, 24],
            [-64, 31],
            ], dtype=float)
        U = np.array([[3, 1], [4, 2]], dtype=float)
        V = np.array([[1, -1/2], [-2, 3/2]], dtype=float)
        w = np.array([-17, -1], dtype=float)
        desired = np.dot(U * np.exp(w), V)
        actual = expm(A)
        assert_allclose(actual, desired)

    def test_burkardt_5(self):
        # This example is due to Moler and Van Loan.
        # This matrix is strictly upper triangular
        # All powers of A are zero beyond some (low) limit.
        # This example will cause problems for Pade approximations.
        A = np.array([
            [0, 6, 0, 0],
            [0, 0, 6, 0],
            [0, 0, 0, 6],
            [0, 0, 0, 0],
            ], dtype=float)
        desired = np.array([
            [1, 6, 18, 36],
            [0, 1, 6, 18],
            [0, 0, 1, 6],
            [0, 0, 0, 1],
            ], dtype=float)
        actual = expm(A)
        assert_allclose(actual, desired)

    def test_burkardt_6(self):
        # This example is due to Moler and Van Loan.
        # This matrix does not have a complete set of eigenvectors.
        # That means the eigenvector approach will fail.
        exp1 = np.exp(1)
        A = np.array([
            [1, 1],
            [0, 1],
            ], dtype=float)
        desired = np.array([
            [exp1, exp1],
            [0, exp1],
            ], dtype=float)
        actual = expm(A)
        assert_allclose(actual, desired)