import pygame

# Constants
WIDTH, HEIGHT = 800, 600
CELL_SIZE = 10
GRID_COLS = WIDTH // CELL_SIZE
GRID_ROWS = HEIGHT // CELL_SIZE

# Colors
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
RED   = (255, 0, 0)

# Directions: 0=Up, 1=Right, 2=Down, 3=Left
direction_vectors = {
    0: (0, -1),
    1: (1, 0),
    2: (0, 1),
    3: (-1, 0)
}
direction_names = {
    0: "Up",
    1: "Right",
    2: "Down",
    3: "Left"
}

pygame.init()
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Langton's Ant Visualization with Debug Info")
clock = pygame.time.Clock()
font = pygame.font.SysFont("Arial", 18)

# Initialize grid: False means white cell, True means black cell.
grid = [[False for _ in range(GRID_COLS)] for _ in range(GRID_ROWS)]

# Initialize ant at center of grid.
ant_col = GRID_COLS // 2
ant_row = GRID_ROWS // 2
ant_direction = 0  # Starting direction: Up
steps = 0

running = True
while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False

    # Run several simulation steps per frame.
    for _ in range(10):
        # Get the current cell's color.
        current_color = grid[ant_row][ant_col]

        # Langton's Ant rules:
        # If on a white cell, flip it to black, turn right.
        # If on a black cell, flip it to white, turn left.
        if not current_color:  # White cell
            grid[ant_row][ant_col] = True
            ant_direction = (ant_direction + 1) % 4
        else:  # Black cell
            grid[ant_row][ant_col] = False
            ant_direction = (ant_direction - 1) % 4

        # Move ant forward one cell in the current direction, with wrap-around.
        dx, dy = direction_vectors[ant_direction]
        ant_col = (ant_col + dx) % GRID_COLS
        ant_row = (ant_row + dy) % GRID_ROWS
        steps += 1

    # Clear screen (fill with white).
    screen.fill(WHITE)

    # Draw grid: fill black cells.
    for row in range(GRID_ROWS):
        for col in range(GRID_COLS):
            if grid[row][col]:
                rect = (col * CELL_SIZE, row * CELL_SIZE, CELL_SIZE, CELL_SIZE)
                pygame.draw.rect(screen, BLACK, rect)

    # Draw the ant as a red cell.
    ant_rect = (ant_col * CELL_SIZE, ant_row * CELL_SIZE, CELL_SIZE, CELL_SIZE)
    pygame.draw.rect(screen, RED, ant_rect)

    # Debug text (no background, just text):
    debug_text1 = font.render(f"Steps: {steps}", True, BLACK)
    debug_text2 = font.render(f"Ant Position: ({ant_col}, {ant_row})", True, BLACK)
    debug_text3 = font.render(f"Ant Direction: {direction_names[ant_direction]}", True, BLACK)
    current_cell_color = "Black" if grid[ant_row][ant_col] else "White"
    debug_text4 = font.render(f"Current Cell: {current_cell_color}", True, BLACK)

    # Blit debug text in the top-left corner.
    screen.blit(debug_text1, (10, 10))
    screen.blit(debug_text2, (10, 30))
    screen.blit(debug_text3, (10, 50))
    screen.blit(debug_text4, (10, 70))

    pygame.display.flip()
    clock.tick(60)

pygame.quit()