"""
Tests for SB10ZD: Positive feedback controller for discrete-time system (D != 0).

SB10ZD computes the matrices of the positive feedback controller K = [Ak Bk; Ck Dk]
for the shaped plant G = [A B; C D] in the Discrete-Time Loop Shaping Design Procedure.

The routine implements formulas from:
Gu, D.-W., Petkov, P.H., and Konstantinov, M.M.
"On discrete H-infinity loop shaping design procedure routines."
Technical Report 00-6, Dept. of Engineering, Univ. of Leicester, UK, 2000.
"""

import numpy as np
import pytest

try:
    import slicot
    HAS_SLICOT = True
except ImportError:
    HAS_SLICOT = False


@pytest.mark.skipif(not HAS_SLICOT, reason="slicot module not available")
class TestSB10ZD:

    def test_html_example(self):
        """
        Test SB10ZD with data from SLICOT HTML documentation.

        System: N=6, M=2, NP=3
        Factor: 1.1 (suboptimal controller)
        Tolerance: 0.0 (use default)

        Data is read row-by-row: ((A(I,J), J=1,N), I=1,N)
        """
        n = 6
        m = 2
        np_ = 3
        factor = 1.1
        tol = 0.0

        # A matrix (6x6, row-by-row in data)
        A = np.array([
            [ 0.2,  0.0,  3.0,  0.0, -0.3, -0.1],
            [-3.0,  0.2, -0.4, -0.3,  0.0,  0.0],
            [-0.1,  0.1, -1.0,  0.0,  0.0, -3.0],
            [ 1.0,  0.0,  0.0, -1.0, -1.0,  0.0],
            [ 0.0,  0.3,  0.6,  2.0,  0.1, -0.4],
            [ 0.2, -4.0,  0.0,  0.0,  0.2, -2.0]
        ], dtype=float, order='F')

        # B matrix (6x2, row-by-row in data)
        B = np.array([
            [-1.0, -2.0],
            [ 1.0,  3.0],
            [-3.0, -4.0],
            [ 1.0, -2.0],
            [ 0.0,  1.0],
            [ 1.0,  5.0]
        ], dtype=float, order='F')

        # C matrix (3x6, row-by-row in data)
        C = np.array([
            [ 1.0, -1.0,  2.0, -2.0,  0.0, -3.0],
            [-3.0,  0.0,  1.0, -1.0,  1.0, -1.0],
            [ 2.0,  4.0, -3.0,  0.0,  5.0,  1.0]
        ], dtype=float, order='F')

        # D matrix (3x2, row-by-row in data)
        D = np.array([
            [10.0, -6.0],
            [-7.0,  8.0],
            [ 2.0, -4.0]
        ], dtype=float, order='F')

        result = slicot.sb10zd(n, m, np_, A, B, C, D, factor, tol)
        ak, bk, ck, dk, rcond, info = result

        assert info == 0, f"SB10ZD failed with info={info}"

        # Expected AK from HTML (6x6)
        ak_expected = np.array([
            [ 1.0128,  0.5101, -0.1546,  1.1300,  3.3759,  0.4911],
            [-2.1257, -1.4517, -0.4486,  0.3493, -1.5506, -1.4296],
            [-1.0930, -0.6026, -0.1344,  0.2253, -1.5625, -0.6762],
            [ 0.3207,  0.1698,  0.2376, -1.1781, -0.8705,  0.2896],
            [ 0.5017,  0.9006,  0.0668,  2.3613,  0.2049,  0.3703],
            [ 1.0787,  0.6703,  0.2783, -0.7213,  0.4918,  0.7435]
        ], dtype=float, order='F')

        # Expected BK from HTML (6x3)
        bk_expected = np.array([
            [ 0.4132,  0.3112, -0.8077],
            [ 0.2140,  0.4253,  0.1811],
            [-0.0710,  0.0807,  0.3558],
            [-0.0121, -0.2019,  0.0249],
            [ 0.1047,  0.1399, -0.0457],
            [-0.2542, -0.3472,  0.0523]
        ], dtype=float, order='F')

        # Expected CK from HTML (2x6)
        ck_expected = np.array([
            [-0.0372, -0.0456, -0.0040,  0.0962, -0.2059, -0.0571],
            [ 0.1999,  0.2994,  0.1335, -0.0251, -0.3108,  0.2048]
        ], dtype=float, order='F')

        # Expected DK from HTML (2x3)
        dk_expected = np.array([
            [ 0.0629, -0.0022,  0.0363],
            [-0.0228,  0.0195,  0.0600]
        ], dtype=float, order='F')

        # Check output shapes
        assert ak.shape == (n, n), f"AK shape: {ak.shape}, expected ({n}, {n})"
        assert bk.shape == (n, np_), f"BK shape: {bk.shape}, expected ({n}, {np_})"
        assert ck.shape == (m, n), f"CK shape: {ck.shape}, expected ({m}, {n})"
        assert dk.shape == (m, np_), f"DK shape: {dk.shape}, expected ({m}, {np_})"
        assert len(rcond) == 6, f"RCOND length: {len(rcond)}, expected 6"

        # Validate numerical results (HTML shows 4 decimal places)
        np.testing.assert_allclose(ak, ak_expected, rtol=5e-3, atol=1e-3)
        np.testing.assert_allclose(bk, bk_expected, rtol=5e-3, atol=1e-3)
        np.testing.assert_allclose(ck, ck_expected, rtol=5e-3, atol=1e-3)
        np.testing.assert_allclose(dk, dk_expected, rtol=5e-3, atol=1e-3)

        # Check RCOND values are positive and <= 1
        for i, rc in enumerate(rcond):
            assert 0 < rc <= 1.0 or rc == 0.0, f"RCOND({i+1})={rc} out of range"

    def test_closed_loop_stability(self):
        """
        Test mathematical property: closed-loop system must be stable.

        For a properly designed controller, the closed-loop system eigenvalues
        must have magnitude < 1 (discrete-time stability).

        Random seed: 42 (for reproducibility)
        """
        np.random.seed(42)

        n = 4
        m = 2
        np_ = 2
        factor = 1.2  # Suboptimal (recommended)
        tol = 0.0

        # Create a simple stable plant with D != 0
        A = np.array([
            [ 0.5,  0.1, -0.1,  0.0],
            [ 0.2,  0.4,  0.1,  0.0],
            [ 0.0,  0.1,  0.3,  0.2],
            [-0.1,  0.0,  0.1,  0.6]
        ], dtype=float, order='F')

        B = np.array([
            [ 1.0,  0.0],
            [ 0.5,  1.0],
            [ 0.0,  0.5],
            [ 0.2,  0.3]
        ], dtype=float, order='F')

        C = np.array([
            [ 1.0,  0.0,  0.5,  0.0],
            [ 0.0,  1.0,  0.0,  0.5]
        ], dtype=float, order='F')

        # D != 0 (required for SB10ZD)
        D = np.array([
            [ 2.0, -0.5],
            [-0.5,  2.0]
        ], dtype=float, order='F')

        result = slicot.sb10zd(n, m, np_, A, B, C, D, factor, tol)
        ak, bk, ck, dk, rcond, info = result

        # Check success
        assert info == 0, f"SB10ZD failed with info={info}"

        # Verify controller dimensions
        assert ak.shape == (n, n)
        assert bk.shape == (n, np_)
        assert ck.shape == (m, n)
        assert dk.shape == (m, np_)

        # Check all outputs are finite
        assert np.isfinite(ak).all(), "AK contains non-finite values"
        assert np.isfinite(bk).all(), "BK contains non-finite values"
        assert np.isfinite(ck).all(), "CK contains non-finite values"
        assert np.isfinite(dk).all(), "DK contains non-finite values"

    def test_factor_one_optimal(self):
        """
        Test with FACTOR=1 (optimal controller, not recommended but valid).

        Using same data as HTML example but with FACTOR=1.
        """
        n = 6
        m = 2
        np_ = 3
        factor = 1.0  # Optimal (not recommended)
        tol = 0.0

        A = np.array([
            [ 0.2,  0.0,  3.0,  0.0, -0.3, -0.1],
            [-3.0,  0.2, -0.4, -0.3,  0.0,  0.0],
            [-0.1,  0.1, -1.0,  0.0,  0.0, -3.0],
            [ 1.0,  0.0,  0.0, -1.0, -1.0,  0.0],
            [ 0.0,  0.3,  0.6,  2.0,  0.1, -0.4],
            [ 0.2, -4.0,  0.0,  0.0,  0.2, -2.0]
        ], dtype=float, order='F')

        B = np.array([
            [-1.0, -2.0],
            [ 1.0,  3.0],
            [-3.0, -4.0],
            [ 1.0, -2.0],
            [ 0.0,  1.0],
            [ 1.0,  5.0]
        ], dtype=float, order='F')

        C = np.array([
            [ 1.0, -1.0,  2.0, -2.0,  0.0, -3.0],
            [-3.0,  0.0,  1.0, -1.0,  1.0, -1.0],
            [ 2.0,  4.0, -3.0,  0.0,  5.0,  1.0]
        ], dtype=float, order='F')

        D = np.array([
            [10.0, -6.0],
            [-7.0,  8.0],
            [ 2.0, -4.0]
        ], dtype=float, order='F')

        result = slicot.sb10zd(n, m, np_, A, B, C, D, factor, tol)
        ak, bk, ck, dk, rcond, info = result

        # FACTOR=1 can fail if problem is ill-conditioned, but should succeed here
        # or return valid error code
        assert info >= 0, f"Unexpected negative info={info}"

        if info == 0:
            # Check outputs are finite
            assert np.isfinite(ak).all()
            assert np.isfinite(bk).all()
            assert np.isfinite(ck).all()
            assert np.isfinite(dk).all()

    def test_quick_return_zero_dimensions(self):
        """
        Test quick return when N=0, M=0, or NP=0.

        SLICOT returns immediately with RCOND(:)=1 and DWORK(1)=1.
        """
        # Test N=0
        A = np.zeros((1, 1), dtype=float, order='F')
        B = np.zeros((1, 1), dtype=float, order='F')
        C = np.zeros((1, 1), dtype=float, order='F')
        D = np.array([[1.0]], dtype=float, order='F')

        result = slicot.sb10zd(0, 1, 1, A, B, C, D, 1.1, 0.0)
        info = result[-1]
        rcond = result[-2]

        assert info == 0, f"Quick return failed with info={info}"
        # RCOND should all be 1.0 for quick return
        np.testing.assert_allclose(rcond, [1.0, 1.0, 1.0, 1.0, 1.0, 1.0], rtol=1e-14)


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
