import os
import pytest

import numpy as np
from numpy.testing import assert_allclose
from pytest import raises as assert_raises
from scipy.sparse.linalg._svdp import _svdp
from scipy.sparse import csr_array, csc_array


# dtype_flavour to tolerance
TOLS = {
    np.float32: 1e-4,
    np.float64: 1e-8,
    np.complex64: 1e-4,
    np.complex128: 1e-8,
}


def is_complex_type(dtype):
    return np.dtype(dtype).kind == "c"


_dtypes = []
for dtype_flavour in TOLS.keys():
    marks = []
    if is_complex_type(dtype_flavour):
        marks = [pytest.mark.slow]
    _dtypes.append(pytest.param(dtype_flavour, marks=marks,
                                id=dtype_flavour.__name__))
_dtypes = tuple(_dtypes)  # type: ignore[assignment]


# The test function here is adapted from the original Fortran PROPACK tests.
# It is not very robust to arbitrary seeding since partial reorthogonalization
# does not have a predictable upperbound on the number of iterations.
def check_svdp(n, m, constructor, dtype, k, irl_mode, which, f=0.8, rng=None):
    tol = TOLS[dtype]

    if rng is None:
        rng = np.random.default_rng(0)

    # Legacy clamp for the generator
    rng2 = np.random.default_rng(0)
    if is_complex_type(dtype):
        M = (- 5 + 10 * rng2.uniform(size=[n, m])
             - 5j + 10j * rng2.uniform(size=[n, m])).astype(dtype)
    else:
        M = (-5 + 10 * rng2.uniform(size=[n, m])).astype(dtype)
    M[M.real > 10 * f - 5] = 0
    Msp = constructor(M)

    u1, sigma1, vt1 = np.linalg.svd(M, full_matrices=False)
    u2, sigma2, vt2, _ = _svdp(Msp, k=k,which=which, irl_mode=irl_mode,
                               tol=tol, rng=rng)

    # check the which
    if which.upper() == 'SM':
        u1 = np.roll(u1, k, 1)
        vt1 = np.roll(vt1, k, 0)
        sigma1 = np.roll(sigma1, k)

    # check that singular values agree
    assert_allclose(sigma1[:k], sigma2, rtol=tol, atol=tol)

    # check that singular vectors are orthogonal
    assert_allclose(np.abs(u1.conj().T @ u2), np.eye(n, k), rtol=tol, atol=tol)
    assert_allclose(np.abs(vt1.conj() @ vt2.T), np.eye(n, k), rtol=tol, atol=tol)


@pytest.mark.parametrize('ctor', (np.array, csr_array, csc_array))
@pytest.mark.parametrize('dtype', [np.float32, np.float64,
                                   np.complex64, np.complex128])
@pytest.mark.parametrize('irl', (True, False))
@pytest.mark.parametrize('which', ('LM', 'SM'))
def test_svdp(ctor, dtype, irl, which):
    rng = np.random.default_rng(1757937293955503)
    n, m, k = 10, 20, 3
    if which == 'SM' and not irl:
        message = "`which`='SM' requires irl_mode=True"
        with assert_raises(ValueError, match=message):
            check_svdp(n, m, ctor, dtype, k, irl, which, rng=rng)
    else:
        check_svdp(n, m, ctor, dtype, k, irl, which, rng=rng)


@pytest.mark.xslow
@pytest.mark.parametrize('dtype', _dtypes)
@pytest.mark.parametrize('irl', (False, True))
def test_examples(dtype, irl):
    # Note: atol for complex64 bumped from 1e-4 to 1e-3 due to test failures
    # with BLIS, Netlib, and MKL+AVX512 - see
    # https://github.com/conda-forge/scipy-feedstock/pull/198#issuecomment-999180432
    atol = {
        np.float32: 1.3e-4,
        np.float64: 1e-9,
        np.complex64: 1e-3,
        np.complex128: 1e-9,
    }[dtype]

    path_prefix = os.path.dirname(__file__)
    # Test matrices from `illc1850.coord` and `mhd1280b.cua` distributed with
    # PROPACK 2.1: http://sun.stanford.edu/~rmunk/PROPACK/
    relative_path = "propack_test_data.npz"
    filename = os.path.join(path_prefix, relative_path)
    with np.load(filename, allow_pickle=True) as data:
        if is_complex_type(dtype):
            A = data['A_complex'].item().astype(dtype)
        else:
            A = data['A_real'].item().astype(dtype)

    k = 200
    u, s, vh, _ = _svdp(A, k, irl_mode=irl, rng=np.random.default_rng(0))

    # complex example matrix has many repeated singular values, so check only
    # beginning non-repeated singular vectors to avoid permutations
    sv_check = 27 if is_complex_type(dtype) else k
    u = u[:, :sv_check]
    vh = vh[:sv_check, :]
    s = s[:sv_check]

    # Check orthogonality of singular vectors
    assert_allclose(np.eye(u.shape[1]), u.conj().T @ u, atol=atol)
    assert_allclose(np.eye(vh.shape[0]), vh @ vh.conj().T, atol=atol)

    # Ensure the norm of the difference between the np.linalg.svd and
    # PROPACK reconstructed matrices is small
    u3, s3, vh3 = np.linalg.svd(A.todense())
    u3 = u3[:, :sv_check]
    s3 = s3[:sv_check]
    vh3 = vh3[:sv_check, :]
    A3 = u3 @ np.diag(s3) @ vh3
    recon = u @ np.diag(s) @ vh
    assert_allclose(np.linalg.norm(A3 - recon), 0, atol=atol)


