"""
Tests for MB03JZ - Reorder eigenvalues of complex skew-Hamiltonian/Hamiltonian pencil.

Moves eigenvalues with strictly negative real parts to the leading subpencil
while keeping triangular form.
"""

import numpy as np
import pytest

from slicot import mb03jz


class TestMB03JZBasic:
    """Basic functionality tests for MB03JZ."""

    def test_basic_reordering(self):
        """
        Test basic functionality with a simple pencil.

        The full N×N pencil has structure that determines the eigenvalue count.
        We verify the routine maintains triangular structure and unitary Q.

        Random seed: 42 (for reproducibility)
        """
        np.random.seed(42)
        m = 2  # n = 2*m = 4
        n = 2 * m

        # Upper triangular A and B
        a = np.array([[1.0 + 0.0j, 0.3 - 0.2j],
                      [0.0, 2.0 + 0.0j]], dtype=complex, order='F')
        b = np.array([[1.0 + 0.0j, 0.2 + 0.1j],
                      [0.0, 3.0 + 0.0j]], dtype=complex, order='F')

        # D is skew-Hermitian (D = -D^H)
        d = np.array([[0.5j, 0.1 + 0.2j],
                      [0.0, -0.3j]], dtype=complex, order='F')

        # F is Hermitian (F = F^H)
        f = np.array([[1.0, 0.2 - 0.1j],
                      [0.0, 0.8]], dtype=complex, order='F')

        a_out, d_out, b_out, f_out, q_out, neig, info = mb03jz(
            'I', n, a.copy(order='F'), d.copy(order='F'),
            b.copy(order='F'), f.copy(order='F')
        )

        assert info == 0
        assert 0 <= neig <= m  # neig in valid range
        # A and B remain upper triangular
        np.testing.assert_allclose(np.tril(a_out, -1), 0.0, atol=1e-14)
        np.testing.assert_allclose(np.tril(b_out, -1), 0.0, atol=1e-14)
        # Q should be unitary
        np.testing.assert_allclose(q_out @ q_out.conj().T, np.eye(n, dtype=complex),
                                   rtol=1e-14, atol=1e-14)

    def test_mixed_eigenvalue_signs(self):
        """
        Test with a pencil having mixed eigenvalue signs in subpencil.

        Random seed: 123 (for reproducibility)
        """
        np.random.seed(123)
        m = 2
        n = 2 * m

        # First eigenvalue positive in subpencil, second negative
        a = np.array([[1.0 + 0.0j, 0.5 - 0.3j],
                      [0.0, 1.0 + 0.0j]], dtype=complex, order='F')
        b = np.array([[1.0 + 0.0j, 0.2 + 0.1j],
                      [0.0, -1.0 + 0.0j]], dtype=complex, order='F')

        d = np.array([[0.2j, 0.1 + 0.05j],
                      [0.0, -0.1j]], dtype=complex, order='F')

        f = np.array([[0.5, 0.1 - 0.05j],
                      [0.0, 0.3]], dtype=complex, order='F')

        a_out, d_out, b_out, f_out, q_out, neig, info = mb03jz(
            'I', n, a.copy(order='F'), d.copy(order='F'),
            b.copy(order='F'), f.copy(order='F')
        )

        assert info == 0
        assert 0 <= neig <= m  # neig in valid range
        # A and B remain upper triangular
        np.testing.assert_allclose(np.tril(a_out, -1), 0.0, atol=1e-14)
        np.testing.assert_allclose(np.tril(b_out, -1), 0.0, atol=1e-14)
        # Q should be unitary
        np.testing.assert_allclose(q_out @ q_out.conj().T, np.eye(n, dtype=complex),
                                   rtol=1e-14, atol=1e-14)

    def test_all_negative_eigenvalues(self):
        """
        Test with all eigenvalues having negative real parts.

        Random seed: 456 (for reproducibility)
        """
        np.random.seed(456)
        m = 2
        n = 2 * m

        # Create pencil where both eigenvalues have negative real parts
        a = np.array([[1.0 + 1.0j, 0.3 - 0.2j],
                      [0.0, 0.5 + 0.5j]], dtype=complex, order='F')
        b = np.array([[-1.0 + 0.5j, 0.2 + 0.1j],
                      [0.0, -0.8 + 0.3j]], dtype=complex, order='F')

        d = np.array([[0.3j, 0.15 + 0.1j],
                      [0.0, -0.2j]], dtype=complex, order='F')

        f = np.array([[0.6, 0.2 - 0.1j],
                      [0.0, 0.4]], dtype=complex, order='F')

        a_out, d_out, b_out, f_out, q_out, neig, info = mb03jz(
            'I', n, a.copy(order='F'), d.copy(order='F'),
            b.copy(order='F'), f.copy(order='F')
        )

        assert info == 0
        assert neig == 2  # Both eigenvalues have negative real parts
        # A and B remain upper triangular
        np.testing.assert_allclose(np.tril(a_out, -1), 0.0, atol=1e-14)
        np.testing.assert_allclose(np.tril(b_out, -1), 0.0, atol=1e-14)
        # Q should be unitary
        np.testing.assert_allclose(q_out @ q_out.conj().T, np.eye(n, dtype=complex),
                                   rtol=1e-14, atol=1e-14)


