sense_hat_worm.py

#!/usr/bin/env python3
"""
NetWare 4.11 Worm Screensaver for Raspberry Pi Sense HAT
Recreates the classic red worm screensaver with CPU load responsiveness
For multi-core systems, displays different colored worms for each core
"""

from sense_hat import SenseHat
import psutil
import time
import random
import math
import argparse

# Initialize the Sense HAT
hat = SenseHat()

# Define colors for different CPU cores
CORE_COLORS = [
    (255, 0, 0),      # Core 0: Red
    (0, 255, 0),      # Core 1: Green
    (0, 0, 255),      # Core 2: Blue
    (255, 255, 0),    # Core 3: Yellow
    (0, 255, 255),    # Core 4: Cyan
    (255, 0, 255),    # Core 5: Magenta
    (255, 165, 0),    # Core 6: Orange
    (255, 128, 128),  # Core 7: Pink
]

BLACK = (0, 0, 0)


class Worm:
    """Represents a single worm moving on the LED matrix"""

    def __init__(self, color, core_id=0):
        """Initialize a worm

        Args:
            color: RGB tuple for worm color
            core_id: CPU core this worm represents
        """
        self.color = color
        self.core_id = core_id

        # Position
        self.x = random.randint(0, 7)
        self.y = random.randint(0, 7)

        # Direction (one of 8 directions)
        self.direction = random.randint(0, 7)

        # Trail - list of (x, y, intensity) tuples
        self.trail = []

        # Base tail length (will be affected by CPU load)
        self.base_tail_length = 5
        self.current_tail_length = 5

        # Movement timing - each worm moves independently
        self.last_move_time = time.time() - (core_id * 0.1)
        self.move_interval = 0.5  # Will be updated based on CPU load

    def _move(self):
        """Move the worm in current direction or change direction"""
        # Randomly decide to change direction (33% chance)
        if random.random() < 0.33:
            self.direction = random.randint(0, 7)

        # Calculate new position based on direction
        dx, dy = self._get_direction_vector()

        new_x = self.x + dx
        new_y = self.y + dy

        # Bounce off edges instead of wrapping
        if new_x < 0 or new_x > 7:
            new_x = max(0, min(7, new_x))
            # Reflect direction - flip horizontal movement
            if self.direction == 0:      # Right -> Left
                self.direction = 4
            elif self.direction == 1:    # Diagonal right-down -> Diagonal left-down
                self.direction = 3
            elif self.direction == 3:    # Diagonal left-down -> Diagonal right-down
                self.direction = 1
            elif self.direction == 4:    # Left -> Right
                self.direction = 0
            elif self.direction == 5:    # Diagonal left-up -> Diagonal right-up
                self.direction = 7
            elif self.direction == 7:    # Diagonal right-up -> Diagonal left-up
                self.direction = 5

        if new_y < 0 or new_y > 7:
            new_y = max(0, min(7, new_y))
            # Reflect direction - flip vertical movement
            if self.direction == 2:      # Down -> Up
                self.direction = 6
            elif self.direction == 1:    # Diagonal right-down -> Diagonal right-up
                self.direction = 7
            elif self.direction == 3:    # Diagonal left-down -> Diagonal left-up
                self.direction = 5
            elif self.direction == 5:    # Diagonal left-up -> Diagonal left-down
                self.direction = 3
            elif self.direction == 6:    # Up -> Down
                self.direction = 2
            elif self.direction == 7:    # Diagonal right-up -> Diagonal right-down
                self.direction = 1

        # Add current position to trail
        self.trail.append((self.x, self.y))

        # Update position
        self.x = new_x
        self.y = new_y

        # Trim trail to maximum length
        if len(self.trail) > self.current_tail_length:
            self.trail = self.trail[-self.current_tail_length:]

    def _get_direction_vector(self):
        """Get dx, dy vector for current direction

        Directions (0-7):
        0=Right, 1=DiagRightDown, 2=Down, 3=DiagLeftDown,
        4=Left, 5=DiagLeftUp, 6=Up, 7=DiagRightUp
        """
        directions = [
            (1, 0),    # 0: Right
            (1, 1),    # 1: Diagonal right-down
            (0, 1),    # 2: Down
            (-1, 1),   # 3: Diagonal left-down
            (-1, 0),   # 4: Left
            (-1, -1),  # 5: Diagonal left-up
            (0, -1),   # 6: Up
            (1, -1),   # 7: Diagonal right-up
        ]
        return directions[self.direction]

    def get_pixels(self):
        """Get list of (x, y, color) tuples for this worm

        Returns:
            List of pixels representing worm head and trail
        """
        pixels = []

        # Head - full brightness
        pixels.append((self.x, self.y, self.color))

        # Trail - fade out from bright to dark
        trail_length = len(self.trail)
        for i, (x, y) in enumerate(self.trail):
            # Calculate fade: 0% at head (start of trail), 100% at end of trail (oldest pixel)
            # Reverse the position so oldest pixels fade to black
            fade_position = (trail_length - i) / max(1, trail_length)

            # Calculate intensity (0-255, but reduce base so it fades to black)
            intensity = int(255 * (1.0 - fade_position) * 0.5)

            # Apply intensity to color
            faded_color = (
                int(self.color[0] * intensity / 255),
                int(self.color[1] * intensity / 255),
                int(self.color[2] * intensity / 255)
            )

            pixels.append((x, y, faded_color))

        return pixels


