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

import numpy as np
from numpy.linalg import norm as npnorm
from numpy.testing import assert_allclose, assert_equal
from pytest import raises as assert_raises

import scipy.sparse
from scipy.sparse.linalg import norm as spnorm


# https://github.com/scipy/scipy/issues/16031
# https://github.com/scipy/scipy/issues/21690
def test_sparray_norm():
    row = np.array([0, 0, 1, 1])
    col = np.array([0, 1, 2, 3])
    data = np.array([4, 5, 7, 9])
    test_arr = scipy.sparse.coo_array((data, (row, col)), shape=(2, 4))
    test_mat = scipy.sparse.coo_matrix((data, (row, col)), shape=(2, 4))
    for ord in (1, np.inf, None):
        for ax in [0, 1, None, (0, 1), (1, 0)]:
            for A in (test_arr, test_mat):
                expected = npnorm(A.toarray(), ord=ord, axis=ax)
                actual = spnorm(A, ord=ord, axis=ax)
                assert hasattr(actual, "dtype")
                assert_equal(actual, expected)
    # test 1d array and 1d-like (column) matrix
    test_arr_1d = scipy.sparse.coo_array((data, (col,)), shape=(4,))
    test_mat_col = scipy.sparse.coo_matrix((data, (col, [0, 0, 0, 0])), shape=(4, 1))
    for ord in (1, np.inf, None):
        for ax in [0, None]:
            for A in (test_arr_1d, test_mat_col):
                expected = npnorm(A.toarray(), ord=ord, axis=ax)
                assert_equal(spnorm(A, ord=ord, axis=ax), expected)


class TestNorm:
    def setup_method(self):
        a = np.arange(9) - 4
        b = a.reshape((3, 3))
        self.b = scipy.sparse.csr_array(b)

    def test_matrix_norm(self):

        # Frobenius norm is the default
        assert_allclose(spnorm(self.b), 7.745966692414834)
        assert_allclose(spnorm(self.b, 'fro'), 7.745966692414834)

        assert_allclose(spnorm(self.b, np.inf), 9)
        assert_allclose(spnorm(self.b, -np.inf), 2)
        assert_allclose(spnorm(self.b, 1), 7)
        assert_allclose(spnorm(self.b, -1), 6)
        assert_allclose(spnorm(self.b.astype(np.float64), 2), 7.348469228349534)

        # _multi_svd_norm is not implemented for sparse array
        assert_raises(NotImplementedError, spnorm, self.b, -2)

    def test_matrix_norm_axis(self):
        for m, axis in ((self.b, None), (self.b, (0, 1)), (self.b.T, (1, 0))):
            assert_allclose(spnorm(m, axis=axis), 7.745966692414834)
            assert_allclose(spnorm(m, 'fro', axis=axis), 7.745966692414834)
            assert_allclose(spnorm(m, np.inf, axis=axis), 9)
            assert_allclose(spnorm(m, -np.inf, axis=axis), 2)
            assert_allclose(spnorm(m, 1, axis=axis), 7)
            assert_allclose(spnorm(m, -1, axis=axis), 6)

    def test_vector_norm(self):
        v = [4.5825756949558398, 4.2426406871192848, 4.5825756949558398]
        for m, a in (self.b, 0), (self.b.T, 1):
            for axis in a, (a, ), a-2, (a-2, ):
                assert_allclose(spnorm(m, 1, axis=axis), [7, 6, 7])
                assert_allclose(spnorm(m, np.inf, axis=axis), [4, 3, 4])
                assert_allclose(spnorm(m, axis=axis), v)
                assert_allclose(spnorm(m, ord=2, axis=axis), v)
                assert_allclose(spnorm(m, ord=None, axis=axis), v)

    def test_norm_exceptions(self):
        m = self.b
        assert_raises(TypeError, spnorm, m, None, 1.5)
        assert_raises(TypeError, spnorm, m, None, [2])
        assert_raises(ValueError, spnorm, m, None, ())
        assert_raises(ValueError, spnorm, m, None, (0, 1, 2))
        assert_raises(ValueError, spnorm, m, None, (0, 0))
        assert_raises(ValueError, spnorm, m, None, (0, 2))
        assert_raises(ValueError, spnorm, m, None, (-3, 0))
        assert_raises(ValueError, spnorm, m, None, 2)
        assert_raises(ValueError, spnorm, m, None, -3)
        assert_raises(ValueError, spnorm, m, 'plate_of_shrimp', 0)
        assert_raises(ValueError, spnorm, m, 'plate_of_shrimp', (0, 1))


class TestVsNumpyNorm:
    _sparse_types = (
            scipy.sparse.bsr_array,
            scipy.sparse.coo_array,
            scipy.sparse.csc_array,
            scipy.sparse.csr_array,
            scipy.sparse.dia_array,
            scipy.sparse.dok_array,
            scipy.sparse.lil_array,
            )
    _test_matrices = (
            (np.arange(9) - 4).reshape((3, 3)),
            [
                [1, 2, 3],
                [-1, 1, 4]],
            [
                [1, 0, 3],
                [-1, 1, 4j]],
            )

    def test_sparse_matrix_norms(self):
        for sparse_type in self._sparse_types:
            for M in self._test_matrices:
                S = sparse_type(M)
                assert_allclose(spnorm(S), npnorm(M))
                assert_allclose(spnorm(S, 'fro'), npnorm(M, 'fro'))
                assert_allclose(spnorm(S, np.inf), npnorm(M, np.inf))
                assert_allclose(spnorm(S, -np.inf), npnorm(M, -np.inf))
                assert_allclose(spnorm(S, 1), npnorm(M, 1))
                assert_allclose(spnorm(S, -1), npnorm(M, -1))

    def test_sparse_matrix_norms_with_axis(self):
        for sparse_type in self._sparse_types:
            for M in self._test_matrices:
                S = sparse_type(M)
                for axis in None, (0, 1), (1, 0):
                    assert_allclose(spnorm(S, axis=axis), npnorm(M, axis=axis))
                    for ord in 'fro', np.inf, -np.inf, 1, -1:
                        assert_allclose(spnorm(S, ord, axis=axis),
                                        npnorm(M, ord, axis=axis))
                # Some numpy matrix norms are allergic to negative axes.
                for axis in (-2, -1), (-1, -2), (1, -2):
                    assert_allclose(spnorm(S, axis=axis), npnorm(M, axis=axis))
                    assert_allclose(spnorm(S, 'f', axis=axis),
                                    npnorm(M, 'f', axis=axis))
                    assert_allclose(spnorm(S, 'fro', axis=axis),
                                    npnorm(M, 'fro', axis=axis))

    def test_sparse_vector_norms(self):
        for sparse_type in self._sparse_types:
            for M in self._test_matrices:
                S = sparse_type(M)
                for axis in (0, 1, -1, -2, (0, ), (1, ), (-1, ), (-2, )):
                    assert_allclose(spnorm(S, axis=axis), npnorm(M, axis=axis))
                    for ord in None, 2, np.inf, -np.inf, 1, 0.5, 0.42:
                        assert_allclose(spnorm(S, ord, axis=axis),
                                        npnorm(M, ord, axis=axis))
