import pygame
pygame.init()

# Constants
WIDTH, HEIGHT = (800, 480)  
LIGHT_SIZE = 60
LIGHT_COLORS = [(255, 0, 0), (0, 255, 0)]  # Red for off, green for on
BACKGROUNDS = [(0,0,0),(40,40,40)] #background colors for unsolved/solved

# Initial light status, you can change the length of the vector
lights = [0,1,1,1,1,1]
n = len(lights)

screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Light Puzzle")

#Game loop
running = True
while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        if event.type == pygame.MOUSEBUTTONDOWN:
            x, y = pygame.mouse.get_pos()
            start_x = (WIDTH - (n * LIGHT_SIZE + (n - 1) * 10)) // 2
            if y >= (HEIGHT - LIGHT_SIZE) / 2 and y <= (HEIGHT + LIGHT_SIZE) // 2:
                index = (x - start_x) // (LIGHT_SIZE + 10)
                if 0 <= index < n:
                    for i in range(max(0, index - 1), min(n, index + 2)):
                        lights[i] = 1 - lights[i]
    #Draw
    screen.fill(BACKGROUNDS[all(light == 1 for light in lights)])
    for i, light in enumerate(lights):
        x = i * (LIGHT_SIZE + 10) + (WIDTH - (n * LIGHT_SIZE + (n - 1) * 10)) / 2
        y = (HEIGHT - LIGHT_SIZE) / 2
        pygame.draw.rect(screen, LIGHT_COLORS[light], (x, y, LIGHT_SIZE, LIGHT_SIZE))
    pygame.display.flip()
pygame.quit()