import numpy
from numpy.linalg import norm
import pytest

from ._solver_cases import diagonal_eigen_case
from .utils import assert_close
from ..eigensolvers import power_iteration, rayleigh_quotient_iteration, signed_eigensolve


def test_rayleigh_quotient_iteration_with_linear_operator() -> None:
    case = diagonal_eigen_case()

    def dense_solver(
            shifted_operator,
            rhs: numpy.ndarray,
            ) -> numpy.ndarray:
        basis = numpy.eye(case.operator.shape[0], dtype=complex)
        columns = [shifted_operator.matvec(basis[:, ii]) for ii in range(case.operator.shape[0])]
        dense_matrix = numpy.column_stack(columns)
        return numpy.linalg.lstsq(dense_matrix, rhs, rcond=None)[0]

    eigval, eigvec = rayleigh_quotient_iteration(
        case.linear_operator,
        case.guess_default,
        iterations=8,
        solver=dense_solver,
        )

    residual = norm(case.operator @ eigvec - eigval * eigvec)
    assert abs(eigval - 3.0) < 1e-12
    assert residual < 1e-12


def test_signed_eigensolve_negative_returns_largest_negative_mode() -> None:
    case = diagonal_eigen_case()

    eigvals, eigvecs = signed_eigensolve(case.operator, how_many=1, negative=True)

    assert eigvals.shape == (1,)
    assert eigvecs.shape == (4, 1)
    assert abs(eigvals[0] + 2.0) < 1e-12
    assert abs(eigvecs[3, 0]) > 0.99


def test_rayleigh_quotient_iteration_uses_default_linear_operator_solver() -> None:
    case = diagonal_eigen_case()

    eigval, eigvec = rayleigh_quotient_iteration(
        case.linear_operator,
        case.guess_default,
        iterations=8,
    )

    residual = norm(case.operator @ eigvec - eigval * eigvec)
    assert abs(eigval - 3.0) < 1e-12
    assert residual < 1e-12


def test_signed_eigensolve_linear_operator_fallback_returns_dominant_positive_mode() -> None:
    case = diagonal_eigen_case()

    eigvals, eigvecs = signed_eigensolve(case.linear_operator, how_many=1)

    assert eigvals.shape == (1,)
    assert eigvecs.shape == (4, 1)
    assert_close(eigvals[0], 5.0, atol=1e-12, rtol=1e-12)
    assert abs(eigvecs[0, 0]) > 0.99


def test_power_iteration_finds_dominant_mode() -> None:
    case = diagonal_eigen_case()
    eigval, eigvec = power_iteration(case.operator, guess_vector=numpy.ones(4, dtype=complex), iterations=20)

    assert eigval == pytest.approx(5.0, rel=1e-6)
    assert abs(eigvec[0]) > abs(eigvec[1])


def test_rayleigh_quotient_iteration_refines_known_sparse_mode() -> None:
    case = diagonal_eigen_case()

    def solver(matrix, rhs: numpy.ndarray) -> numpy.ndarray:
        return numpy.linalg.lstsq(matrix.toarray(), rhs, rcond=None)[0]

    eigval, eigvec = rayleigh_quotient_iteration(
        case.operator,
        case.guess_sparse,
        iterations=8,
        solver=solver,
    )

    residual = numpy.linalg.norm(case.operator @ eigvec - eigval * eigvec)
    assert eigval == pytest.approx(3.0, rel=1e-6)
    assert residual < 1e-8


def test_signed_eigensolve_returns_largest_positive_modes() -> None:
    case = diagonal_eigen_case()
    eigvals, eigvecs = signed_eigensolve(case.operator, how_many=2)

    assert_close(eigvals, [3.0, 5.0], atol=1e-6)
    assert eigvecs.shape == (4, 2)