Python Pygame OpenGL Particle Demo

import pygame
from pygame.locals import *
import random
import math
import colorsys
import numpy as np
import ctypes # Needed for c_void_p

# Import OpenGL specific modules
from OpenGL.GL import *
from OpenGL.GL import shaders
# PyOpenGL_accelerate might improve performance if available
try:
    from OpenGL.GL.ARB.vertex_buffer_object import *
    _VBO_SUPPORT = True
except ImportError:
    _VBO_SUPPORT = False
    print("VBO support not found. Performance may be reduced.")


# --- Shader Code ---
VERTEX_SHADER = """
#version 330 core
layout (location = 0) in vec2 aPos;      // Vertex position (x, y)
layout (location = 1) in vec4 aColor;    // Vertex color (r, g, b, a)
layout (location = 2) in float aSize;     // Vertex size

out vec4 vColor; // Pass color to fragment shader

uniform mat4 projection; // Projection matrix

void main()
{
    gl_Position = projection * vec4(aPos.x, aPos.y, 0.0, 1.0);
    gl_PointSize = aSize; // Set the point size
    vColor = aColor;      // Pass color through
}
"""

FRAGMENT_SHADER = """
#version 330 core
in vec4 vColor; // Receive color from vertex shader
out vec4 FragColor; // Output fragment color

void main()
{
    // Optional: Make points round instead of square
    // vec2 coord = gl_PointCoord - vec2(0.5);
    // if(length(coord) > 0.5) {
    //     discard;
    // }
    FragColor = vColor; // Output the interpolated color
}
"""

# --- OpenGL Helper Functions ---

def create_shader_program():
    """Compiles and links the vertex and fragment shaders."""
    try:
        vertex_shader = shaders.compileShader(VERTEX_SHADER, GL_VERTEX_SHADER)
        fragment_shader = shaders.compileShader(FRAGMENT_SHADER, GL_FRAGMENT_SHADER)
        program = shaders.compileProgram(vertex_shader, fragment_shader)
        return program
    except shaders.ShaderCompilationError as e:
        print("Shader compilation error:")
        # Attempt to decode error message if bytes
        error_msg = e.args[0]
        if isinstance(error_msg, bytes):
            try:
                error_msg = error_msg.decode()
            except UnicodeDecodeError:
                error_msg = str(error_msg) # Fallback if decoding fails
        print(error_msg)
        print("Source:")
        # Attempt to decode source if bytes
        shader_source = e.source
        if isinstance(shader_source, bytes):
             try:
                 shader_source = shader_source.decode()
             except UnicodeDecodeError:
                 shader_source = str(shader_source) # Fallback
        print(shader_source)
        raise

def create_orthographic_matrix(left, right, bottom, top, near, far):
    """Creates an orthographic projection matrix."""
    matrix = np.identity(4, dtype=np.float32)
    matrix[0, 0] = 2.0 / (right - left)
    matrix[1, 1] = 2.0 / (top - bottom)
    matrix[2, 2] = -2.0 / (far - near)
    matrix[0, 3] = -(right + left) / (right - left)
    matrix[1, 3] = -(top + bottom) / (top - bottom)
    matrix[2, 3] = -(far + near) / (far - near)
    return matrix

# --- Constants ---
WIDTH, HEIGHT = 1000, 700

