"""Test interleaved ordering functionality.

Tests the create_interleaved_order_delay and create_binary_tree_order functions from ordering_utils.py.
"""

import sys
from pathlib import Path

# Add parent directory to path
sys.path.insert(0, str(Path(__file__).parent.parent))

from ordering_utils import create_interleaved_order_delay, create_binary_tree_order


def test_basic_interleave():
    """Test basic interleaving with K=5 groups, DELAY=2.

    Input: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24]
    Groups: [0, 5, 10, 15, 20], [1, 6, 11, 16, 21], [2, 7, 12, 17, 22], [3, 8, 13, 18, 23], [4, 9, 14, 19, 24]
    """
    N = 25
    K = 5
    DELAY = 2
    PAD = -1

    order = create_interleaved_order_delay(N, K, DELAY, PAD)

    # fmt: off
    expected = [
        0, -1, -1, -1, -1,
        5, -1, -1, -1, -1,
        10, 1, -1, -1, -1,
        15, 6, -1, -1, -1,
        20, 11, 2, -1, -1,
        -1, 16, 7, -1, -1,
        -1, 21, 12, 3, -1,
        -1, -1, 17, 8, -1,
        -1, -1, 22, 13, 4,
        -1, -1, -1, 18, 9,
        -1, -1, -1, 23, 14,
        -1, -1, -1, -1, 19,
        -1, -1, -1, -1, 24,
    ]
    # fmt: on

    assert order == expected, f"Expected:\n{expected}\n\nGot:\n{order}"

    # Verify all original indices appear exactly once
    non_pad = [x for x in order if x != PAD]
    assert sorted(non_pad) == list(range(N)), "All original indices must appear exactly once"


def test_different_k():
    """Test with different K values."""
    N = 20

    for K in [2, 4, 5, 10]:
        order = create_interleaved_order_delay(N, K, DELAY=2, PAD=-1)
        non_pad = [x for x in order if x != -1]

        # All elements should be present
        assert len(non_pad) == N
        assert sorted(non_pad) == list(range(N))

        # Each group should maintain order
        for i in range(K):
            group = [x for x in order if x != -1 and x % K == i]
            expected = list(range(i, N, K))
            assert group == expected


def test_no_delay():
    """Test with DELAY=0 (no padding).

    Input: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
    Groups: [0, 4, 8], [1, 5, 9], [2, 6, 10], [3, 7, 11]
    With DELAY=0, groups are simply interleaved with no padding
    """
    N = 12
    K = 4
    DELAY = 0
    PAD = -1

    order = create_interleaved_order_delay(N, K, DELAY, PAD)

    # fmt: off
    expected = [
        0, 1, 2, 3,
        4, 5, 6, 7,
        8, 9, 10, 11,
    ]
    # fmt: on

    assert order == expected, f"Expected:\n{expected}\n\nGot:\n{order}"

    # With no delay, we should have exactly N elements (no padding)
    assert len(order) == N
    assert sorted(order) == list(range(N))


def test_edge_case_k_equals_n():
    """Test when K equals N (each element is its own group)."""
    N = 5
    K = 5
    DELAY = 1
    PAD = -1

    order = create_interleaved_order_delay(N, K, DELAY, PAD)

    # All elements should appear
    non_pad = [x for x in order if x != PAD]
    assert len(non_pad) == N
    assert sorted(non_pad) == list(range(N))


def test_custom_pad_value():
    """Test with custom PAD value."""
    N = 10
    K = 5
    DELAY = 1
    PAD = -999

    order = create_interleaved_order_delay(N, K, DELAY, PAD)

    # Check padding uses custom value
    pad_values = [x for x in order if x == PAD]
    assert all(x == PAD for x in pad_values)

    # Non-pad elements should be 0 to N-1
    non_pad = [x for x in order if x != PAD]
    assert sorted(non_pad) == list(range(N))


def test_binary_tree_order():
    """Test binary partitioning order for various N values."""
    assert create_binary_tree_order(0) == []
    assert create_binary_tree_order(1) == [0]
    assert create_binary_tree_order(3) == [1, 0, 2]
    assert create_binary_tree_order(5) == [2, 1, 4, 0, 3]
    assert create_binary_tree_order(7) == [3, 1, 5, 0, 2, 4, 6]
    assert create_binary_tree_order(8) == [4, 2, 6, 1, 3, 5, 7, 0]
    assert create_binary_tree_order(10) == [5, 2, 8, 1, 4, 7, 9, 0, 3, 6]
    assert create_binary_tree_order(15) == [7, 3, 11, 1, 5, 9, 13, 0, 2, 4, 6, 8, 10, 12, 14]
    assert create_binary_tree_order(16) == [8, 4, 12, 2, 6, 10, 14, 1, 3, 5, 7, 9, 11, 13, 15, 0]


if __name__ == "__main__":
    # Run tests manually if not using pytest
    test_basic_interleave()
    print("✓ test_basic_interleave passed")

    test_different_k()
    print("✓ test_different_k passed")

    test_no_delay()
    print("✓ test_no_delay passed")

    test_edge_case_k_equals_n()
    print("✓ test_edge_case_k_equals_n passed")

    test_custom_pad_value()
    print("✓ test_custom_pad_value passed")

    test_binary_tree_order()
    print("✓ test_binary_tree_order passed")

    print("\nAll tests passed! ✓")