"""Test swap."""
import numpy as np

from toqito.perms import swap


def test_swap_matrix_4_by_4():
    """Tests swap operation on a 4-by-4 matrix."""
    test_mat = np.arange(1, 17).reshape(4, 4).transpose()

    # Default argument (systems [1, 2]):
    expected_res = np.array([[1, 9, 5, 13], [3, 11, 7, 15], [2, 10, 6, 14], [4, 12, 8, 16]])
    res = swap(test_mat)

    bool_mat = np.isclose(res, expected_res)
    np.testing.assert_equal(np.all(bool_mat), True)

    # Hard-coded argument (systems [1, 2]):
    res = swap(test_mat, [1, 2], [2, 2])
    bool_mat = np.isclose(res, expected_res)
    np.testing.assert_equal(np.all(bool_mat), True)

    # Hard-coded argument (systems [2, 1])--this should be identical to the prior two cases.
    res = swap(test_mat, [2, 1], [2, 2])
    bool_mat = np.isclose(res, expected_res)
    np.testing.assert_equal(np.all(bool_mat), True)


def test_swap_vector_1():
    """Tests swap operation on vector."""
    test_vec = np.array([1, 2, 3, 4])

    expected_res = np.array([1, 3, 2, 4])

    res = swap(test_vec)

    bool_mat = np.isclose(res, expected_res)
    np.testing.assert_equal(np.all(bool_mat), True)


def test_swap_int_dim():
    """Test swap operation when int is provided."""
    test_mat = np.array([[1, 5, 9, 13], [2, 6, 10, 14], [3, 7, 11, 15], [4, 8, 12, 16]])

    expected_res = np.array([[1, 9, 5, 13], [3, 11, 7, 15], [2, 10, 6, 14], [4, 12, 8, 16]])

    res = swap(test_mat, [1, 2], 2)

    bool_mat = np.isclose(res, expected_res)
    np.testing.assert_equal(np.all(bool_mat), True)


def test_swap_8_by_8():
    """Test swap on an 8-by-8 matrix."""
    test_mat = np.arange(1, 65).reshape(8, 8).transpose()

    # Swapping on systems "1" and "2":
    res_1_2 = swap(test_mat, [1, 2], [2, 2, 2])
    expected_res_1_2 = np.array(
        [
            [1, 9, 33, 41, 17, 25, 49, 57],
            [2, 10, 34, 42, 18, 26, 50, 58],
            [5, 13, 37, 45, 21, 29, 53, 61],
            [6, 14, 38, 46, 22, 30, 54, 62],
            [3, 11, 35, 43, 19, 27, 51, 59],
            [4, 12, 36, 44, 20, 28, 52, 60],
            [7, 15, 39, 47, 23, 31, 55, 63],
            [8, 16, 40, 48, 24, 32, 56, 64],
        ]
    )
    bool_mat = np.isclose(res_1_2, expected_res_1_2)
    np.testing.assert_equal(np.all(bool_mat), True)

    # Swapping on systems "2" and "1" (should be same as swapping on "1" and "2"):
    res_2_1 = swap(test_mat, [2, 1], [2, 2, 2])
    bool_mat = np.isclose(res_2_1, expected_res_1_2)
    np.testing.assert_equal(np.all(bool_mat), True)

    # Swapping on systems "1" and "3":
    res_1_3 = swap(test_mat, [1, 3], [2, 2, 2])
    expected_res_1_3 = np.array(
        [
            [1, 33, 17, 49, 9, 41, 25, 57],
            [5, 37, 21, 53, 13, 45, 29, 61],
            [3, 35, 19, 51, 11, 43, 27, 59],
            [7, 39, 23, 55, 15, 47, 31, 63],
            [2, 34, 18, 50, 10, 42, 26, 58],
            [6, 38, 22, 54, 14, 46, 30, 62],
            [4, 36, 20, 52, 12, 44, 28, 60],
            [8, 40, 24, 56, 16, 48, 32, 64],
        ]
    )
    bool_mat = np.isclose(res_1_3, expected_res_1_3)
    np.testing.assert_equal(np.all(bool_mat), True)

    # Swapping on systems "3" and "1" (should be same as swapping on "1" and "3"):
    res_3_1 = swap(test_mat, [3, 1], [2, 2, 2])
    bool_mat = np.isclose(res_3_1, expected_res_1_3)
    np.testing.assert_equal(np.all(bool_mat), True)

    # Swapping on systems "2" and "3":
    res_2_3 = swap(test_mat, [2, 3], [2, 2, 2])
    expected_res_2_3 = np.array(
        [
            [1, 17, 9, 25, 33, 49, 41, 57],
            [3, 19, 11, 27, 35, 51, 43, 59],
            [2, 18, 10, 26, 34, 50, 42, 58],
            [4, 20, 12, 28, 36, 52, 44, 60],
            [5, 21, 13, 29, 37, 53, 45, 61],
            [7, 23, 15, 31, 39, 55, 47, 63],
            [6, 22, 14, 30, 38, 54, 46, 62],
            [8, 24, 16, 32, 40, 56, 48, 64],
        ]
    )
    bool_mat = np.isclose(res_2_3, expected_res_2_3)
    np.testing.assert_equal(np.all(bool_mat), True)

    # Swapping on systems "3" and "2" (should be same as swapping on "2" and "3"):
    res_3_2 = swap(test_mat, [3, 2], [2, 2, 2])
    bool_mat = np.isclose(res_3_2, expected_res_2_3)
    np.testing.assert_equal(np.all(bool_mat), True)


