main.py

import operator
import pygame
import random
pygame.init()


class DrawInformation:
    BLACK = 0, 0, 0
    WHITE = 255, 255, 255
    GREEN = 0, 255, 0
    RED = 255, 0, 0
    GREY = 128, 128, 128
    BACKGROUND_COLOR = WHITE
    GRADIENTS = [

        (128, 128, 128), (160, 160, 160), (192, 192, 192)
    ]
    FONT = pygame.font.SysFont('comicsans', 20)
    LARGEFONT = pygame.font.SysFont('comicsans', 30)
    SIDE_PAD = 100
    TOP_PAD = 120

    def __init__(self, width, height, lst):
        self.width = width
        self.height = height
        self.window = pygame.display.set_mode(
            (self.width, self.height), pygame.RESIZABLE)
        pygame.display.set_caption("Sorting Algorithm Visualizer")
        self.set_list(lst)

    def set_dimensions(self):
        self.width, self.height = self.window.get_size()
        self.set_list(self.lst)

    def set_list(self, lst):
        self.lst = lst
        self.max_val = max(lst)
        self.min_val = min(lst)
        self.block_width = round((self.width-self.SIDE_PAD)/len(lst))
        self.block_height = ((self.height-self.TOP_PAD) /
                             (self.max_val-self.min_val))
        self.start_x = self.SIDE_PAD//2


def draw(draw_info, algo_name, ascending):
    draw_info.window.fill(draw_info.BACKGROUND_COLOR)

    title = draw_info.FONT.render(
        f"{algo_name}-{'Ascending' if ascending else 'Descending'}", 1, draw_info.BLACK)
    draw_info.window.blit(
        title, (draw_info.width/2-title.get_width()/2, 5))
    controls = draw_info.FONT.render(
        "R-RESET | SPACE- Start Sorting | A- Ascending | D- Descending", 1, draw_info.BLACK)
    draw_info.window.blit(
        controls, (draw_info.width/2-controls.get_width()/2, 35))
    sorting = draw_info.FONT.render(
        "I-Insertion Sort|B-Bubble Sort|S-Selection Sort", 1, draw_info.BLACK)
    draw_info.window.blit(
        sorting, (draw_info.width/2-sorting.get_width()/2, 65))
    draw_list(draw_info)
    pygame.display.update()


def draw_list(draw_info, color_positions={}, clear_bg=False):
    lst = draw_info.lst

    if clear_bg:
        clear_rect = (draw_info.SIDE_PAD//2, draw_info.TOP_PAD,
                      draw_info.width - draw_info.SIDE_PAD,
                      draw_info.height-draw_info.TOP_PAD)
        pygame.draw.rect(draw_info.window,
                         draw_info.BACKGROUND_COLOR, clear_rect)
    for i, val in enumerate(lst):
        x = draw_info.start_x+i*draw_info.block_width
        y = draw_info.height-(val-draw_info.min_val)*draw_info.block_height
        color = draw_info.GRADIENTS[i % 3]
        if i in color_positions:
            color = color_positions[i]
        pygame.draw.rect(draw_info.window, color,
                         (x, y, draw_info.block_width, draw_info.height))
    if clear_bg:
        pygame.display.update()


def generate_starting_list(n, min_val, max_val):
    lst = []
    for _ in range(n):
        val = random.randint(min_val, max_val)
        lst.append(val)
    return lst


def bubble_sort(draw_info, ascending=True):
    lst = draw_info.lst
    for i in range(len(lst)-1):
        for j in range(len(lst)-1-i):
            num1 = lst[j]
            num2 = lst[j+1]
            if (num1 > num2 and ascending) or (num1 < num2 and not ascending):
                lst[j], lst[j+1] = lst[j+1], lst[j]
                draw_list(draw_info, {j: draw_info.BLACK,
                          j+1: draw_info.RED}, True)
                yield True
    return lst


def insertion_sort(draw_info, ascending=True):
    lst = draw_info.lst
    for i in range(1, len(lst)):
        curr = lst[i]
        while True:
            ascending_sort = i > 0 and lst[i-1] > curr and ascending
            descending_sort = i > 0 and lst[i-1] < curr and not ascending
            if not ascending_sort and not descending_sort:
                break
            lst[i] = lst[i-1]
            i = i-1
            lst[i] = curr
            draw_list(draw_info, {i-1: draw_info.BLACK,
                      i: draw_info.RED}, True)
            yield True
    return lst


def selection_sort(draw_info, ascending=True):
    lst = draw_info.lst
    op = operator.ge if ascending else operator.le
    for i in range(len(lst)):
        min = i
        for j in range(i+1, len(lst)):
            if op(lst[min], lst[j]):
                min = j
        draw_list(draw_info, {i: draw_info.RED,
                              min: draw_info.BLACK}, True)
        lst[i], lst[min] = lst[min], lst[i]

        yield True
    return lst


def main():
    run = True
    clock = pygame.time.Clock()

    n = 50
    min_val = 0
    max_val = 100
    lst = generate_starting_list(n, min_val, max_val)

    draw_info = DrawInformation(800, 600, lst)

    sorting = False
    ascending = True
    sorting_algorithm = bubble_sort
    sorting_algo_name = "Bubble Sort"
    sorting_algorithm_generator = None
    while run:
        clock.tick(60)

        if sorting:
            try:
                next(sorting_algorithm_generator)  #yield returns to here and is called again 
            except StopIteration:
                sorting = False
        else:
            draw(draw_info, sorting_algo_name, ascending)
        pygame.display.update()
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                run = False
            if event.type == pygame.VIDEORESIZE:
                draw_info.set_dimensions()
            if event.type != pygame.KEYDOWN:
                continue

            if event.key == pygame.K_r:
                lst = generate_starting_list(n, min_val, max_val)
                draw_info.set_list(lst)
                sorting = False
            elif event.key == pygame.K_SPACE and sorting == False:
                sorting = True
                sorting_algorithm_generator = sorting_algorithm(
                    draw_info, ascending)
            elif event.key == pygame.K_a and not sorting:
                ascending = True
            elif event.key == pygame.K_d and not sorting:
                ascending = False
            elif event.key == pygame.K_i and not sorting:
                sorting_algorithm = insertion_sort
                sorting_algo_name = "Insertion Sort"
            elif event.key == pygame.K_b and not sorting:
                sorting_algorithm = bubble_sort
                sorting_algo_name = "Bubble Sort"
            elif event.key == pygame.K_s and not sorting:
                sorting_algorithm = selection_sort
                sorting_algo_name = "Selection Sort"

    pygame.quit()


if __name__ == "__main__":
    main()