from manim import *

class BubbleSort(Scene):
    def construct(self):
        self.camera.background_color = WHITE  # Set background color to white

        # Step 1: Create an array of numbers to be sorted
        array = [5, 1, 4, 3, 2, 8]
        box_size = 2.0  # Increase the size for each box to make them larger
        box_elements = []

        # Step 2: Create a larger boxed text for each number and position them
        for i, num in enumerate(array):
            num_text = Text(str(num), font_size=48, weight=BOLD, color=BLACK)  # Increase font size for larger text
            box = Square(side_length=box_size, color=BLACK)  # Increase the side length for uniform larger boxes
            box_group = VGroup(box, num_text)  # Group the box and number together
            num_text.move_to(box.get_center())  # Center the number within the box
            box_group.move_to(ORIGIN + RIGHT * box_size * (i - len(array) / 2))  # Center the whole array on the screen
            box_elements.append(box_group)

        # Display the array with boxes on screen
        self.add(*box_elements)

        # Step 3: Bubble Sort logic with animations and optimization
        n = len(array)
        for i in range(n):
            swapped = False  # Reset swapped flag at the start of each outer loop
            for j in range(n - i - 1):
                # Highlight the pair being compared
                self.play(box_elements[j].animate.set_color(ORANGE),
                          box_elements[j + 1].animate.set_color(ORANGE))

                # If the current item is larger, swap the elements
                if array[j] > array[j + 1]:
                    array[j], array[j + 1] = array[j + 1], array[j]
                    swapped = True  # Set swapped to True since a swap occurred

                    # Define the start and end positions for the swap
                    start_left = box_elements[j].get_center()
                    end_left = box_elements[j + 1].get_center()
                    start_right = box_elements[j + 1].get_center()
                    end_right = box_elements[j].get_center()

                    # Define control points for a Bezier curve for each box
                    control_left_up = start_left + UP * box_size
                    control_right_up = end_left + UP * box_size
                    control_left_down = start_right + DOWN * box_size
                    control_right_down = end_right + DOWN * box_size

                    # Move boxes along the Bezier curves
                    self.play(
                        MoveAlongPath(box_elements[j], CubicBezier(start_left, control_left_up, control_right_up, end_left)),
                        MoveAlongPath(box_elements[j + 1], CubicBezier(start_right, control_left_down, control_right_down, end_right))
                    )

                    # Update box_elements to reflect the swap
                    box_elements[j], box_elements[j + 1] = box_elements[j + 1], box_elements[j]

                # Reset colors after comparison
                self.play(box_elements[j].animate.set_color(BLACK),
                          box_elements[j + 1].animate.set_color(BLACK))

            # Highlight the last element of the current pass in red (indicating it's in the correct position)
            self.play(box_elements[n - i - 1].animate.set_color(BLUE))

            # If no swaps were made in this pass, break early as the array is already sorted
            if not swapped:
                break

        # Pause for a moment before ending the animation
        self.wait(2)