Flip game in pygame

flip_game.py

'''
Implementation of the "Flip" game, in which clicking on a tile flips itself and the four tiles touching it.
'''

import random

import pygame
pygame.init()
pygame.display.set_caption('flip game')

clock = pygame.time.Clock()

square_size = 100
board_wh = 5
screen_wh = board_wh * square_size

dimensions = (screen_wh, screen_wh)
screen = pygame.display.set_mode(dimensions)

colors = {
    'white' : (255, 255, 255),
    'black' : (0, 0, 0),
    'red' : (255, 0, 0),
    'green' : (0, 255, 0),
    'blue' : (0, 0, 255),
    'light' : (240, 217, 181),
    'dark' : (181, 136, 99)
}

def square_from_coords(x, y):
    c = x // square_size
    r = -(y % (square_size * board_wh) // square_size + 1 - board_wh)
    return r, c
    
def random_board():
    # for each square, there is a 50% chance it is black
    out = set()
    for row in range(board_wh):
        for col in range(board_wh):
            if random.random() < 0.5:
                out.add((row, col))
    if len(out) == 0: # just in case
        return random_board()
    else:
        return out

def vec_apply(point, d):
    dx, dy = d
    x, y = point
    return (x + dx, y + dy)

def affected_squares(square):
    dirs_5 = [(0,1), (0,-1), (-1,0), (1,0), (0,0)]
    applied = [vec_apply(dir_5, square) for dir_5 in dirs_5]
    return [(x,y) for (x,y) in applied if 0 <= x < board_wh and 0 <= y < board_wh]

def apply_move(board, square):
    affected = affected_squares(square)
    for a in affected:
        if a in board:
            board.remove(a)
        else:
            board.add(a)
    return board

def copy_board(board):
    return {square for square in board}


# BOARD is a set of square tuples which are black (and not white)
BOARD = random_board()

mouse_click_down = False
key_is_down = False

prev_moves = []
next_moves = []

running = True
while running:
    key_down_count = 0
    for event in pygame.event.get():
        current_square = square_from_coords(*pygame.mouse.get_pos())
        if event.type == pygame.KEYUP and event.key == pygame.K_ESCAPE:
            running = False
        if event.type == pygame.QUIT:
            running = False
        if event.type == pygame.MOUSEBUTTONDOWN:
            if not mouse_click_down:
                mouse_click_down = True
                mouse_down_square = current_square
        if event.type == pygame.MOUSEBUTTONUP:
            mouse_click_down = False
            if current_square == mouse_down_square:
                # square click confirmed
                prev_moves.append(copy_board(BOARD))
                # apply move
                BOARD = apply_move(BOARD, current_square)
                # kill redo stack
                next_moves = []
                # print('BOARD:', BOARD)
                # print('prev_moves:', prev_moves)
                # print('next_moves:', next_moves)

        if event.type == pygame.KEYDOWN:
            key_down_count += 1
            if key_is_down:
                continue
            if event.key == pygame.K_r:
                # reset (randomize board)
                BOARD = random_board()
                mouse_click_down = False
                prev_moves = []
                next_moves = []
            if event.key == pygame.K_LEFT:
                mouse_click_down = False
                try:
                    next_moves.append(copy_board(BOARD))
                    BOARD = prev_moves.pop()
                except IndexError:
                    # do nothing (ignore keypress)
                    pass
            if event.key == pygame.K_RIGHT:
                mouse_click_down = False
                if next_moves:
                    prev_moves.append(copy_board(BOARD))
                    BOARD = next_moves.pop()

    if key_down_count == 0:
        key_is_down = False
        

    screen.fill(colors['white'])

    # squares
    for r in range(board_wh):
        for c in range(board_wh):
            square_number = board_wh * r + c
            if (r, c) in BOARD:
                this_color = colors['black']
            else:
                this_color = colors['white']
            pygame.draw.rect(
                screen, 
                this_color,
                pygame.Rect(
                    square_size * c,
                    square_size * (board_wh - r - 1),
                    square_size,
                    square_size
                ) 
            )

    # check for checkmate
    if len(BOARD) == 0:
        pygame.draw.rect(
            screen,
            colors['green'],
            pygame.Rect(0, 0, *dimensions), 
            width=10
        )
    
    clock.tick(60)
    pygame.display.flip()