Particles

import random
from dataclasses import dataclass

import pygame

PARTICLE_COLORS = [tuple(color[:3]) for color in list(pygame.color.THECOLORS.values())]
GLOW_SHRINK_RATE = 0.15
GLOW_STARTING_RADIUS = 30.0
BLUR_RADIUS = 5
GRAVITY = 0.015


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


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, pygame.Color(color).lerp("black", 0.7), rect, 0)

    return background_surface


class ImageCache:
    def __init__(self):
        self.cache = {}
        self.blurred_cache = {}

        self._render_glow_surface()

    def _render_glow_surface(self):
        steps = int(GLOW_STARTING_RADIUS / GLOW_SHRINK_RATE)
        glow_radii = [GLOW_STARTING_RADIUS - i * GLOW_SHRINK_RATE for i in range(steps)]
        alpha_value = 150
        color = (*[255] * 3, alpha_value)

        for radius in glow_radii:
            size = [radius * 5] * 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, color, center, radius)
            pygame.transform.gaussian_blur(
                glow_surface,
                BLUR_RADIUS,
                repeat_edge_pixels=False,
                dest_surface=base_surface
            )
            self.store_blurred(radius, base_surface)

    def store_blurred(self, glow_radius, surface):
        key = round(glow_radius, 2)
        self.blurred_cache[key] = surface

    def get_surface(self, glow_radius, color):
        key = round(glow_radius, 2), color
        if key not in self.cache:
            blurred_cache = round(glow_radius, 2)
            glow_surface = self.blurred_cache[blurred_cache]
            colored_surface = glow_surface.copy()
            colored_surface.fill(color)
            colored_surface.blit(glow_surface, (0, 0), special_flags=pygame.BLEND_RGB_MULT)

            # return colored_surface  # Less caching
            self.cache[key] = colored_surface

        return self.cache[key]

    def __len__(self):
        return len(self.cache.keys())


class ParticleManager:
    def __init__(self):
        self.particles = []

    def update_particles(self, delta_time):
        for particle in self.particles[:]:
            particle.float_duration -= delta_time
            if particle.float_duration <= 0:
                particle.direction.y += GRAVITY

            particle.radius -= 0.03
            particle.glow_radius -= 0.15

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

            if particle.glow_radius <= 1.0:
                particle.alive = False
                self.particles.remove(particle)

    def emit_particles(self, particle_count=1):
        position = pygame.mouse.get_pos()

        for count in range(particle_count):
            color = random.choice(PARTICLE_COLORS)
            particle = Particle(
                pygame.Vector2(position), pygame.Vector2(random.uniform(-1, 1), -1), color=color
            )
            self.particles.append(particle)

    def draw_particle(self, surface, image_cache):
        for particle in self.particles:
            glow_surface = image_cache.get_surface(particle.glow_radius, particle.color)
            center = glow_surface.get_rect(center=particle.position)
            pygame.draw.circle(surface, particle.color, particle.position, particle.radius)
            surface.blit(glow_surface, center, special_flags=pygame.BLEND_RGB_ADD)


def main():
    pygame.init()
    display = pygame.display.set_mode((1500, 900))
    clock = pygame.Clock()
    background_surface = create_background(display)
    debug_font = pygame.font.SysFont("Consolas", 12, bold=True)

    image_cache = ImageCache()
    particle_manager = ParticleManager()

    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

        particle_manager.emit_particles(particle_count=1)

        elapsed_time += delta_time
        if elapsed_time >= fixed_delta_time:
            elapsed_time = 0
            particle_manager.update_particles(fixed_delta_time)

        debug_string = f"""\
                Cache size: {len(image_cache)}
                Glow cache size: {len(image_cache.blurred_cache)}
                Colors: {len(PARTICLE_COLORS)}
                Total particles: {len(particle_manager.particles)}
                FPS: {clock.get_fps():.2f}
                """
        debug_surface = debug_font.render(debug_string, True, "white")
        debug_shadow = debug_font.render(debug_string, True, "black")
        mouse_position = pygame.mouse.get_pos()

        display.fill("black")
        display.blit(background_surface, (0, 0))
        particle_manager.draw_particle(display, image_cache)
        display.blit(debug_shadow, (mouse_position[0] + 1, mouse_position[1]))
        display.blit(debug_surface, mouse_position)
        pygame.display.flip()
        delta_time = clock.tick() / 1000

        pygame.display.set_caption(debug_string)

    pygame.quit()


if __name__ == '__main__':
    main()