# --- Particle Class (Physics logic remains mostly the same) ---
class Particle:
    def __init__(self, x, y, color, size=2, velocity=None, life=None):
        self.x = float(x) # Ensure float
        self.y = float(y) # Ensure float
        self.original_color_rgb = color[:3] # Store RGB separately
        self.alpha = float(color[3] / 255.0) # Store alpha as 0.0-1.0
        self.size = float(size)
        self.original_size = float(size)
        self.life = float(life if life else random.randint(50, 200)) # Adjusted life for GL
        self.max_life = float(self.life)
        self.velocity = velocity if velocity else [random.uniform(-1, 1), random.uniform(-1, 1)]
        self.velocity = [float(v) for v in self.velocity] # Ensure float
        self.gravity = 0.03
        self.decay = random.uniform(0.97, 0.995) # Slightly adjusted decay

    def update(self, mouse_pos, particles, attract=True):
        # Update position
        self.x += self.velocity[0]
        self.y += self.velocity[1]

        # Apply gravity toward/away from mouse
        dx = mouse_pos[0] - self.x
        dy = mouse_pos[1] - self.y
        # Add small epsilon to prevent division by zero if distance is exactly 0
        distance_sq = dx*dx + dy*dy
        distance = max(math.sqrt(distance_sq), 1e-6) # Avoid sqrt(0) and div by zero

        force_dir_x = dx / distance
        force_dir_y = dy / distance

        if attract:
            self.velocity[0] += force_dir_x * self.gravity
            self.velocity[1] += force_dir_y * self.gravity
        else:
            self.velocity[0] -= force_dir_x * self.gravity * 2
            self.velocity[1] -= force_dir_y * self.gravity * 2

        # Apply velocity decay
        self.velocity[0] *= self.decay
        self.velocity[1] *= self.decay

        # --- Particle interaction removed for simplicity/performance focus ---

        # Bounce off edges
        bounce_dampening = -0.5
        if self.x <= 0:
            self.velocity[0] *= bounce_dampening
            self.x = 1.0 # Prevent sticking slightly inside
        elif self.x >= WIDTH:
             self.velocity[0] *= bounce_dampening
             self.x = float(WIDTH - 1) # Prevent sticking slightly inside

        if self.y <= 0:
            self.velocity[1] *= bounce_dampening
            self.y = 1.0 # Prevent sticking slightly inside
        elif self.y >= HEIGHT:
            self.velocity[1] *= bounce_dampening
            self.y = float(HEIGHT - 1) # Prevent sticking slightly inside


        # Update life and appearance
        self.life -= 1.0
        # Ensure max_life is not zero before division
        life_ratio = max(0.0, self.life / self.max_life if self.max_life > 0 else 0.0)
        self.size = self.original_size * life_ratio
        self.alpha = life_ratio # Alpha directly tied to life ratio

    # No draw method needed; data is collected externally

    def is_dead(self):
        return self.life <= 0 or self.size < 0.5 # Remove if too small

    def get_data(self):
        """Returns data tuple for VBO: (x, y, r, g, b, a, size)"""
        r, g, b = [c / 255.0 for c in self.original_color_rgb] # Normalize color
        # Ensure all data are floats
        return (float(self.x), float(self.y),
                float(r), float(g), float(b), float(self.alpha),
                float(self.size))

# --- Helper Functions (Color, Creation) ---
def get_color(hue):
    r, g, b = colorsys.hsv_to_rgb(hue, 1.0, 1.0)
    # Return RGBA tuple (alpha is set in Particle based on life)
    return (int(r * 255), int(g * 255), int(b * 255), 255) # Initial alpha 255

# Function to create burst particles
def create_burst(x, y, hue, count=5, vel_scale=2):
    particles = []
    for _ in range(count):
        size = random.uniform(2, 8)
        color = get_color(hue)
        vel = [random.uniform(-vel_scale, vel_scale), random.uniform(-vel_scale, vel_scale)]
        life = random.randint(40, 150) # Adjusted life
        particles.append(Particle(x, y, color, size, vel, life))
    return particles

# Function to create explosion particles
def create_explosion(x, y, hue, count=50):
    return create_burst(x, y, hue, count, vel_scale=5)


