wormhole_game_of_life_logic

import os
import numpy as np
from PIL import Image
import sys

class WormholeGameOfLife:
    def __init__(self, input_dir):
        self.input_dir = input_dir

        # Load input images
        self.starting_position = self.load_binary_image(os.path.join(input_dir, 'starting_position.png'))
        self.horizontal_tunnel = self.load_image(os.path.join(input_dir, 'horizontal_tunnel.png'))
        self.vertical_tunnel = self.load_image(os.path.join(input_dir, 'vertical_tunnel.png'))

        # Check all images have the same dimensions
        height, width = self.starting_position.shape
        assert self.horizontal_tunnel.shape[:2] == (height, width), "Image dimensions don't match"
        assert self.vertical_tunnel.shape[:2] == (height, width), "Image dimensions don't match"

        self.height = height
        self.width = width

        # Find wormhole pairs
        self.h_wormholes = self.find_wormhole_pairs(self.horizontal_tunnel)
        self.v_wormholes = self.find_wormhole_pairs(self.vertical_tunnel)

        print(f"Horizontal wormhole pairs: {len(self.h_wormholes) // 2}")
        print(f"Vertical wormhole pairs: {len(self.v_wormholes) // 2}")

        # Create a set of all wormhole positions for quick lookups
        self.wormhole_positions = set(self.h_wormholes.keys()) | set(self.v_wormholes.keys())

        # Current state of the grid
        self.grid = self.starting_position.copy()

        # Load expected results if available
        try:
            self.expected_1 = self.load_binary_image(os.path.join(input_dir, 'expected-1.png'))
            print(f"Expected alive cells after 1 iteration: {np.sum(self.expected_1)}")
        except:
            self.expected_1 = None

    def load_binary_image(self, filepath):
        """Load a binary image (white=alive, black=dead)"""
        img = Image.open(filepath).convert('L')
        binary = np.array(img) > 128  # Convert to binary using threshold
        return binary

    def load_image(self, filepath):
        """Load a color image"""
        img = Image.open(filepath)
        return np.array(img)

    def find_wormhole_pairs(self, tunnel_img):
        """Find all wormhole pairs by color"""
        wormhole_map = {}  # (y, x) -> (y', x')
        color_to_pos = {}

        # Collect all positions for each color
        for y in range(self.height):
            for x in range(self.width):
                color = tuple(tunnel_img[y, x][:3])
                if color != (0, 0, 0):  # Not black
                    if color not in color_to_pos:
                        color_to_pos[color] = []
                    color_to_pos[color].append((y, x))

        # Create wormhole mappings between pairs
        for color, positions in color_to_pos.items():
            if len(positions) == 2:
                pos1, pos2 = positions
                wormhole_map[pos1] = pos2
                wormhole_map[pos2] = pos1

        return wormhole_map

    def get_neighbors(self, y, x):
        """
        Get the 8 neighbors of a cell, considering wormhole tunnels
        Wormholes connect two points, allowing neighbors to be counted through the wormhole
        """
        neighbors = []

        # Check each of the 8 adjacent cells
        for dy in [-1, 0, 1]:
            for dx in [-1, 0, 1]:
                if dy == 0 and dx == 0:
                    continue  # Skip the cell itself

                ny, nx = y + dy, x + dx

                # Check for out of bounds
                if ny < 0 or ny >= self.height or nx < 0 or nx >= self.width:
                    continue

                # For diagonal neighbors, just add them directly
                if dy != 0 and dx != 0:
                    neighbors.append((ny, nx))
                    continue

                # For orthogonal neighbors, check if they're at a wormhole entrance
                # If so, count both the regular neighbor AND the neighbor through the wormhole

                # Vertical wormhole (top or bottom neighbor)
                if dx == 0 and (ny, nx) in self.v_wormholes:
                    # Get the corresponding position at the other end of the wormhole
                    dest_y, dest_x = self.v_wormholes[(ny, nx)]
                    # Add both the regular neighbor and the wormhole destination
                    neighbors.append((ny, nx))
                    neighbors.append((dest_y, dest_x))

                # Horizontal wormhole (left or right neighbor)
                elif dy == 0 and (ny, nx) in self.h_wormholes:
                    # Get the corresponding position at the other end of the wormhole
                    dest_y, dest_x = self.h_wormholes[(ny, nx)]
                    # Add both the regular neighbor and the wormhole destination
                    neighbors.append((ny, nx))
                    neighbors.append((dest_y, dest_x))

                # Regular neighbor (not at a wormhole)
                else:
                    neighbors.append((ny, nx))

        return neighbors

    def count_live_neighbors(self, y, x):
        """Count the number of live neighbors, with wormhole connections"""
        count = 0
        for ny, nx in self.get_neighbors(y, x):
            if self.grid[ny, nx]:
                count += 1
        return count

    def step(self):
        """Advance the simulation by one step"""
        new_grid = np.zeros_like(self.grid)

        # Process each cell in the grid
        for y in range(self.height):
            for x in range(self.width):
                # Special cases for wormhole positions
                if (y, x) in self.wormhole_positions:
                    # Apply special rule: wormhole positions stay alive with any neighbors
                    if self.grid[y, x] and self.count_live_neighbors(y, x) > 0:
                        new_grid[y, x] = True
                    # Otherwise standard Game of Life rules
                    else:
                        # Apply standard rules (for birth only)
                        if not self.grid[y, x] and self.count_live_neighbors(y, x) == 3:
                            new_grid[y, x] = True
                else:
                    # For non-wormhole positions, apply standard Game of Life rules
                    live_neighbors = self.count_live_neighbors(y, x)

                    if self.grid[y, x]:  # Cell is alive
                        # A live cell with 2 or 3 live neighbors survives
                        if live_neighbors in [2, 3]:
                            new_grid[y, x] = True
                    else:  # Cell is dead
                        # A dead cell with exactly 3 live neighbors becomes alive
                        if live_neighbors == 3:
                            new_grid[y, x] = True

        # Update the grid state
        self.grid = new_grid

    def simulate(self, save_at_iterations, output_dir):
        """Run continuously up to max iteration, saving and verifying at specific steps."""
        max_iter = max(save_at_iterations)

        # Preload expected outputs if available
        expected_outputs = {}
        for iter_num in save_at_iterations:
            expected_path = os.path.join(self.input_dir, f'expected-{iter_num}.png')
            if os.path.exists(expected_path):
                expected_outputs[iter_num] = self.load_binary_image(expected_path)

        for i in range(1, max_iter + 1):
            self.step()

            if i in save_at_iterations:
                output_path = os.path.join(output_dir, f'{i}.png')
                self.save_output(output_path)
                print(f"Saved output at iteration {i}")

                # Compare with expected output if available
                if i in expected_outputs:
                    expected = expected_outputs[i]
                    diff = np.sum(self.grid != expected)
                    if diff > 0:
                        diff_coords = np.where(self.grid != expected)
                        print(f"Warning: Iteration {i} differs from expected by {diff} cells")
                        for idx in range(min(5, len(diff_coords[0]))):
                            y, x = diff_coords[0][idx], diff_coords[1][idx]
                            print(f"  Diff at ({y}, {x}): Got {self.grid[y, x]}, Expected {expected[y, x]}")
                    else:
                        print(f"Iteration {i} matches expected output exactly ✅")


    def save_output(self, output_path):
        """Save the current grid state as an image"""
        img = Image.new('L', (self.width, self.height))
        pixels = []

        for y in range(self.height):
            for x in range(self.width):
                pixels.append(255 if self.grid[y, x] else 0)

        img.putdata(pixels)
        img.save(output_path)

        # Print the count of live cells in the output
        live_count = np.sum(self.grid)
        print(f"Live cells in {output_path}: {live_count}")

        # Display coordinates of alive cells if there aren't too many
        if live_count > 0 and live_count <= 15:
            alive_coords = np.where(self.grid)
            print("Alive cell coordinates:")
            for i in range(len(alive_coords[0])):
                y, x = alive_coords[0][i], alive_coords[1][i]
                is_wormhole = (y, x) in self.wormhole_positions
                print(f"  ({y}, {x}){' (wormhole)' if is_wormhole else ''}")