class TestMB03JZMathematicalProperties:
    """Mathematical property validation tests."""

    def test_unitary_transformation_preserves_structure(self):
        """
        Validate that Q is unitary: Q @ Q^H = I.

        Random seed: 789 (for reproducibility)
        """
        np.random.seed(789)
        m = 3
        n = 2 * m

        # Create random upper triangular matrices
        a = np.triu(np.random.randn(m, m) + 1j * np.random.randn(m, m)).astype(complex, order='F')
        b = np.triu(np.random.randn(m, m) + 1j * np.random.randn(m, m)).astype(complex, order='F')

        # D skew-Hermitian: upper triangular with purely imaginary diagonal
        d = np.triu(np.random.randn(m, m) + 1j * np.random.randn(m, m)).astype(complex, order='F')
        np.fill_diagonal(d, 1j * np.random.randn(m))

        # F Hermitian: upper triangular with real diagonal
        f = np.triu(np.random.randn(m, m) + 1j * np.random.randn(m, m)).astype(complex, order='F')
        np.fill_diagonal(f, np.random.randn(m))

        a_out, d_out, b_out, f_out, q_out, neig, info = mb03jz(
            'I', n, a.copy(order='F'), d.copy(order='F'),
            b.copy(order='F'), f.copy(order='F')
        )

        assert info == 0
        # Q should be unitary
        np.testing.assert_allclose(q_out @ q_out.conj().T, np.eye(n, dtype=complex),
                                   rtol=1e-13, atol=1e-14)
        np.testing.assert_allclose(q_out.conj().T @ q_out, np.eye(n, dtype=complex),
                                   rtol=1e-13, atol=1e-14)

    def test_output_triangular_structure(self):
        """
        Validate that output matrices maintain triangular structure.

        Random seed: 321 (for reproducibility)
        """
        np.random.seed(321)
        m = 4
        n = 2 * m

        a = np.triu(np.random.randn(m, m) + 1j * np.random.randn(m, m)).astype(complex, order='F')
        b = np.triu(np.random.randn(m, m) + 1j * np.random.randn(m, m)).astype(complex, order='F')
        d = np.triu(np.random.randn(m, m) + 1j * np.random.randn(m, m)).astype(complex, order='F')
        f = np.triu(np.random.randn(m, m) + 1j * np.random.randn(m, m)).astype(complex, order='F')

        a_out, d_out, b_out, f_out, q_out, neig, info = mb03jz(
            'I', n, a.copy(order='F'), d.copy(order='F'),
            b.copy(order='F'), f.copy(order='F')
        )

        assert info == 0
        # A and B remain upper triangular
        np.testing.assert_allclose(np.tril(a_out, -1), 0.0, atol=1e-14)
        np.testing.assert_allclose(np.tril(b_out, -1), 0.0, atol=1e-14)

    def test_eigenvalue_count_consistency(self):
        """
        Validate that neig correctly counts eigenvalues with negative real part.

        Random seed: 654 (for reproducibility)
        """
        np.random.seed(654)
        m = 3
        n = 2 * m

        # Create diagonal matrices to easily verify eigenvalues
        # Eigenvalue sign determined by Re(A[k,k] * conj(B[k,k]))
        a_diag = np.array([1.0 + 0.5j, 1.0 + 0.5j, 1.0 + 0.5j], dtype=complex)
        # First two have positive eigenvalue, third has negative
        b_diag = np.array([1.0 + 0.3j, 1.0 - 0.3j, -1.0 + 0.2j], dtype=complex)

        a = np.diag(a_diag).astype(complex, order='F')
        b = np.diag(b_diag).astype(complex, order='F')
        d = np.zeros((m, m), dtype=complex, order='F')
        np.fill_diagonal(d, 1j * np.random.randn(m))
        f = np.zeros((m, m), dtype=complex, order='F')
        np.fill_diagonal(f, np.random.randn(m))

        a_out, d_out, b_out, f_out, q_out, neig, info = mb03jz(
            'I', n, a.copy(order='F'), d.copy(order='F'),
            b.copy(order='F'), f.copy(order='F')
        )

        assert info == 0
        # Count expected negative eigenvalues based on Re(a[k]*conj(b[k]))
        expected_neg = sum(1 for ak, bk in zip(a_diag, b_diag)
                          if (ak * bk.conjugate()).real < 0)
        assert neig == expected_neg


