from modules.arrays.dutch_flag import (
    dutch_flag,
    dutch_flag_partition,
    dutch_flag_partition_2,
)


class TestDutchFlagV1:
    def test_dutch_flag__three_elements_sorted_ascending_pivot_middle(self):
        a, pivot_index = [0, 1, 2], 1
        b = dutch_flag(pivot_index, a)
        assert b == [0, 1, 2]

    def test_dutch_flag__three_elements_sorted_descending(self):
        a, pivot_index = [2, 1, 0], 1
        b = dutch_flag(pivot_index, a)
        assert b == [0, 1, 2]

    def test_dutch_flag__three_elements_sorted_ascending_pivot_start(self):
        a, pivot_index = [0, 1, 2], 0
        b = dutch_flag(pivot_index, a)
        assert b == [0, 2, 1]

    def test_dutch_flag__three_elements_sorted_ascending_pivot_end(self):
        a, pivot_index = [0, 1, 2], 2
        b = dutch_flag(pivot_index, a)
        assert b == [1, 0, 2]

    def test_dutch_flag__palindromic_array(self):
        a, pivot_index = [2, 0, 1, 0, 2], 2
        b = dutch_flag(pivot_index, a)
        assert b == [0, 0, 1, 2, 2]

    def test_dutch_flag__array_2(self):
        a, pivot_index = [2, 2, 0, 1, 1, 1], 2
        b = dutch_flag(pivot_index, a)
        assert b == [0, 2, 1, 2, 1, 1]


class TestDutchFlagV2:
    def test_dutch_flag__three_elements_sorted_ascending_pivot_middle(self):
        a, pivot_index = [0, 1, 2], 1
        b = dutch_flag_partition(pivot_index, a)
        assert b == [0, 1, 2]

    def test_dutch_flag__three_elements_sorted_descending(self):
        a, pivot_index = [2, 1, 0], 1
        b = dutch_flag_partition(pivot_index, a)
        assert b == [0, 1, 2]

    def test_dutch_flag__three_elements_sorted_ascending_pivot_start(self):
        a, pivot_index = [0, 1, 2], 0
        b = dutch_flag_partition(pivot_index, a)
        assert b == [0, 1, 2]

    def test_dutch_flag__three_elements_sorted_ascending_pivot_end(self):
        a, pivot_index = [0, 1, 2], 2
        b = dutch_flag_partition(pivot_index, a)
        assert b == [0, 1, 2]

    def test_dutch_flag__palindromic_array(self):
        a, pivot_index = [2, 0, 1, 0, 2], 2
        b = dutch_flag_partition(pivot_index, a)
        assert b == [0, 0, 1, 2, 2]

    def test_dutch_flag__array_2(self):
        a, pivot_index = [2, 2, 0, 1, 1, 1], 2
        b = dutch_flag_partition(pivot_index, a)
        assert b == [0, 2, 2, 1, 1, 1]


class TestDutchFlagV3:
    def test_dutch_flag__three_elements_sorted_ascending_pivot_middle(self):
        a, pivot_index = [0, 1, 2], 1
        b = dutch_flag_partition_2(pivot_index, a)
        assert b == [0, 1, 2]

    def test_dutch_flag__three_elements_sorted_descending(self):
        a, pivot_index = [2, 1, 0], 1
        b = dutch_flag_partition_2(pivot_index, a)
        assert b == [0, 1, 2]

    def test_dutch_flag__three_elements_sorted_ascending_pivot_start(self):
        a, pivot_index = [0, 1, 2], 0
        b = dutch_flag_partition_2(pivot_index, a)
        assert b == [0, 2, 1]

    def test_dutch_flag__three_elements_sorted_ascending_pivot_end(self):
        a, pivot_index = [0, 1, 2], 2
        b = dutch_flag_partition_2(pivot_index, a)
        assert b == [0, 1, 2]

    def test_dutch_flag__palindromic_array(self):
        a, pivot_index = [2, 0, 1, 0, 2], 2
        b = dutch_flag_partition_2(pivot_index, a)
        assert b == [0, 0, 1, 2, 2]

    def test_dutch_flag__array_2(self):
        a, pivot_index = [2, 2, 0, 1, 1, 1], 2
        b = dutch_flag_partition_2(pivot_index, a)
        assert b == [0, 2, 1, 1, 1, 2]
