Particles

import dataclasses
import random
from collections import deque
from dataclasses import dataclass
from dataclasses import fields

import pygame

POOL_MAX = 2_000
PARTICLE_COLORS = [tuple(color[:3]) for color in list(pygame.color.THECOLORS.values())]
GRAVITY = 0.015

image_cache = {}
particle_pool = deque(maxlen=POOL_MAX)


@dataclass(slots=True)
class Particle:
    position: pygame.Vector2
    direction: pygame.Vector2
    color: tuple
    radius: float = 5.0
    glow_radius: float = 30.0
    float_duration: float = 0.1
    speed: int = 100
    alive: bool = True

    def reset(self, position, direction, color):
        self.position.update(position)
        self.direction.update(direction)
        self.color = color

        for data_field in fields(self):
            if data_field.default is not dataclasses.MISSING:
                setattr(self, data_field.name, data_field.default)

    def update(self, delta_time):
        self.float_duration -= delta_time
        if self.float_duration <= 0:
            self.direction.y += GRAVITY

        self.radius -= 0.03
        self.glow_radius -= 0.15

        if self.glow_radius <= 1.0:
            self.alive = False

        self.position += self.direction * self.speed * delta_time


def create_background(surface):
    rectangle_colors = [
        ["red", "black", "skyblue"],
        ["black", "green", "white"],
        ["grey", "black", "blue"]
    ]
    margin = 50
    background_surface = surface.copy()
    rectangle_width = (background_surface.get_width() - (margin * 2)) // len(rectangle_colors[0])
    rectangle_height = (background_surface.get_height() - (margin * 2)) // len(rectangle_colors)

    for y, colors in enumerate(rectangle_colors):
        for x, color in enumerate(colors):
            rect = pygame.Rect(
                (x * rectangle_width + margin, y * rectangle_height + margin),
                (rectangle_width, rectangle_height)
            )
            pygame.draw.rect(background_surface, color, rect, 0)

    return background_surface


def draw_particle(surface, particles):
    global image_cache

    for particle in particles:
        key = (particle.color, particle.glow_radius)
        if key not in image_cache:
            size = [particle.glow_radius * 2] * 2
            base_surface = pygame.Surface(size, flags=pygame.SRCALPHA)
            glow_surface = base_surface.copy()
            glow_surface.set_colorkey("black")

            center = base_surface.get_rect().center
            pygame.draw.circle(glow_surface, (*particle.color, 150), center, particle.glow_radius)
            glow_surface = pygame.transform.gaussian_blur(glow_surface, 5, repeat_edge_pixels=False)
            base_surface.blit(glow_surface, (0, 0))

            image_cache[key] = base_surface

        base_surface = image_cache[key]
        pygame.draw.circle(surface, particle.color, particle.position, particle.radius)
        surface.blit(base_surface, base_surface.get_rect(center=particle.position), special_flags=pygame.BLEND_RGB_ADD)


def update_particles(particles, delta_time):
    global particle_pool

    for particle in particles[:]:
        particle.update(delta_time)
        if not particle.alive:
            particles.remove(particle)
            particle_pool.append(particle)


def create_particles(particles, particle_count=1):
    global particle_pool

    position = pygame.mouse.get_pos()

    for count in range(particle_count):
        color = random.choice(PARTICLE_COLORS)

        if particle_pool:
            particle = particle_pool.popleft()
            particle.reset(
                position=position,
                direction=(random.uniform(-1, 1), -1),
                color=color
            )
        else:
            particle = Particle(pygame.Vector2(position), pygame.Vector2(random.uniform(-1, 1), -1), color=color)
        particles.append(particle)


def main():
    global image_cache, particle_pool

    pygame.init()
    display = pygame.display.set_mode((1500, 900))
    clock = pygame.Clock()
    background_surface = create_background(display)

    particles = []
    for _ in range(POOL_MAX):
        particle_pool.append(Particle(pygame.Vector2(0, 0), pygame.Vector2(0, 0), (0, 0, 0, 0)))

    elapsed_time = 0
    fixed_delta_time = 1 / 60
    delta_time = 0

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

        create_particles(particles, particle_count=1)

        elapsed_time += delta_time
        if elapsed_time >= fixed_delta_time:
            elapsed_time = 0
            update_particles(particles, fixed_delta_time)

        display.fill("black")
        display.blit(background_surface, (0, 0))
        draw_particle(display, particles)
        pygame.display.flip()
        delta_time = clock.tick() / 1000
        pygame.display.set_caption(
            f"Cache size: {len(image_cache.keys())} "
            f"Particle Pool: {len(particle_pool)} "
            f"FPS: {clock.get_fps():.2f}"
        )


if __name__ == '__main__':
    main()