class TestMB03JZEdgeCases:
    """Edge case tests."""

    def test_zero_dimension(self):
        """Test with n=0 (quick return)."""
        n = 0
        m = n // 2
        a = np.empty((m, m), dtype=complex, order='F')
        d = np.empty((m, m), dtype=complex, order='F')
        b = np.empty((m, m), dtype=complex, order='F')
        f = np.empty((m, m), dtype=complex, order='F')

        a_out, d_out, b_out, f_out, q_out, neig, info = mb03jz(
            'N', n, a, d, b, f
        )

        assert info == 0
        assert neig == 0

    def test_n_equals_2(self):
        """
        Test minimal case with n=2 (m=1, single eigenvalue).

        Random seed: 111 (for reproducibility)
        """
        np.random.seed(111)
        m = 1
        n = 2

        a = np.array([[1.5 + 0.5j]], dtype=complex, order='F')
        b = np.array([[-1.0 + 0.2j]], dtype=complex, order='F')
        d = np.array([[0.3j]], dtype=complex, order='F')
        f = np.array([[0.5]], dtype=complex, order='F')

        a_out, d_out, b_out, f_out, q_out, neig, info = mb03jz(
            'I', n, a.copy(order='F'), d.copy(order='F'),
            b.copy(order='F'), f.copy(order='F')
        )

        assert info == 0
        # Check sign of eigenvalue: Re(a[0,0] * conj(b[0,0]))
        eig_sign = (a[0, 0] * b[0, 0].conjugate()).real
        expected_neig = 1 if eig_sign < 0 else 0
        assert neig == expected_neig


class TestMB03JZCompqModes:
    """Tests for different COMPQ modes."""

    def test_compq_n_no_q_computed(self):
        """Test COMPQ='N' - Q not computed."""
        m = 2
        n = 2 * m

        a = np.array([[1.0 + 0.5j, 0.3 - 0.2j],
                      [0.0, 2.0 - 0.1j]], dtype=complex, order='F')
        b = np.array([[1.0 + 0.1j, 0.2 + 0.1j],
                      [0.0, 1.5 + 0.2j]], dtype=complex, order='F')
        d = np.array([[0.5j, 0.1 + 0.2j],
                      [0.0, -0.3j]], dtype=complex, order='F')
        f = np.array([[1.0, 0.2 - 0.1j],
                      [0.0, 0.8]], dtype=complex, order='F')

        a_out, d_out, b_out, f_out, q_out, neig, info = mb03jz(
            'N', n, a.copy(order='F'), d.copy(order='F'),
            b.copy(order='F'), f.copy(order='F')
        )

        assert info == 0

    def test_compq_u_update_q(self):
        """
        Test COMPQ='U' - update existing Q matrix.

        Random seed: 222 (for reproducibility)
        """
        np.random.seed(222)
        m = 2
        n = 2 * m

        a = np.array([[1.0 + 1.0j, 0.5 - 0.3j],
                      [0.0, 1.0 + 1.0j]], dtype=complex, order='F')
        b = np.array([[1.0 + 0.5j, 0.2 + 0.1j],
                      [0.0, -1.0 + 0.5j]], dtype=complex, order='F')
        d = np.array([[0.2j, 0.1 + 0.05j],
                      [0.0, -0.1j]], dtype=complex, order='F')
        f = np.array([[0.5, 0.1 - 0.05j],
                      [0.0, 0.3]], dtype=complex, order='F')

        # Start with identity matrix for Q0
        q0 = np.eye(n, dtype=complex, order='F')

        a_out, d_out, b_out, f_out, q_out, neig, info = mb03jz(
            'U', n, a.copy(order='F'), d.copy(order='F'),
            b.copy(order='F'), f.copy(order='F'), q=q0.copy(order='F')
        )

        assert info == 0
        # Result should be unitary since Q0 was identity
        np.testing.assert_allclose(q_out @ q_out.conj().T, np.eye(n, dtype=complex),
                                   rtol=1e-13, atol=1e-14)


class TestMB03JZErrorHandling:
    """Error handling tests."""

    def test_invalid_compq(self):
        """Test invalid COMPQ parameter."""
        m = 2
        n = 2 * m
        a = np.zeros((m, m), dtype=complex, order='F')
        d = np.zeros((m, m), dtype=complex, order='F')
        b = np.zeros((m, m), dtype=complex, order='F')
        f = np.zeros((m, m), dtype=complex, order='F')

        a_out, d_out, b_out, f_out, q_out, neig, info = mb03jz(
            'X', n, a, d, b, f
        )

        assert info == -1

    def test_invalid_n_negative(self):
        """Test negative N."""
        m = 2
        a = np.zeros((m, m), dtype=complex, order='F')
        d = np.zeros((m, m), dtype=complex, order='F')
        b = np.zeros((m, m), dtype=complex, order='F')
        f = np.zeros((m, m), dtype=complex, order='F')

        a_out, d_out, b_out, f_out, q_out, neig, info = mb03jz(
            'N', -2, a, d, b, f
        )

        assert info == -2

    def test_invalid_n_odd(self):
        """Test odd N (must be even)."""
        m = 2
        a = np.zeros((m, m), dtype=complex, order='F')
        d = np.zeros((m, m), dtype=complex, order='F')
        b = np.zeros((m, m), dtype=complex, order='F')
        f = np.zeros((m, m), dtype=complex, order='F')

        a_out, d_out, b_out, f_out, q_out, neig, info = mb03jz(
            'N', 3, a, d, b, f
        )

        assert info == -2