def run_simulation(input_dir, output_dir):
    os.makedirs(output_dir, exist_ok=True)

    game = WormholeGameOfLife(input_dir)

    # Save initial state (optional for debugging)
    game.save_output(os.path.join(output_dir, 'initial.png'))

    # Define milestones where we save outputs
    milestones = [1, 10, 100, 1000]

    # Run full simulation
    game.simulate(milestones, output_dir)

def print_example_differences():
    """Print the differences between examples to understand the wormhole impact"""
    import os

    # Load standard Game of Life results from example-1
    std_img_1 = Image.open(os.path.join('std_output/example-1', '1.png')).convert('L')
    std_result_1 = np.array(std_img_1) > 128

    # Load expected results from example-0
    expected_img_0 = Image.open(os.path.join('assets/archive/example-0', 'expected-1.png')).convert('L')
    expected_result_0 = np.array(expected_img_0) > 128

    # Compare
    print(f"Standard Game of Life (ex-1) alive cells: {np.sum(std_result_1)}")
    print(f"Expected with wormholes (ex-0) alive cells: {np.sum(expected_result_0)}")

if __name__ == "__main__":
    if len(sys.argv) < 3:
        print("Usage: python wormhole_game_of_life.py <input_directory> <output_directory>")
        sys.exit(1)

    input_dir = sys.argv[1]
    output_dir = sys.argv[2]

    # Optionally print differences between examples
    # print_example_differences()

    run_simulation(input_dir, output_dir)