import pygame
import random
import math

class Particle:
    def __init__(
        self,
        pos,
        vel=(0, 0),
        acceleration=(0, 0),
        lifetime=2.0,         # seconds the particle lives
        decay_rate=1,         # how fast the lifetime decreases
        color=(255, 255, 255),
        size=5,
        size_decay=0,         # rate at which the size decreases
        glow=False,
        glow_intensity=0,     # multiplier for glow radius
        friction=0,           # slows down velocity over time
        gravity=0,            # constant acceleration downward
        bounce=False,
        bounce_damping=0.5,   # energy loss when bouncing off boundaries
        random_spread=0,      # randomness added to velocity each update
        spin=0,               # current rotation (unused for circles)
        spin_rate=0,          # rotation speed
        spin_decay=0,         # how spin_rate slows down
        trail=False,
        trail_length=10,      # how many previous positions to store
        fade=False,           # if True, the particle will fade (alpha decreases)
        shape='circle'        # could be 'circle' or 'square'
    ):
        # Position, velocity, and acceleration are stored as vectors.
        self.pos = pygame.math.Vector2(pos)
        self.vel = pygame.math.Vector2(vel)
        self.acceleration = pygame.math.Vector2(acceleration)

        # Lifetime and decay
        self.lifetime = lifetime
        self.initial_lifetime = lifetime
        self.decay_rate = decay_rate

        # Appearance settings
        self.color = color
        self.size = size
        self.initial_size = size
        self.size_decay = size_decay
        self.fade = fade
        self.alpha = 255

        # Effects options
        self.glow = glow
        self.glow_intensity = glow_intensity
        self.friction = friction
        self.gravity = gravity
        self.bounce = bounce
        self.bounce_damping = bounce_damping
        self.random_spread = random_spread

        # Spin (if using rotated images or shapes)
        self.spin = spin
        self.spin_rate = spin_rate
        self.spin_decay = spin_decay

        # Trail settings: stores previous positions to create a trailing effect
        self.trail = trail
        self.trail_length = trail_length
        self.positions = [self.pos.copy()] if trail else []

        # Shape of the particle ('circle' or 'square')
        self.shape = shape

    def update(self, dt, screen_rect):
        """
        dt: Delta time in seconds.
        screen_rect: The rectangle representing screen boundaries (for bounce).
        """
        # Decrease lifetime
        self.lifetime -= self.decay_rate * dt
        if self.lifetime < 0:
            self.lifetime = 0

        # Add random variation to velocity if enabled.
        if self.random_spread:
            self.vel.x += random.uniform(-self.random_spread, self.random_spread) * dt
            self.vel.y += random.uniform(-self.random_spread, self.random_spread) * dt

        # Apply gravity (as an extra acceleration on the y-axis)
        self.acceleration.y += self.gravity

        # Update velocity using acceleration.
        self.vel += self.acceleration * dt

        # Apply friction to slow the particle over time.
        self.vel *= (1 - self.friction * dt)

        # Update the position based on velocity.
        self.pos += self.vel * dt

        # Bounce off screen edges if enabled.
        if self.bounce:
            if self.pos.x - self.size < screen_rect.left or self.pos.x + self.size > screen_rect.right:
                self.vel.x = -self.vel.x * self.bounce_damping
                if self.pos.x - self.size < screen_rect.left:
                    self.pos.x = screen_rect.left + self.size
                elif self.pos.x + self.size > screen_rect.right:
                    self.pos.x = screen_rect.right - self.size
            if self.pos.y - self.size < screen_rect.top or self.pos.y + self.size > screen_rect.bottom:
                self.vel.y = -self.vel.y * self.bounce_damping
                if self.pos.y - self.size < screen_rect.top:
                    self.pos.y = screen_rect.top + self.size
                elif self.pos.y + self.size > screen_rect.bottom:
                    self.pos.y = screen_rect.bottom - self.size

        # Update spin
        self.spin += self.spin_rate * dt
        self.spin_rate *= (1 - self.spin_decay * dt)

        # Decrease size if size_decay is set.
        if self.size_decay:
            self.size = max(0, self.size - self.size_decay * dt)

        # Fade out by decreasing alpha based on lifetime.
        if self.fade:
            self.alpha = int(255 * (self.lifetime / self.initial_lifetime))
            if self.alpha < 0:
                self.alpha = 0

        # Update trail history.
        if self.trail:
            self.positions.append(self.pos.copy())
            if len(self.positions) > self.trail_length:
                self.positions.pop(0)

    def draw(self, surface):
        # Set drawing color with alpha if fading.
        if self.fade:
            draw_color = (*self.color, self.alpha)
        else:
            draw_color = self.color

        # Draw the particle's trail if enabled.
        if self.trail and len(self.positions) > 1:
            for i in range(1, len(self.positions)):
                start = self.positions[i - 1]
                end = self.positions[i]
                pygame.draw.line(surface, draw_color, start, end, int(self.size))

        # Draw glow effect if enabled.
        if self.glow:
            glow_radius = int(self.size * (1 + self.glow_intensity))
            glow_color = (*self.color, self.alpha) if self.fade else self.color
            glow_surface = pygame.Surface((glow_radius * 2, glow_radius * 2), pygame.SRCALPHA)
            pygame.draw.circle(glow_surface, glow_color, (glow_radius, glow_radius), glow_radius)
            surface.blit(glow_surface, (self.pos.x - glow_radius, self.pos.y - glow_radius), special_flags=pygame.BLEND_ADD)

        # Draw the particle shape.
        if self.shape == 'circle':
            pygame.draw.circle(surface, draw_color, (int(self.pos.x), int(self.pos.y)), int(self.size))
        elif self.shape == 'square':
            rect = pygame.Rect(self.pos.x - self.size, self.pos.y - self.size, self.size * 2, self.size * 2)
            pygame.draw.rect(surface, draw_color, rect)
        # Additional shapes (e.g., triangles, rotated images) can be added here.

    def is_dead(self):
        """Return True if the particle should be removed."""
        return self.lifetime <= 0 or self.size <= 0