def test_swap_16_by_16():
    """Test swap on an 16-by-16 matrix."""
    test_mat = np.arange(1, 257).reshape(16, 16).transpose()

    # Swapping on systems "1" and "2":
    res_1_2 = swap(test_mat, [1, 2], [2, 2, 2, 2])
    expected_res_1_2 = np.array(
        [
            [1, 17, 33, 49, 129, 145, 161, 177, 65, 81, 97, 113, 193, 209, 225, 241],
            [2, 18, 34, 50, 130, 146, 162, 178, 66, 82, 98, 114, 194, 210, 226, 242],
            [3, 19, 35, 51, 131, 147, 163, 179, 67, 83, 99, 115, 195, 211, 227, 243],
            [4, 20, 36, 52, 132, 148, 164, 180, 68, 84, 100, 116, 196, 212, 228, 244],
            [9, 25, 41, 57, 137, 153, 169, 185, 73, 89, 105, 121, 201, 217, 233, 249],
            [10, 26, 42, 58, 138, 154, 170, 186, 74, 90, 106, 122, 202, 218, 234, 250],
            [11, 27, 43, 59, 139, 155, 171, 187, 75, 91, 107, 123, 203, 219, 235, 251],
            [12, 28, 44, 60, 140, 156, 172, 188, 76, 92, 108, 124, 204, 220, 236, 252],
            [5, 21, 37, 53, 133, 149, 165, 181, 69, 85, 101, 117, 197, 213, 229, 245],
            [6, 22, 38, 54, 134, 150, 166, 182, 70, 86, 102, 118, 198, 214, 230, 246],
            [7, 23, 39, 55, 135, 151, 167, 183, 71, 87, 103, 119, 199, 215, 231, 247],
            [8, 24, 40, 56, 136, 152, 168, 184, 72, 88, 104, 120, 200, 216, 232, 248],
            [13, 29, 45, 61, 141, 157, 173, 189, 77, 93, 109, 125, 205, 221, 237, 253],
            [14, 30, 46, 62, 142, 158, 174, 190, 78, 94, 110, 126, 206, 222, 238, 254],
            [15, 31, 47, 63, 143, 159, 175, 191, 79, 95, 111, 127, 207, 223, 239, 255],
            [16, 32, 48, 64, 144, 160, 176, 192, 80, 96, 112, 128, 208, 224, 240, 256],
        ]
    )
    bool_mat = np.isclose(res_1_2, expected_res_1_2)
    np.testing.assert_equal(np.all(bool_mat), True)

    # Swapping on systems "2" and "1" (should be same as swapping on "1" and "2"):
    res_2_1 = swap(test_mat, [2, 1], [2, 2, 2, 2])
    bool_mat = np.isclose(res_2_1, expected_res_1_2)
    np.testing.assert_equal(np.all(bool_mat), True)

    # Swapping on systems "1" and "3":
    res_1_3 = swap(test_mat, [1, 3], [2, 2, 2, 2])
    expected_res_1_3 = np.array(
        [
            [1, 17, 129, 145, 65, 81, 193, 209, 33, 49, 161, 177, 97, 113, 225, 241],
            [2, 18, 130, 146, 66, 82, 194, 210, 34, 50, 162, 178, 98, 114, 226, 242],
            [9, 25, 137, 153, 73, 89, 201, 217, 41, 57, 169, 185, 105, 121, 233, 249],
            [10, 26, 138, 154, 74, 90, 202, 218, 42, 58, 170, 186, 106, 122, 234, 250],
            [5, 21, 133, 149, 69, 85, 197, 213, 37, 53, 165, 181, 101, 117, 229, 245],
            [6, 22, 134, 150, 70, 86, 198, 214, 38, 54, 166, 182, 102, 118, 230, 246],
            [13, 29, 141, 157, 77, 93, 205, 221, 45, 61, 173, 189, 109, 125, 237, 253],
            [14, 30, 142, 158, 78, 94, 206, 222, 46, 62, 174, 190, 110, 126, 238, 254],
            [3, 19, 131, 147, 67, 83, 195, 211, 35, 51, 163, 179, 99, 115, 227, 243],
            [4, 20, 132, 148, 68, 84, 196, 212, 36, 52, 164, 180, 100, 116, 228, 244],
            [11, 27, 139, 155, 75, 91, 203, 219, 43, 59, 171, 187, 107, 123, 235, 251],
            [12, 28, 140, 156, 76, 92, 204, 220, 44, 60, 172, 188, 108, 124, 236, 252],
            [7, 23, 135, 151, 71, 87, 199, 215, 39, 55, 167, 183, 103, 119, 231, 247],
            [8, 24, 136, 152, 72, 88, 200, 216, 40, 56, 168, 184, 104, 120, 232, 248],
            [15, 31, 143, 159, 79, 95, 207, 223, 47, 63, 175, 191, 111, 127, 239, 255],
            [16, 32, 144, 160, 80, 96, 208, 224, 48, 64, 176, 192, 112, 128, 240, 256],
        ]
    )
    bool_mat = np.isclose(res_1_3, expected_res_1_3)
    np.testing.assert_equal(np.all(bool_mat), True)

    # Swapping on systems "3" and "1" (should be same as swapping on "1" and "3"):
    res_3_1 = swap(test_mat, [3, 1], [2, 2, 2, 2])
    bool_mat = np.isclose(res_3_1, expected_res_1_3)
    np.testing.assert_equal(np.all(bool_mat), True)

    # Swapping on systems "1" and "4":
    res_1_4 = swap(test_mat, [1, 4], [2, 2, 2, 2])
    expected_res_1_4 = np.array(
        [
            [1, 129, 33, 161, 65, 193, 97, 225, 17, 145, 49, 177, 81, 209, 113, 241],
            [9, 137, 41, 169, 73, 201, 105, 233, 25, 153, 57, 185, 89, 217, 121, 249],
            [3, 131, 35, 163, 67, 195, 99, 227, 19, 147, 51, 179, 83, 211, 115, 243],
            [11, 139, 43, 171, 75, 203, 107, 235, 27, 155, 59, 187, 91, 219, 123, 251],
            [5, 133, 37, 165, 69, 197, 101, 229, 21, 149, 53, 181, 85, 213, 117, 245],
            [13, 141, 45, 173, 77, 205, 109, 237, 29, 157, 61, 189, 93, 221, 125, 253],
            [7, 135, 39, 167, 71, 199, 103, 231, 23, 151, 55, 183, 87, 215, 119, 247],
            [15, 143, 47, 175, 79, 207, 111, 239, 31, 159, 63, 191, 95, 223, 127, 255],
            [2, 130, 34, 162, 66, 194, 98, 226, 18, 146, 50, 178, 82, 210, 114, 242],
            [10, 138, 42, 170, 74, 202, 106, 234, 26, 154, 58, 186, 90, 218, 122, 250],
            [4, 132, 36, 164, 68, 196, 100, 228, 20, 148, 52, 180, 84, 212, 116, 244],
            [12, 140, 44, 172, 76, 204, 108, 236, 28, 156, 60, 188, 92, 220, 124, 252],
            [6, 134, 38, 166, 70, 198, 102, 230, 22, 150, 54, 182, 86, 214, 118, 246],
            [14, 142, 46, 174, 78, 206, 110, 238, 30, 158, 62, 190, 94, 222, 126, 254],
            [8, 136, 40, 168, 72, 200, 104, 232, 24, 152, 56, 184, 88, 216, 120, 248],
            [16, 144, 48, 176, 80, 208, 112, 240, 32, 160, 64, 192, 96, 224, 128, 256],
        ]
    )
    bool_mat = np.isclose(res_1_4, expected_res_1_4)
    np.testing.assert_equal(np.all(bool_mat), True)

    # Swapping on systems "4" and "1" (should be same as swapping on "1" and "4"):
    res_4_1 = swap(test_mat, [4, 1], [2, 2, 2, 2])
    bool_mat = np.isclose(res_4_1, expected_res_1_4)
    np.testing.assert_equal(np.all(bool_mat), True)

    # Swapping on systems "2" and "3":
    res_2_3 = swap(test_mat, [2, 3], [2, 2, 2, 2])
    expected_res_2_3 = np.array(
        [
            [1, 17, 65, 81, 33, 49, 97, 113, 129, 145, 193, 209, 161, 177, 225, 241],
            [2, 18, 66, 82, 34, 50, 98, 114, 130, 146, 194, 210, 162, 178, 226, 242],
            [5, 21, 69, 85, 37, 53, 101, 117, 133, 149, 197, 213, 165, 181, 229, 245],
            [6, 22, 70, 86, 38, 54, 102, 118, 134, 150, 198, 214, 166, 182, 230, 246],
            [3, 19, 67, 83, 35, 51, 99, 115, 131, 147, 195, 211, 163, 179, 227, 243],
            [4, 20, 68, 84, 36, 52, 100, 116, 132, 148, 196, 212, 164, 180, 228, 244],
            [7, 23, 71, 87, 39, 55, 103, 119, 135, 151, 199, 215, 167, 183, 231, 247],
            [8, 24, 72, 88, 40, 56, 104, 120, 136, 152, 200, 216, 168, 184, 232, 248],
            [9, 25, 73, 89, 41, 57, 105, 121, 137, 153, 201, 217, 169, 185, 233, 249],
            [10, 26, 74, 90, 42, 58, 106, 122, 138, 154, 202, 218, 170, 186, 234, 250],
            [13, 29, 77, 93, 45, 61, 109, 125, 141, 157, 205, 221, 173, 189, 237, 253],
            [14, 30, 78, 94, 46, 62, 110, 126, 142, 158, 206, 222, 174, 190, 238, 254],
            [11, 27, 75, 91, 43, 59, 107, 123, 139, 155, 203, 219, 171, 187, 235, 251],
            [12, 28, 76, 92, 44, 60, 108, 124, 140, 156, 204, 220, 172, 188, 236, 252],
            [15, 31, 79, 95, 47, 63, 111, 127, 143, 159, 207, 223, 175, 191, 239, 255],
            [16, 32, 80, 96, 48, 64, 112, 128, 144, 160, 208, 224, 176, 192, 240, 256],
        ]
    )
    bool_mat = np.isclose(res_2_3, expected_res_2_3)
    np.testing.assert_equal(np.all(bool_mat), True)

    # Swapping on systems "3" and "2" (should be same as swapping on "2" and "3"):
    res_3_2 = swap(test_mat, [3, 2], [2, 2, 2, 2])
    bool_mat = np.isclose(res_3_2, expected_res_2_3)
    np.testing.assert_equal(np.all(bool_mat), True)

    # Swapping on systems "2" and "4":
    res_2_4 = swap(test_mat, [2, 4], [2, 2, 2, 2])
    expected_res_2_4 = np.array(
        [
            [1, 65, 33, 97, 17, 81, 49, 113, 129, 193, 161, 225, 145, 209, 177, 241],
            [5, 69, 37, 101, 21, 85, 53, 117, 133, 197, 165, 229, 149, 213, 181, 245],
            [3, 67, 35, 99, 19, 83, 51, 115, 131, 195, 163, 227, 147, 211, 179, 243],
            [7, 71, 39, 103, 23, 87, 55, 119, 135, 199, 167, 231, 151, 215, 183, 247],
            [2, 66, 34, 98, 18, 82, 50, 114, 130, 194, 162, 226, 146, 210, 178, 242],
            [6, 70, 38, 102, 22, 86, 54, 118, 134, 198, 166, 230, 150, 214, 182, 246],
            [4, 68, 36, 100, 20, 84, 52, 116, 132, 196, 164, 228, 148, 212, 180, 244],
            [8, 72, 40, 104, 24, 88, 56, 120, 136, 200, 168, 232, 152, 216, 184, 248],
            [9, 73, 41, 105, 25, 89, 57, 121, 137, 201, 169, 233, 153, 217, 185, 249],
            [13, 77, 45, 109, 29, 93, 61, 125, 141, 205, 173, 237, 157, 221, 189, 253],
            [11, 75, 43, 107, 27, 91, 59, 123, 139, 203, 171, 235, 155, 219, 187, 251],
            [15, 79, 47, 111, 31, 95, 63, 127, 143, 207, 175, 239, 159, 223, 191, 255],
            [10, 74, 42, 106, 26, 90, 58, 122, 138, 202, 170, 234, 154, 218, 186, 250],
            [14, 78, 46, 110, 30, 94, 62, 126, 142, 206, 174, 238, 158, 222, 190, 254],
            [12, 76, 44, 108, 28, 92, 60, 124, 140, 204, 172, 236, 156, 220, 188, 252],
            [16, 80, 48, 112, 32, 96, 64, 128, 144, 208, 176, 240, 160, 224, 192, 256],
        ]
    )
    bool_mat = np.isclose(res_2_4, expected_res_2_4)
    np.testing.assert_equal(np.all(bool_mat), True)

    # Swapping on systems "4" and "2" (should be same as swapping on "2" and "4"):
    res_4_2 = swap(test_mat, [4, 2], [2, 2, 2, 2])
    bool_mat = np.isclose(res_4_2, expected_res_2_4)
    np.testing.assert_equal(np.all(bool_mat), True)

    # Swapping on systems "3" and "4":
    res_3_4 = swap(test_mat, [3, 4], [2, 2, 2, 2])
    expected_res_3_4 = np.array(
        [
            [1, 33, 17, 49, 65, 97, 81, 113, 129, 161, 145, 177, 193, 225, 209, 241],
            [3, 35, 19, 51, 67, 99, 83, 115, 131, 163, 147, 179, 195, 227, 211, 243],
            [2, 34, 18, 50, 66, 98, 82, 114, 130, 162, 146, 178, 194, 226, 210, 242],
            [4, 36, 20, 52, 68, 100, 84, 116, 132, 164, 148, 180, 196, 228, 212, 244],
            [5, 37, 21, 53, 69, 101, 85, 117, 133, 165, 149, 181, 197, 229, 213, 245],
            [7, 39, 23, 55, 71, 103, 87, 119, 135, 167, 151, 183, 199, 231, 215, 247],
            [6, 38, 22, 54, 70, 102, 86, 118, 134, 166, 150, 182, 198, 230, 214, 246],
            [8, 40, 24, 56, 72, 104, 88, 120, 136, 168, 152, 184, 200, 232, 216, 248],
            [9, 41, 25, 57, 73, 105, 89, 121, 137, 169, 153, 185, 201, 233, 217, 249],
            [11, 43, 27, 59, 75, 107, 91, 123, 139, 171, 155, 187, 203, 235, 219, 251],
            [10, 42, 26, 58, 74, 106, 90, 122, 138, 170, 154, 186, 202, 234, 218, 250],
            [12, 44, 28, 60, 76, 108, 92, 124, 140, 172, 156, 188, 204, 236, 220, 252],
            [13, 45, 29, 61, 77, 109, 93, 125, 141, 173, 157, 189, 205, 237, 221, 253],
            [15, 47, 31, 63, 79, 111, 95, 127, 143, 175, 159, 191, 207, 239, 223, 255],
            [14, 46, 30, 62, 78, 110, 94, 126, 142, 174, 158, 190, 206, 238, 222, 254],
            [16, 48, 32, 64, 80, 112, 96, 128, 144, 176, 160, 192, 208, 240, 224, 256],
        ]
    )
    bool_mat = np.isclose(res_3_4, expected_res_3_4)
    np.testing.assert_equal(np.all(bool_mat), True)

    # Swapping on systems "4" and "3" (should be same as swapping on "3" and "4"):
    res_4_3 = swap(test_mat, [4, 3], [2, 2, 2, 2])
    bool_mat = np.isclose(res_4_3, expected_res_3_4)
    np.testing.assert_equal(np.all(bool_mat), True)


def test_swap_invalid_dim():
    """Invalid dim parameters."""
    test_mat = np.array([[1, 5, 9, 13], [2, 6, 10, 14], [3, 7, 11, 15], [4, 8, 12, 16]])
    with np.testing.assert_raises(ValueError):
        swap(test_mat, [1, 2], 5)


def test_swap_invalid_sys():
    """Invalid sys parameters."""
    test_mat = np.array([[1, 5, 9, 13], [2, 6, 10, 14], [3, 7, 11, 15], [4, 8, 12, 16]])
    with np.testing.assert_raises(ValueError):
        swap(test_mat, [0])


def test_swap_invalid_sys_len():
    """Invalid sys parameters."""
    test_mat = np.array([[1, 5, 9, 13], [2, 6, 10, 14], [3, 7, 11, 15], [4, 8, 12, 16]])
    with np.testing.assert_raises(ValueError):
        swap(test_mat, [1], 2)


if __name__ == "__main__":
    np.testing.run_module_suite()