@pytest.mark.parametrize('shifts', (None, -10, 0, 1, 10, 70))
@pytest.mark.parametrize('dtype', _dtypes[:2])
def test_shifts(shifts, dtype):
    rng = np.random.default_rng(0)
    n, k = 70, 10
    A = rng.random((n, n))
    if shifts is not None and ((shifts < 0) or (k > min(n-1-shifts, n))):
        with pytest.raises(ValueError):
            _svdp(A, k, shifts=shifts, kmax=5*k, irl_mode=True, rng=rng)
    else:
        _svdp(A, k, shifts=shifts, kmax=5*k, irl_mode=True, rng=rng)


@pytest.mark.slow
@pytest.mark.xfail()
def test_shifts_accuracy():
    rng = np.random.default_rng(0)
    n, k = 70, 10
    A = rng.random((n, n)).astype(np.float64)
    u1, s1, vt1, _ = _svdp(A, k, shifts=None, which='SM', irl_mode=True, rng=rng)
    u2, s2, vt2, _ = _svdp(A, k, shifts=32, which='SM', irl_mode=True, rng=rng)
    # shifts <= 32 doesn't agree with shifts > 32
    # Does agree when which='LM' instead of 'SM'
    assert_allclose(s1, s2)


@pytest.mark.parametrize('irl_mode', [False, True])
@pytest.mark.parametrize('dtype', (np.float32, np.float64))
def test_thin_hilbert(irl_mode, dtype):
    rng = np.random.default_rng(1757951587606893)
    m, n = 200, 4

    # Generate a Hilbert matrix of size m x n
    A = np.array([[1 / (i + j + 1) for j in range(n)] for i in range(m)], dtype=dtype)
    uu, ss, vv = np.linalg.svd(A, full_matrices=False)
    u, s, vt, _ = _svdp(A, k=4, which='LM', irl_mode=irl_mode, rng=rng)
    assert_allclose(s, ss, atol=TOLS[dtype])

    # Check orthogonality of singular vectors
    assert_allclose(np.eye(u.shape[1]), u.T @ u, atol=TOLS[dtype])
    assert_allclose(np.eye(vt.shape[0]), vt @ vt.T, atol=TOLS[dtype])

    # Check orthogonality against numpy svd results
    assert_allclose(np.abs(uu.T @ u), np.eye(n), atol=TOLS[dtype])
    assert_allclose(np.abs(vv @ vt.T), np.eye(n), atol=TOLS[dtype])


@pytest.mark.parametrize('dtype', (np.float32, np.float64, np.complex64, np.complex128))
def test_fat_random(dtype):
    rng = np.random.default_rng(1758046113948869)
    m, n = 3, 100

    A = rng.uniform(size=(m, n)).astype(dtype)
    if dtype in (np.complex64, np.complex128):
        A += dtype(1j) * rng.uniform(size=(m, n)).astype(dtype)

    uu, ss, vv = np.linalg.svd(A, full_matrices=False)
    u, s, vt, _ = _svdp(A, k=3, which='LM', irl_mode=True, rng=rng)
    assert_allclose(s, ss, atol=TOLS[dtype])

    # Check orthogonality of singular vectors
    assert_allclose(np.eye(u.shape[1]), u.conj().T @ u, atol=TOLS[dtype])
    assert_allclose(np.eye(vt.shape[0]), vt @ vt.conj().T, atol=TOLS[dtype])

    # Check orthogonality against numpy svd results
    assert_allclose(np.abs(uu.conj().T @ u), np.eye(m), atol=TOLS[dtype])
    assert_allclose(np.abs(vv @ vt.conj().T), np.eye(m), atol=TOLS[dtype])