class WormScreensaver:
    """Main screensaver controller"""

    def __init__(self):
        """Initialize the screensaver with one worm per CPU core"""
        self.num_cores = psutil.cpu_count()
        self.worms = []

        # Create a worm for each CPU core
        for i in range(self.num_cores):
            color = CORE_COLORS[i % len(CORE_COLORS)]
            worm = Worm(color, core_id=i)
            self.worms.append(worm)

        print(f"Initialized {self.num_cores} worms (one per CPU core)")

    def get_cpu_load(self):
        """Get overall CPU load percentage

        Returns:
            float: CPU usage percentage (0-100)
        """
        return psutil.cpu_percent(interval=0.05)

    def get_core_loads(self):
        """Get per-core CPU load percentages

        Returns:
            list: CPU usage percentages for each core
        """
        return psutil.cpu_percent(interval=0.05, percpu=True)

    def update(self):
        """Update all worms based on CPU load"""
        core_loads = self.get_core_loads()

        for i, worm in enumerate(self.worms):
            # Use per-core load if available, otherwise use overall load
            cpu_percent = core_loads[i] if i < len(core_loads) else core_loads[0]
            worm.update(cpu_percent)

    def render(self):
        """Render all worms to the LED matrix"""
        # Start with black canvas
        pixels = [BLACK] * 64

        # Collect all worm pixels
        for worm in self.worms:
            worm_pixels = worm.get_pixels()
            for x, y, color in worm_pixels:
                # Check bounds
                if 0 <= x < 8 and 0 <= y < 8:
                    idx = y * 8 + x

                    # Blend colors if multiple worms overlap
                    current = pixels[idx]

                    # Simple additive blending (cap at 255)
                    blended = (
                        min(255, current[0] + color[0]),
                        min(255, current[1] + color[1]),
                        min(255, current[2] + color[2])
                    )

                    pixels[idx] = blended

        hat.set_pixels(pixels)

    def print_stats(self):
        """Print current stats to console"""
        core_loads = self.get_core_loads()
        avg_load = sum(core_loads) / len(core_loads)

        # Format core loads
        core_str = " | ".join([f"Core {i}: {load:5.1f}%" for i, load in enumerate(core_loads)])
        tail_lengths = " | ".join([f"W{i}: {worm.current_tail_length:.0f}" for i, worm in enumerate(self.worms)])

        print(f"Avg: {avg_load:5.1f}% | {core_str} | Tail: {tail_lengths}", end='\r')


def main():
    """Main screensaver loop"""
    try:
        print("NetWare 4.11 Worm Screensaver for Sense HAT")
        print("==========================================")

        # Parse command-line arguments
        parser = argparse.ArgumentParser(description='NetWare 4.11 Worm Screensaver')
        parser.add_argument('--stats', action='store_true', help='Enable stats output')
        parser.add_argument('--interval', type=float, default=0.3,help='Update interval in seconds (0..whatever)')
        args = parser.parse_args()

        print("\nStarting screensaver...")
        print("Red worm adapts speed and tail length to CPU load")
        print("Multi-core systems show different colored worms per core")
        print("Press Ctrl+C to exit\n")

        screensaver = WormScreensaver()

        # Main loop - worms move with speed dependent on CPU load
        last_move = time.time()

        while True:
            current_time = time.time()

            # Get current CPU load for tail length and speed adjustment
            core_loads = screensaver.get_core_loads()

            for i, worm in enumerate(screensaver.worms):
                cpu_percent = core_loads[i] if i < len(core_loads) else core_loads[0]

                # Update tail length based on CPU load
                load_factor = cpu_percent / 100.0
                worm.current_tail_length = worm.base_tail_length + int(load_factor * 10)

                # Calculate dynamic movement interval based on this worm's CPU load
                # Interval: 0.5 seconds at idle, 0.2 seconds at 100% load
                worm.move_interval = 0.5 - (load_factor * 0.5)

                # Move worm if its interval has elapsed
                if current_time - worm.last_move_time >= worm.move_interval:
                    worm.last_move_time = current_time
                    worm._move()

            # Render to LED matrix
            screensaver.render()

            # Print stats if enabled
            if args.stats:
                screensaver.print_stats()

            # Sleep briefly to avoid busy waiting
            time.sleep(args.interval)

    except KeyboardInterrupt:
        print("\n\nShutting down gracefully...")
        hat.clear()
    except Exception as e:
        print(f"\nError: {e}")
        import traceback
        traceback.print_exc()
        hat.clear()


if __name__ == "__main__":
    main()