# --- Main Function ---
def main():
    pygame.init()
    # Setup OpenGL context flags
    pygame.display.gl_set_attribute(pygame.GL_CONTEXT_MAJOR_VERSION, 3)
    pygame.display.gl_set_attribute(pygame.GL_CONTEXT_MINOR_VERSION, 3)
    pygame.display.gl_set_attribute(pygame.GL_CONTEXT_PROFILE_MASK, pygame.GL_CONTEXT_PROFILE_CORE)
    pygame.display.gl_set_attribute(pygame.GL_DOUBLEBUFFER, 1) # Use double buffering
    pygame.display.gl_set_attribute(pygame.GL_DEPTH_SIZE, 24) # Optional depth buffer

    # Create window with OpenGL support
    screen = pygame.display.set_mode((WIDTH, HEIGHT), DOUBLEBUF | OPENGL)
    pygame.display.set_caption("Cosmic Particle Galaxy (OpenGL - Fixed)")
    clock = pygame.time.Clock()

    # --- OpenGL Initialization ---
    print("OpenGL Vendor:", glGetString(GL_VENDOR).decode())
    print("OpenGL Renderer:", glGetString(GL_RENDERER).decode())
    print("OpenGL Version:", glGetString(GL_VERSION).decode())
    print("GLSL Version:", glGetString(GL_SHADING_LANGUAGE_VERSION).decode())

    # Compile and link shaders
    shader_program = create_shader_program()
    glUseProgram(shader_program) # Use program once to set uniforms

    # Get uniform location for projection matrix
    proj_location = glGetUniformLocation(shader_program, "projection")
    if proj_location == -1:
         print("Warning: 'projection' uniform not found in shader.")

    # Create orthographic projection matrix (maps pixel coords to OpenGL coords)
    projection_matrix = create_orthographic_matrix(0, WIDTH, HEIGHT, 0, -1, 1) # Top-left origin
    # Apply the projection matrix
    glUniformMatrix4fv(proj_location, 1, GL_TRUE, projection_matrix)
    glUseProgram(0) # Unbind shader after setting uniform


    # --- VAO & VBO Setup ---
    if not _VBO_SUPPORT:
        print("Cannot run without VBO support. Exiting.")
        return

    # ** 1. Generate VAO **
    vao = glGenVertexArrays(1)
    # 2. Generate VBO
    vbo = glGenBuffers(1)

    # ** 3. Bind VAO **
    glBindVertexArray(vao)

    # 4. Bind VBO
    glBindBuffer(GL_ARRAY_BUFFER, vbo)

    # --- Get attribute locations (need shader bound) ---
    glUseProgram(shader_program)
    pos_loc = glGetAttribLocation(shader_program, "aPos")
    col_loc = glGetAttribLocation(shader_program, "aColor")
    size_loc = glGetAttribLocation(shader_program, "aSize")
    glUseProgram(0) # Unbind shader again

    if -1 in (pos_loc, col_loc, size_loc):
        print("Warning: Could not get all attribute locations (aPos, aColor, aSize).")
        print(f"Pos: {pos_loc}, Color: {col_loc}, Size: {size_loc}")
        # Consider exiting or handling this more gracefully if locations are missing

    # ** 5. Configure vertex attribute pointers **
    #    (These settings are stored in the bound VAO)
    stride = 7 * np.dtype(np.float32).itemsize # 7 floats, calculate size robustly
    offset_pos = ctypes.c_void_p(0)
    offset_col = ctypes.c_void_p(2 * np.dtype(np.float32).itemsize) # Offset after 2 pos floats
    offset_size = ctypes.c_void_p(6 * np.dtype(np.float32).itemsize) # Offset after pos (2) + color (4) floats

    # Position (Attribute 0)
    glVertexAttribPointer(pos_loc, 2, GL_FLOAT, GL_FALSE, stride, offset_pos)
    glEnableVertexAttribArray(pos_loc)
    # Color (Attribute 1)
    glVertexAttribPointer(col_loc, 4, GL_FLOAT, GL_FALSE, stride, offset_col)
    glEnableVertexAttribArray(col_loc)
    # Size (Attribute 2)
    glVertexAttribPointer(size_loc, 1, GL_FLOAT, GL_FALSE, stride, offset_size)
    glEnableVertexAttribArray(size_loc)

    # ** 6. Unbind VBO first, then VAO **
    #    (Order matters less here, but good practice)
    glBindBuffer(GL_ARRAY_BUFFER, 0)
    glBindVertexArray(0)


    # --- OpenGL State ---
    glClearColor(0.0, 0.0, 0.0, 1.0) # Black background
    glEnable(GL_BLEND) # Enable blending for transparency
    glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA) # Standard alpha blending
    glEnable(GL_PROGRAM_POINT_SIZE) # Allow shaders to set point size


    # --- Game Variables ---
    particles = []
    hue = 0.0
    attract_mode = True
    last_click_time = 0
    max_particles = 10000 # Can usually handle more particles now

    # --- Text Setup ---
    font = pygame.font.Font(None, 24)
    ui_dirty = True # Flag to rerender text only when needed
    text_surfaces = [] # Initialize list

    # --- Main Loop ---
    running = True
    while running:
        current_time = pygame.time.get_ticks()
        mouse_pos = pygame.mouse.get_pos()
        dt = clock.tick(60) / 1000.0 # Delta time in seconds

        # --- Event Handling ---
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_ESCAPE:
                     running = False
                elif event.key == pygame.K_SPACE:
                    attract_mode = not attract_mode
                    ui_dirty = True
                elif event.key == pygame.K_c:
                    particles = []
            elif event.type == pygame.MOUSEBUTTONDOWN:
                if event.button == 1: # Left click
                     if current_time - last_click_time > 100:
                         particles.extend(create_explosion(mouse_pos[0], mouse_pos[1], hue))
                         last_click_time = current_time

        # --- Particle Creation ---
        if pygame.mouse.get_rel() != (0, 0):
             if random.random() < 0.5:
                 particles.extend(create_burst(mouse_pos[0], mouse_pos[1], hue, count=2, vel_scale=1.5))

        hue = (hue + 0.002) % 1.0

        if random.random() < 0.05:
            angle = random.uniform(0, math.pi * 2)
            dist = random.uniform(50, 200)
            x = mouse_pos[0] + math.cos(angle) * dist
            y = mouse_pos[1] + math.sin(angle) * dist
            if 0 <= x < WIDTH and 0 <= y < HEIGHT:
                 color = get_color(hue)
                 particles.append(Particle(x, y, color, random.uniform(2, 5), life=random.randint(60,180)))


        # --- Particle Update and Data Collection ---
        live_particle_data = []
        num_live_particles = 0
        for i in range(len(particles) - 1, -1, -1):
            p = particles[i]
            p.update(mouse_pos, particles, attract_mode)
            if p.is_dead():
                del particles[i]
            else:
                live_particle_data.extend(p.get_data())
                num_live_particles += 1 # Count live particles correctly

        # Limit particle count
        if num_live_particles > max_particles:
            # Remove oldest particles (those added first)
            num_to_remove = num_live_particles - max_particles
            del particles[0 : num_to_remove]
            # Need to regenerate data if particles were removed this way
            live_particle_data = []
            num_live_particles = 0
            for p in particles:
                 live_particle_data.extend(p.get_data())
                 num_live_particles += 1

        # Set flag to update UI particle count
        if 'particle_count_text' not in locals() or particle_count_text != f"Particles: {num_live_particles}":
            ui_dirty = True


        # --- OpenGL Rendering ---
        glClear(GL_COLOR_BUFFER_BIT) # Clear the screen (color buffer)

        if live_particle_data and num_live_particles > 0:
             # Convert collected data to a NumPy array
             vertex_data = np.array(live_particle_data, dtype=np.float32)

             # Bind the VBO and upload the data
             glBindBuffer(GL_ARRAY_BUFFER, vbo)
             # Use glBufferData: potentially faster for rapidly changing data
             glBufferData(GL_ARRAY_BUFFER, vertex_data.nbytes, vertex_data, GL_DYNAMIC_DRAW)
             # Alternative: glBufferSubData if buffer size rarely changes (might be slower here)
             # glBufferSubData(GL_ARRAY_BUFFER, 0, vertex_data.nbytes, vertex_data)
             glBindBuffer(GL_ARRAY_BUFFER, 0) # Unbind VBO after upload

             # Activate the shader program
             glUseProgram(shader_program)

             # ** BIND THE VAO **
             # This sets up the VBO binding and attribute pointers automatically
             glBindVertexArray(vao)

             # Draw the points!
             glDrawArrays(GL_POINTS, 0, num_live_particles) # Use correct count

             # ** Unbind VAO ** and shader program
             glBindVertexArray(0)
             glUseProgram(0)


             # Bind the VBO and upload the data
             glBindBuffer(GL_ARRAY_BUFFER, vbo)
             glBufferData(GL_ARRAY_BUFFER, vertex_data.nbytes, vertex_data, GL_DYNAMIC_DRAW)
             glBindBuffer(GL_ARRAY_BUFFER, 0) # Unbind VBO after upload

             # Activate the shader program
             glUseProgram(shader_program)

             # ** BIND THE VAO **
             glBindVertexArray(vao)

             # Draw the points!
             glDrawArrays(GL_POINTS, 0, num_live_particles) # Use correct count

             # ** Unbind VAO ** and shader program
             glBindVertexArray(0)
             glUseProgram(0)

        # --- UI Text Rendering (Using Pygame Surface Blitting) ---
        if ui_dirty: # Only render text if mode or count changed
             mode_text = f"{'Attraction' if attract_mode else 'Repulsion'} Mode (SPACE)"
             clear_text = "Press C to clear"
             click_text = "Click to create explosions"
             particle_count_text = f"Particles: {num_live_particles}"

             text_surfaces = [
                 font.render(mode_text, True, (255, 255, 255)),
                 font.render(clear_text, True, (255, 255, 255)),
                 font.render(click_text, True, (255, 255, 255)),
                 font.render(particle_count_text, True, (255, 255, 255))
             ]
             ui_dirty = False # Reset flag

        # Blit cached text surfaces
        for i, text_surface in enumerate(text_surfaces):
             screen.blit(text_surface, (10, 10 + i * 25))


        # --- Display Update ---
        pygame.display.flip() # Swaps the front and back buffers

    # --- Cleanup ---
    print("Cleaning up OpenGL objects...")
    # Delete OpenGL objects (VAO first, then VBO, then Shader)
    glDeleteVertexArrays(1, [vao])
    glDeleteBuffers(1, [vbo])
    glDeleteProgram(shader_program)
    print("Cleanup complete.")
    pygame.quit()


if __name__ == "__main__":
    main()