Create a lightning in pygame

Question

I'm trying to create a lightning in my app. I did this so far

import pygame
import random as rd
import copy

pygame.init()


color = 250,250,250
starting_point =0,0
ending_point = 400,400

a = [[400,0]]

depth = rd.randrange(50,80)

for i in range(depth):
    rand1 = rd.randrange(-20,20)
    rand2 = rd.randrange(2,15)
    a.append([a[i][0]+rand1,a[i][1]+rand2])

screen = pygame.display.set_mode((800,600))

running = True

while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False

    screen.fill((0,0,0))

    for j in range(0,depth):
        pygame.draw.aaline(screen, (200-2*j,200-2*j,200-2*j), [ a[j][0] , a[j][1] ], [ a[j+1][0] , a[j+1][1] ],1)

    pygame.draw.line(screen, color, a[0], a[1], 1)
    pygame.display.update()

This code creates a randomly generated line of lightning like this:

I want to do something like this:

I want more lines to go from the initially generated line

All I need is the initial idea, should I use a class or another lists ?

Answer 1

You could put code in function and use recursion to start next line in random moment when you generate previous line. But sometimes you can get one line and sometimes you can get too many lines.

I keep it on flat list so it doesn't know what color use to draw next line - code could add color or depth to every point, or maybe it should add depth to every line as first element on list.

import pygame
import random as rd

def generate(x, y, depth):

    result = [ [(x,y)] ] # list for many lines

    for i in range(depth):
        x = result[0][-1][0] + rd.randrange(-20,20)
        y = result[0][-1][1] + rd.randrange(2,15)
        result[0].append((x, y))

        # should I generate next line from here?
        if rd.randrange(0,100) < 5:
            # add all generated lines to result
            result.extend(generate(x, y, depth-i))

    return result

# ---

pygame.init()

depth = rd.randrange(50, 80)
result = generate(400, 0, depth)
print(result)

screen = pygame.display.set_mode((800, 600))

screen.fill((0,0,0))

for line in result:
    for n, (point1, point2) in enumerate(zip(line, line[1:]), 1):
        #x1, y1 = point1
        #x2, y2 = point2
        c = 200-2*n
        pygame.draw.aaline(screen, (c,c,c), point1, point2, 1)

pygame.display.flip()

running = True

while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False

pygame.quit()    

---

EDIT: I add depth as first element in list so now it has better colors. If I use line(..., width=) then it is thicker at the top and thinner at the bottom

I was thinking to add random value to rd.randrange(-20+R,20+R) so some lines would go more left or more right.

import pygame
import random as rd


def generate(x, y, depth, c):

    result = [[c, (x,y)]]
    for i in range(depth):
        x = result[0][-1][0] + rd.randrange(-20,20)
        y = result[0][-1][1] + rd.randrange(2,15)
        result[0].append((x, y))

        if rd.randrange(0,100) < 5:
            result.extend(generate(x, y, depth-i, c+i))

    return result

# ---

pygame.init()

depth = rd.randrange(50, 80)
result = generate(400, 0, depth, 0)
print(result)

screen = pygame.display.set_mode((800, 600))

screen.fill((0,0,0))

for line in result:
    depth = line[0]
    line = line[1:]
    for n, (point1, point2) in enumerate(zip(line, line[1:]), depth):
        c = 200-2*n
        #pygame.draw.aaline(screen, (c,c,c), point1, point2, 1)
        pygame.draw.line(screen, (c,c,c), point1, point2, 12-n//6)

pygame.display.flip()

running = True

while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False

pygame.quit()