python snow screensaver

import math
import random
import sys
import pygame


class Snowflake:
    __slots__ = ("x", "y", "r", "vy", "vx", "layer", "phase", "amp")

    def __init__(self, W, H, layer):
        self.layer = layer
        self.reset(W, H, random_y=True)

    def reset(self, W, H, random_y=False):
        if self.layer == 0:
            self.r = random.uniform(1.0, 2.0)
            self.vy = random.uniform(25, 50)
        elif self.layer == 1:
            self.r = random.uniform(1.8, 3.0)
            self.vy = random.uniform(55, 90)
        else:
            self.r = random.uniform(2.8, 4.8)
            self.vy = random.uniform(90, 140)

        self.x = random.uniform(0, W)
        self.y = random.uniform(0, H) if random_y else random.uniform(-H * 0.2, 0)
        self.vx = random.uniform(-8, 8)
        self.phase = random.uniform(0, math.tau)
        self.amp = random.uniform(8, 28) * (0.4 + 0.4 * self.layer)

    def update(self, W, H, dt, wind, t):
        self.y += self.vy * dt
        self.x += (self.vx + wind) * dt
        self.x += math.sin(t + self.phase) * self.amp * dt * 0.15

        if self.x < -20:
            self.x = W + 20
        elif self.x > W + 20:
            self.x = -20

        if self.y > H + 10:
            self.reset(W, H)

    def draw(self, surf):
        alpha = 120 if self.layer == 0 else 180 if self.layer == 1 else 230
        pygame.draw.circle(
            surf, (255, 255, 255, alpha), (int(self.x), int(self.y)), int(self.r)
        )


def main():
    pygame.init()
    screen = pygame.display.set_mode((0, 0), pygame.FULLSCREEN)
    pygame.display.set_caption("Snowfall Screensaver")
    clock = pygame.time.Clock()

    W, H = screen.get_size()
    pygame.mouse.set_visible(False)
    pygame.event.set_grab(True)

    snow_layer = pygame.Surface((W, H), pygame.SRCALPHA)

    area = W * H
    base = int(area / (1920 * 1080) * 260)

    flakes = (
        [Snowflake(W, H, 0) for _ in range(int(base * 0.4))]
        + [Snowflake(W, H, 1) for _ in range(int(base * 0.35))]
        + [Snowflake(W, H, 2) for _ in range(int(base * 0.25))]
    )

    t = 0.0
    last_mouse = pygame.mouse.get_pos()

    running = True
    while running:
        dt = clock.tick(60) / 1000.0
        t += dt

        for event in pygame.event.get():
            if event.type in (
                pygame.QUIT,
                pygame.KEYDOWN,
                pygame.MOUSEBUTTONDOWN,
                pygame.MOUSEWHEEL,
            ):
                running = False
            elif event.type == pygame.MOUSEMOTION and event.rel != (0, 0):
                running = False

        if pygame.key.get_pressed()[pygame.K_ESCAPE]:
            running = False

        if pygame.mouse.get_pos() != last_mouse:
            running = False
        last_mouse = pygame.mouse.get_pos()

        # Pure black background
        screen.fill((0, 0, 0))

        # Gentle wind
        wind = 14 * math.sin(t * 0.25) + random.uniform(-1.5, 1.5)

        snow_layer.fill((0, 0, 0, 0))
        for flake in flakes:
            flake.update(W, H, dt, wind * (0.4 + 0.35 * flake.layer), t)
            flake.draw(snow_layer)

        screen.blit(snow_layer, (0, 0))
        pygame.display.flip()

    pygame.event.set_grab(False)
    pygame.mouse.set_visible(True)
    pygame.quit()
    sys.exit()


if __name__ == "__main__":
    main()