environment.py

"""
The Environment class manages the entire sheep herding simulation.
It coordinates interactions between sheep, dog, and environmental factors.
"""

import pygame
import sys
import random
from vector_utils import Vector2D
from sheep import Sheep
from dog import ShepherdDog
import config

class Environment:
    """
    Main environment class that manages the sheep herding simulation.

    Handles:
    - Simulation loop and timing
    - Entity creation and management
    - Physics updates
    - Rendering and visualization
    - User input and controls
    - Performance monitoring
    """

    
    
    def __init__(self):
        print("DEBUG: Initializing Environment...")
        # Initialize Pygame
        pygame.init()
        print("DEBUG: Pygame initialized")

        try:
            self.screen = pygame.display.set_mode((config.SCREEN_WIDTH, config.SCREEN_HEIGHT))
            print(f"DEBUG: Display set to {config.SCREEN_WIDTH}x{config.SCREEN_HEIGHT}")
            pygame.display.set_caption("AI Sheep Herding Simulation - Advanced Neighbor Selection")
            self.clock = pygame.time.Clock()
            print("DEBUG: Clock initialized")
        except Exception as e:
            print(f"DEBUG: Error initializing display: {str(e)}")
            raise

        # Initialize fonts for UI and debug text
        self.font = pygame.font.Font(None, 24)
        self.small_font = pygame.font.Font(None, 18)

        # Create entities - initialize them as None first
        self.sheep_flock = []
        self.dog = None  # Initialize as None, will be created in _initialize_entities
        try:
            self.target_position = Vector2D(*config.ENVIRONMENT_CONFIG['target_position'])
            print(f"DEBUG: Target position set to {self.target_position}")
        except Exception as e:
            print(f"DEBUG: Error setting target position: {str(e)}")
            raise

        # Simulation state
        self.running = True
        self.paused = False
        self.frame_count = 0
        self.simulation_time = 0.0

        # Performance tracking
        self.fps_history = []
        self.performance_metrics = {}

        # Initialize simulation entities
        print("DEBUG: Initializing entities...")
        try:
            self._initialize_entities()
            print("DEBUG: Entities initialized")
        except Exception as e:
            print(f"DEBUG: Error initializing entities: {str(e)}")
            raise

    def _initialize_entities(self):
        """Initialize sheep flock and shepherd dog."""
        print("DEBUG: Initializing sheep flock...")
        self.sheep_flock = []
        # Use the count from SHEEP_CONFIG
        sheep_count = config.SHEEP_CONFIG['count']
        print(f"DEBUG: Creating {sheep_count} sheep...")

        # Print some sheep configuration details for debugging
        print("DEBUG: Sheep configuration:")
        print(f"DEBUG:   size: {config.SHEEP_CONFIG['size']}")
        print(f"DEBUG:   max_speed: {config.SHEEP_CONFIG['max_speed']}")
        print(f"DEBUG:   max_force: {config.SHEEP_CONFIG['max_force']}")
        print("DEBUG:   Flocking parameters:")
        flocking_params = config.SHEEP_CONFIG['flocking_params']
        for param, value in flocking_params.items():
            print(f"DEBUG:     {param}: {value}")
        print("DEBUG:   Fear response:")
        fear_response = config.SHEEP_CONFIG['fear_response']
        for param, value in fear_response.items():
            print(f"DEBUG:     {param}: {value}")

        for i in range(sheep_count):
            # Random position away from target
            x = random.uniform(50, config.SCREEN_WIDTH - 300)
            y = random.uniform(50, config.SCREEN_HEIGHT - 50)

            # Ensure sheep don't start too close to target
            while Vector2D(x, y).distance_to(self.target_position) < 150:
                x = random.uniform(50, config.SCREEN_WIDTH - 300)
                y = random.uniform(50, config.SCREEN_HEIGHT - 50)

            # Pass the entire SHEEP_CONFIG to the Sheep constructor
            try:
                sheep = Sheep(x, y, config.SHEEP_CONFIG)
                self.sheep_flock.append(sheep)
                if i < 3:  # Print position for first few sheep to avoid too much output
                    print(f"DEBUG: Created sheep {i} at position ({x:.1f}, {y:.1f})")
            except Exception as e:
                print(f"DEBUG: Error creating sheep {i}: {str(e)}")
                raise

        # Create shepherd dog positioned away from the sheep
        print("DEBUG: Creating shepherd dog...")
        try:
            # Define dog position variables - these must be defined before use
            dog_x = random.uniform(50, 200)
            dog_y = random.uniform(50, config.SCREEN_HEIGHT - 50)

            # Create the shepherd dog with its target position
            self.dog = ShepherdDog(dog_x, dog_y, config.DOG_CONFIG, self.target_position)
            print(f"DEBUG: Dog created at position ({dog_x:.1f}, {dog_y:.1f}) with target at {self.target_position}")
        except Exception as e:
            print(f"DEBUG: Error creating dog: {str(e)}")
            raise

    def run(self):
        """Main simulation loop."""
        print("DEBUG: Starting simulation loop...")
        frame_counter = 0
        while self.running:
            frame_counter += 1
            # Print every 60 frames (approximately once per second at 60 FPS)
            if frame_counter % 60 == 0:
                print(f"DEBUG: Simulation running - Frame: {self.frame_count}, Time: {self.simulation_time:.1f}s")

            try:
                self._handle_events()

                if not self.paused:
                    self._update()

                self._render()
                self._update_performance_metrics()

                # Control frame rate
                self.clock.tick(config.FPS)
                self.frame_count += 1
                self.simulation_time += 1.0 / config.FPS
            except Exception as e:
                print(f"DEBUG: Error in main loop: {str(e)}")
                import traceback
                traceback.print_exc()
                self.running = False

    def _handle_events(self):
        """Handle user input and window events."""
        try:
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    print("DEBUG: QUIT event detected")
                    self.running = False

                elif event.type == pygame.KEYDOWN:
                    print(f"DEBUG: Key pressed: {pygame.key.name(event.key)}")
                    if event.key == pygame.K_SPACE:
                        self.paused = not self.paused
                        print(f"DEBUG: Simulation {'paused' if self.paused else 'resumed'}")

                    elif event.key == pygame.K_r:
                        # Reset simulation
                        print("DEBUG: Resetting simulation...")
                        self._initialize_entities()
                        self.frame_count = 0
                        self.simulation_time = 0.0

                    elif event.key == pygame.K_d:
                        # Toggle AI debug info
                        config.DEBUG_CONFIG['show_ai_info'] = not config.DEBUG_CONFIG['show_ai_info']
                        print(f"DEBUG: AI debug info {'enabled' if config.DEBUG_CONFIG['show_ai_info'] else 'disabled'}")

                    elif event.key == pygame.K_v:
                        # Toggle dog's vision display
                        config.DEBUG_CONFIG['show_dog_vision'] = not config.DEBUG_CONFIG['show_dog_vision']
                        print(f"DEBUG: Dog vision display {'enabled' if config.DEBUG_CONFIG['show_dog_vision'] else 'disabled'}")

                    elif event.key == pygame.K_c:
                        # Toggle connections/target lines display
                        config.DEBUG_CONFIG['show_target_lines'] = not config.DEBUG_CONFIG['show_target_lines']
                        print(f"DEBUG: Target lines display {'enabled' if config.DEBUG_CONFIG['show_target_lines'] else 'disabled'}")

                    elif event.key == pygame.K_ESCAPE:
                        print("DEBUG: ESC key pressed - exiting")
                        self.running = False

                elif event.type == pygame.MOUSEBUTTONDOWN:
                    if event.button == 1:  # Left click
                        # Move target to mouse position
                        new_target = Vector2D(*event.pos)
                        self.target_position = new_target
                        if self.dog:
                            self.dog.target_position = new_target
                        print(f"DEBUG: Target position moved to {new_target}")
        except Exception as e:
            print(f"DEBUG: Error handling events: {str(e)}")
            import traceback
            traceback.print_exc()
            self.running = False

    def _update(self):
        """Update all simulation entities."""
        try:
            if self.dog and self.sheep_flock:
                # Verify we have valid sheep objects
                valid_sheep = []
                for sheep in self.sheep_flock:
                    if hasattr(sheep, 'position'):
                        valid_sheep.append(sheep)
                    else:
                        print("DEBUG: Found sheep without position attribute")

                # Pass only valid sheep to the dog
                self.dog.update(valid_sheep)

                # Get herding pressure from dog
                dog_pressure = self.dog.get_herding_pressure()

                # Update each sheep
                for i, sheep in enumerate(self.sheep_flock):
                    if hasattr(sheep, 'update'):
                        sheep.update(self.sheep_flock, self.dog.position, dog_pressure)
                    else:
                        print(f"DEBUG: Sheep {i} has no update method")
        except Exception as e:
            print(f"DEBUG: Error in update: {str(e)}")
            import traceback
            traceback.print_exc()
            self.running = False

    def _render(self):
        """Render all visual elements."""
        try:
            # Clear screen using the background color from COLORS
            self.screen.fill(config.COLORS['background'])

            # Draw target
            target_size = config.ENVIRONMENT_CONFIG['target_size']
            pygame.draw.circle(self.screen, config.COLORS['target'],
                            self.target_position.to_tuple(), target_size)
            pygame.draw.circle(self.screen, (255, 255, 255),
                            self.target_position.to_tuple(), target_size, 2)

            # Draw each sheep
            for sheep in self.sheep_flock:
                sheep.draw(self.screen)

            # Draw the shepherd dog if it exists
            if self.dog:
                self.dog.draw(self.screen)

            # Draw UI and debug information
            self._draw_ui()

            # Update display
            pygame.display.flip()
            if self.frame_count % 60 == 0:  # Print render info every second
                print(f"DEBUG: Rendered frame {self.frame_count}")
        except Exception as e:
            print(f"DEBUG: Error in render: {str(e)}")
            import traceback
            traceback.print_exc()
            self.running = False

    def _draw_ui(self):
        """Draw user interface and debug information."""
        y_offset = 10
        line_height = 25

        # Basic simulation info
        info_texts = [
            f"Frame: {self.frame_count}",
            f"Time: {self.simulation_time:.1f}s",
            f"FPS: {self.clock.get_fps():.1f}",
            f"Sheep: {len(self.sheep_flock)}",
        ]

        for text in info_texts:
            surface = self.font.render(text, True, (255, 255, 255))
            self.screen.blit(surface, (10, y_offset))
            y_offset += line_height

        # AI Debug information
        if config.DEBUG_CONFIG['show_ai_info']:
            y_offset += 10
            try:
                debug_info = self.dog.get_debug_info()

                debug_texts = [
                    "=== AI DEBUG INFO ===",
                    f"Visible Sheep: {debug_info['visible_sheep']}",
                    f"Selected Sheep: {debug_info['selected_sheep']}",
                    f"Herding Pressure: {debug_info['herding_pressure']}",
                    f"Performance: {debug_info['performance']:.3f}",
                    "",
                    "Feature Weights:",
                    f"  Distance: {debug_info['feature_weights'][0]}",
                    f"  Alignment: {debug_info['feature_weights'][1]}",
                    f"  Density: {debug_info['feature_weights'][2]}",
                    f"  Boundary: {debug_info['feature_weights'][3]}",
                ]

                for text in debug_texts:
                    if text:  # Skip empty strings
                        surface = self.small_font.render(text, True, (255, 255, 0))
                        self.screen.blit(surface, (10, y_offset))
                    y_offset += 20
            except Exception as e:
                print(f"DEBUG: Error getting debug info: {str(e)}")
                # Continue without debug info display

        # Control instructions
        controls = [
            "CONTROLS:",
            "SPACE - Pause/Resume",
            "R - Reset Simulation",
            "D - Toggle Debug Info",
            "V - Toggle Vision Radius",
            "C - Toggle Connections",
            "Left Click - Move Target",
            "ESC - Exit"
        ]

        y_offset_controls = config.SCREEN_HEIGHT - len(controls) * 18 - 10
        for control in controls:
            surface = self.small_font.render(control, True, (200, 200, 200))
            self.screen.blit(surface, (config.SCREEN_WIDTH - 200, y_offset_controls))
            y_offset_controls += 18

        # Performance metrics
        if config.DEBUG_CONFIG['show_ai_info']:
            try:
                self._draw_performance_metrics()
            except Exception as e:
                print(f"DEBUG: Error drawing performance metrics: {str(e)}")

    def _draw_performance_metrics(self):
        """Draw performance metrics and statistics."""
        metrics_x = config.SCREEN_WIDTH - 350
        metrics_y = 10

        try:
            avg_distance_to_target = self._calculate_average_distance_to_target()
            flock_compactness = self._calculate_flock_compactness()

            metrics_texts = [
                "=== PERFORMANCE METRICS ===",
                f"Avg Distance to Target: {avg_distance_to_target:.1f}",
                f"Flock Compactness: {flock_compactness:.1f}",
                f"Sheep in Target Zone: {self._count_sheep_in_target_zone()}",
            ]

            for text in metrics_texts:
                surface = self.small_font.render(text, True, (255, 255, 255))
                self.screen.blit(surface, (metrics_x, metrics_y))
                metrics_y += 20
        except Exception as e:
            print(f"DEBUG: Error calculating performance metrics: {str(e)}")
            import traceback
            traceback.print_exc()

    def _calculate_average_distance_to_target(self):
        """Calculate average distance of all sheep to target."""
        if not self.sheep_flock:
            return 0.0
        total_distance = sum(sheep.position.distance_to(self.target_position)
                           for sheep in self.sheep_flock)
        return total_distance / len(self.sheep_flock)

    def _calculate_flock_compactness(self):
        """Calculate how compact the flock is (lower value means more compact)."""
        if len(self.sheep_flock) < 2:
            return 0.0
        center = Vector2D()
        for sheep in self.sheep_flock:
            center += sheep.position
        center = center / len(self.sheep_flock)
        total_distance = sum(sheep.position.distance_to(center)
                           for sheep in self.sheep_flock)
        return total_distance / len(self.sheep_flock)

    def _count_sheep_in_target_zone(self):
        """Count sheep within target zone."""
        target_zone_radius = config.ENVIRONMENT_CONFIG['target_size'] * 3
        count = 0
        for sheep in self.sheep_flock:
            if sheep.position.distance_to(self.target_position) < target_zone_radius:
                count += 1
        return count

    def _update_performance_metrics(self):
        """Update performance tracking."""
        current_fps = self.clock.get_fps()
        self.fps_history.append(current_fps)
        if len(self.fps_history) > 60:
            self.fps_history.pop(0)
        self.performance_metrics = {
            'avg_fps': sum(self.fps_history) / len(self.fps_history) if self.fps_history else 0,
            'avg_distance_to_target': self._calculate_average_distance_to_target(),
            'flock_compactness': self._calculate_flock_compactness(),
            'sheep_in_target': self._count_sheep_in_target_zone(),
            'simulation_time': self.simulation_time
        }

    def cleanup(self):
        """Clean up resources before exit."""
        print("DEBUG: Starting cleanup...")
        pygame.quit()
        print("DEBUG: Pygame quit called")
        sys.exit()

def main():
    """Main entry point for the simulation."""
    try:
        print("DEBUG: Starting main function...")
        env = Environment()
        print("DEBUG: Environment created")
        env.run()
        print("DEBUG: Simulation run completed (this should not be printed unless running=False)")
    except KeyboardInterrupt:
        print("\nDEBUG: Simulation interrupted by user")
    except Exception as e:
        print(f"\nDEBUG: An error occurred: {str(e)}")
        print(f"DEBUG: Error type: {type(e).__name__}")
        import traceback
        print("DEBUG: Stack trace:")
        traceback.print_exc()
    finally:
        print("DEBUG: Entering finally block")
        pygame.quit()
        print("DEBUG: Pygame quit called")
        sys.exit()

if __name__ == "__main__":
    main()