import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from matplotlib.colors import ListedColormap
import os
from typing import Dict, List, Tuple, Optional, Union
from .maze_wrapper import MazeWrapper
from .maze_generation import MazeGenerator

def get_agent_offset_x(agent_id: str, agent_index: int) -> float:
    """
    Calculate horizontal offset for agent visualization based on agent ID.
    
    Args:
        agent_id: Unique identifier for the agent (e.g., 'agent_0', 'agent_1')
        agent_index: Index of the agent in the iteration (0, 1, 2, ...)
    
    Returns:
        Horizontal offset in pixels (converted to matplotlib coordinates)
    """
    # Extract numeric ID from agent_id if possible, otherwise use agent_index
    try:
        if 'agent_' in agent_id:
            numeric_id = int(agent_id.replace('agent_', ''))
        else:
            # Try to extract number from end of string
            import re
            match = re.search(r'(\d+)$', agent_id)
            numeric_id = int(match.group(1)) if match else agent_index
    except (ValueError, AttributeError):
        numeric_id = agent_index
    
    # Create offset pattern: alternate between positive and negative offsets
    # Agent 0: 0, Agent 1: +10px, Agent 2: -10px, Agent 3: +20px, Agent 4: -20px, etc.
    if numeric_id == 0:
        offset_pixels = 0
    elif numeric_id % 2 == 1:  # Odd agents get positive offset
        offset_pixels = ((numeric_id + 1) // 2) * 10
    else:  # Even agents get negative offset
        offset_pixels = -(numeric_id // 2) * 10
    
    # Convert pixels to matplotlib coordinate units
    # Assume each cell is approximately 50 pixels wide for conversion
    return offset_pixels / 50.0


class MazeAggregator:
    """
    Aggregator class that manages multiple agents (MazeWrapper instances) in the same maze.
    
    Provides comprehensive visualization and tracking of multiple agents navigating 
    the same maze simultaneously, with individual agent statistics and combined views.
    """
    
    def __init__(self, maze_generator: MazeGenerator):
        """
        Initialize the maze aggregator with a base maze.
        
        Args:
            maze_generator: MazeGenerator instance that all agents will use
        """
        self.maze_generator = maze_generator
        self.maze = maze_generator.maze
        self.height, self.width = self.maze.shape
        
        # Dictionary to store agents by their unique IDs
        self.agents: Dict[str, MazeWrapper] = {}
        
        # Agent color mapping for visualization
        self.agent_colors = [
            '#FF0000',  # Red
            '#0000FF',  # Blue
            '#00FF00',  # Green
            '#FF00FF',  # Magenta
            '#FFA500',  # Orange
            '#800080',  # Purple
            '#00FFFF',  # Cyan
            '#FF69B4',  # Hot Pink
            '#32CD32',  # Lime Green
            '#FFD700',  # Gold
            '#FF1493',  # Deep Pink
            '#4169E1',  # Royal Blue
            '#FF6347',  # Tomato
            '#9370DB',  # Medium Purple
            '#20B2AA',  # Light Sea Green
        ]
        
        # Base maze visualization colors
        self.base_colors = {
            'W': '#000000',  # Black for walls
            'O': '#FFFFFF',  # White for open paths  
            'E': '#FFD700',  # Yellow for exit (matching image)
            'X': '#808080',  # Gray for outer frame
        }
        
        # Track global statistics
        self.total_moves = 0
        self.total_failed_moves = 0
        
    def add_agent(self, agent_id: str, start_position: Optional[Tuple[int, int]] = None,
                  maze_wrapper: Optional[MazeWrapper] = None, avoid_occupied: bool = True) -> bool:
        """
        Add a new agent to the aggregator.
        
        Args:
            agent_id: Unique identifier for the agent
            start_position: Starting position for the agent (optional)
            maze_wrapper: Existing MazeWrapper instance (optional)
            avoid_occupied: If True, avoid starting positions already occupied by other agents
            
        Returns:
            True if agent was added successfully, False if agent_id already exists
        """
        if agent_id in self.agents:
            print(f"Warning: Agent '{agent_id}' already exists. Use update_agent() to modify.")
            return False
        
        if maze_wrapper is not None:
            # Verify the maze wrapper uses the same maze
            if not np.array_equal(maze_wrapper.maze, self.maze):
                print(f"Warning: Agent '{agent_id}' uses a different maze. Creating new wrapper.")
                maze_wrapper = MazeWrapper(self.maze_generator, start_position)
        else:
            # Handle start position logic
            if start_position is not None and avoid_occupied:
                # Check if the requested position is already occupied
                occupied_positions = {agent.get_agent_position() for agent in self.agents.values()}
                if start_position in occupied_positions:
                    print(f"Warning: Position {start_position} is occupied. Finding alternative...")
                    start_position = self._find_alternative_start_position(start_position, occupied_positions)
                    if start_position is None:
                        print(f"Error: Could not find alternative starting position for agent '{agent_id}'")
                        return False
            elif start_position is None and avoid_occupied:
                # Find an unoccupied starting position
                occupied_positions = {agent.get_agent_position() for agent in self.agents.values()}
                start_position = self._find_unoccupied_start_position(occupied_positions)
                if start_position is None:
                    print(f"Error: Could not find unoccupied starting position for agent '{agent_id}'")
                    return False
            
            # Create new maze wrapper
            maze_wrapper = MazeWrapper(self.maze_generator, start_position)
        
        self.agents[agent_id] = maze_wrapper
        print(f"Agent '{agent_id}' added at position {maze_wrapper.get_agent_position()}")
        return True
    
    def remove_agent(self, agent_id: str) -> bool:
        """
        Remove an agent from the aggregator.
        
        Args:
            agent_id: ID of the agent to remove
            
        Returns:
            True if agent was removed, False if agent didn't exist
        """
        if agent_id in self.agents:
            del self.agents[agent_id]
            print(f"Agent '{agent_id}' removed")
            return True
        else:
            print(f"Agent '{agent_id}' not found")
            return False
    
    def add_multiple_agents(self, agent_configs: List[Union[str, Tuple[str, Tuple[int, int]]]], 
                           avoid_occupied: bool = True, distribute_evenly: bool = False) -> Dict[str, bool]:
        """
        Add multiple agents with different starting positions.
        
        Args:
            agent_configs: List of agent configurations. Each can be:
                          - str: agent_id (auto-find position)
                          - Tuple[str, Tuple[int, int]]: (agent_id, start_position)
            avoid_occupied: If True, avoid starting positions already occupied by other agents
            distribute_evenly: If True, try to distribute agents evenly across the maze
            
        Returns:
            Dictionary mapping agent_id to success status
        """
        results = {}
        
        if distribute_evenly and len(agent_configs) > 1:
            # Get all open positions and distribute agents evenly
            open_positions = self._find_all_open_positions()
            if len(open_positions) < len(agent_configs):
                print(f"Warning: Only {len(open_positions)} open positions available for {len(agent_configs)} agents")
            
            # Distribute positions evenly
            step = max(1, len(open_positions) // len(agent_configs))
            distributed_positions = [open_positions[i * step] for i in range(min(len(agent_configs), len(open_positions)))]
            
            # Add agents with distributed positions
            for i, config in enumerate(agent_configs):
                if isinstance(config, str):
                    agent_id = config
                    start_pos = distributed_positions[i] if i < len(distributed_positions) else None
                else:
                    agent_id, requested_pos = config
                    # Use requested position if provided, otherwise use distributed position
                    start_pos = requested_pos if requested_pos is not None else (distributed_positions[i] if i < len(distributed_positions) else None)
                
                results[agent_id] = self.add_agent(agent_id, start_pos, avoid_occupied=avoid_occupied)
        else:
            # Add agents with their specified or auto-found positions
            for config in agent_configs:
                if isinstance(config, str):
                    agent_id = config
                    start_pos = None
                else:
                    agent_id, start_pos = config
                
                results[agent_id] = self.add_agent(agent_id, start_pos, avoid_occupied=avoid_occupied)
        
        return results
    
    def get_available_start_positions(self, exclude_occupied: bool = True) -> List[Tuple[int, int]]:
        """
        Get list of available starting positions in the maze.
        
        Args:
            exclude_occupied: If True, exclude positions currently occupied by agents
            
        Returns:
            List of available (row, col) positions
        """
        open_positions = self._find_all_open_positions()
        
        if exclude_occupied:
            occupied_positions = {agent.get_agent_position() for agent in self.agents.values()}
            return [pos for pos in open_positions if pos not in occupied_positions]
        
        return open_positions
    
    def get_agent(self, agent_id: str) -> Optional[MazeWrapper]:
        """
        Get a specific agent by ID.
        
        Args:
            agent_id: ID of the agent to retrieve
            
        Returns:
            MazeWrapper instance or None if not found
        """
        return self.agents.get(agent_id)
    
    def get_agent_ids(self) -> List[str]:
        """
        Get list of all agent IDs.
        
        Returns:
            List of agent ID strings
        """
        return list(self.agents.keys())
    
    def move_agent(self, agent_id: str, direction: str) -> bool:
        """
        Move a specific agent in the given direction.
        
        Args:
            agent_id: ID of the agent to move
            direction: Direction to move ('up', 'down', 'left', 'right')
            
        Returns:
            True if move was successful, False otherwise
        """
        if agent_id not in self.agents:
            print(f"Agent '{agent_id}' not found")
            return False
        
        success = self.agents[agent_id].try_move(direction)
        if success:
            self.total_moves += 1
        else:
            self.total_failed_moves += 1
        
        return success
    
    def move_all_agents(self, directions: Dict[str, str]) -> Dict[str, bool]:
        """
        Move multiple agents simultaneously.
        
        Args:
            directions: Dictionary mapping agent_id to direction
            
        Returns:
            Dictionary mapping agent_id to success status
        """
        results = {}
        for agent_id, direction in directions.items():
            results[agent_id] = self.move_agent(agent_id, direction)
        return results
    
    def reset_agent(self, agent_id: str, position: Optional[Tuple[int, int]] = None) -> bool:
        """
        Reset a specific agent to a position.
        
        Args:
            agent_id: ID of the agent to reset
            position: Position to reset to (None for original start)
            
        Returns:
            True if reset was successful, False if agent not found
        """
        if agent_id not in self.agents:
            print(f"Agent '{agent_id}' not found")
            return False
        
        self.agents[agent_id].reset_agent_position(position)
        print(f"Agent '{agent_id}' reset to position {self.agents[agent_id].get_agent_position()}")
        return True
    
    def reset_all_agents(self) -> None:
        """Reset all agents to their starting positions."""
        for agent_id in self.agents:
            self.reset_agent(agent_id)
    
    def get_agent_positions(self) -> Dict[str, Tuple[int, int]]:
        """
        Get current positions of all agents.
        
        Returns:
            Dictionary mapping agent_id to (row, col) position
        """
        return {agent_id: agent.get_agent_position() for agent_id, agent in self.agents.items()}
    
    def get_agents_at_exit(self) -> List[str]:
        """
        Get list of agents currently at the exit.
        
        Returns:
            List of agent IDs at the exit
        """
        return [agent_id for agent_id, agent in self.agents.items() if agent.is_at_exit()]
    
    def get_global_statistics(self) -> Dict:
        """
        Get aggregated statistics for all agents.
        
        Returns:
            Dictionary containing global statistics
        """
        if not self.agents:
            return {
                'total_agents': 0,
                'agents_at_exit': 0,
                'total_moves': 0,
                'total_failed_moves': 0,
                'average_moves_per_agent': 0,
                'success_rate': 0
            }
        
        total_agent_moves = sum(agent.get_move_count() for agent in self.agents.values())
        total_agent_failed = sum(agent.get_failed_move_count() for agent in self.agents.values())
        agents_at_exit = len(self.get_agents_at_exit())
        
        total_attempts = total_agent_moves + total_agent_failed
        success_rate = (total_agent_moves / total_attempts * 100) if total_attempts > 0 else 0
        
        return {
            'total_agents': len(self.agents),
            'agents_at_exit': agents_at_exit,
            'total_moves': total_agent_moves,
            'total_failed_moves': total_agent_failed,
            'average_moves_per_agent': total_agent_moves / len(self.agents),
            'success_rate': success_rate,
            'completion_rate': (agents_at_exit / len(self.agents) * 100)
        }
    
    def get_agent_statistics(self) -> Dict[str, Dict]:
        """
        Get individual statistics for each agent.
        
        Returns:
            Dictionary mapping agent_id to individual statistics
        """
        stats = {}
        for agent_id, agent in self.agents.items():
            stats[agent_id] = {
                'position': agent.get_agent_position(),
                'moves': agent.get_move_count(),
                'failed_moves': agent.get_failed_move_count(),
                'at_exit': agent.is_at_exit(),
                'possible_moves': agent.get_possible_moves(),
                'path_length': len(agent.move_history),
                'unique_positions': len(set(agent.move_history))
            }
        return stats
    
    def print_status(self) -> None:
        """Print comprehensive status of all agents and global statistics."""
        print("\n" + "="*80)
        print("MULTI-AGENT MAZE STATUS")
        print("="*80)
        
        # Global statistics
        global_stats = self.get_global_statistics()
        print(f"Total Agents: {global_stats['total_agents']}")
        print(f"Agents at Exit: {global_stats['agents_at_exit']}")
        print(f"Completion Rate: {global_stats['completion_rate']:.1f}%")
        print(f"Total Moves: {global_stats['total_moves']}")
        print(f"Total Failed Moves: {global_stats['total_failed_moves']}")
        print(f"Average Moves per Agent: {global_stats['average_moves_per_agent']:.1f}")
        print(f"Overall Success Rate: {global_stats['success_rate']:.1f}%")
        
        # Individual agent statistics
        print("\n" + "-"*80)
        print("INDIVIDUAL AGENT STATUS")
        print("-"*80)
        
        agent_stats = self.get_agent_statistics()
        for agent_id, stats in agent_stats.items():
            status = "🎉 AT EXIT" if stats['at_exit'] else "🏃 NAVIGATING"
            print(f"\nAgent '{agent_id}' {status}")
            print(f"  Position: {stats['position']}")
            print(f"  Moves: {stats['moves']}, Failed: {stats['failed_moves']}")
            print(f"  Path Length: {stats['path_length']}, Unique Positions: {stats['unique_positions']}")
            print(f"  Possible Moves: {stats['possible_moves']}")
        
        print("="*80)
    
    def get_combined_maze_visualization(self, show_paths: bool = True, show_agents: bool = True,
                                      show_starts: bool = True, show_failed_moves: bool = False,
                                      agent_labels: bool = True) -> plt.Figure:
        """
        Create a combined visualization showing all agents in the same maze.
        
        Args:
            show_paths: Whether to show agent paths
            show_agents: Whether to show current agent positions
            show_starts: Whether to show starting positions
            show_failed_moves: Whether to show failed move attempts
            agent_labels: Whether to show agent ID labels
            
        Returns:
            matplotlib Figure object
        """
        # Create figure with appropriate size
        fig, ax = plt.subplots(figsize=(15, 12))
        
        # Start with base maze visualization
        visual_maze = self.maze.copy().astype(object)
        
        # Create a color map for the visualization
        height, width = self.maze.shape
        
        # Draw base maze
        for y in range(height):
            for x in range(width):
                cell = self.maze[y, x]
                color = self.base_colors.get(cell, '#FFFFFF')
                rect = plt.Rectangle((x, height-y-1), 1, 1, 
                                   facecolor=color, edgecolor='black', linewidth=0.5)
                ax.add_patch(rect)
        
        # Draw agent paths
        if show_paths:
            for i, (agent_id, agent) in enumerate(self.agents.items()):
                color = self.agent_colors[i % len(self.agent_colors)]
                path_color = self._lighten_color(color, 0.6)  # Lighter version for path

                # Apply the same horizontal offset used for agent positions to path markers
                # This reduces overlap when multiple agents traverse the same cell
                offset_x = get_agent_offset_x(agent_id, i)

                # Draw path (excluding current position)
                for j, (row, col) in enumerate(agent.move_history[:-1]):
                    circle = plt.Circle((col + 0.5 + offset_x, height - row - 0.5), 0.15,
                                      facecolor=path_color, edgecolor=color, linewidth=1,
                                      alpha=0.7)
                    ax.add_patch(circle)
        
        # Draw starting positions
        if show_starts:
            for i, (agent_id, agent) in enumerate(self.agents.items()):
                if agent.move_history:
                    start_row, start_col = agent.move_history[0]
                    color = self.agent_colors[i % len(self.agent_colors)]
                    
                    # Draw start marker as a square
                    square = plt.Rectangle((start_col + 0.2, height - start_row - 0.8), 0.6, 0.6,
                                         facecolor='white', edgecolor=color, linewidth=3)
                    ax.add_patch(square)
                    
                    # Add 'S' label
                    ax.text(start_col + 0.5, height - start_row - 0.5, 'S', 
                           ha='center', va='center', fontsize=8, fontweight='bold', color=color)
        
        # Draw failed moves
        if show_failed_moves:
            for i, (agent_id, agent) in enumerate(self.agents.items()):
                color = self.agent_colors[i % len(self.agent_colors)]
                failed_color = self._darken_color(color, 0.3)
                
                for fail_row, fail_col in agent.failed_moves:
                    if (0 <= fail_row < height and 0 <= fail_col < width):
                        x_mark = ax.plot([fail_col + 0.2, fail_col + 0.8], 
                                       [height - fail_row - 0.2, height - fail_row - 0.8], 
                                       color=failed_color, linewidth=2)[0]
                        ax.plot([fail_col + 0.2, fail_col + 0.8], 
                               [height - fail_row - 0.8, height - fail_row - 0.2], 
                               color=failed_color, linewidth=2)
        
        # Draw current agent positions (on top)
        if show_agents:
            for i, (agent_id, agent) in enumerate(self.agents.items()):
                row, col = agent.get_agent_position()
                color = self.agent_colors[i % len(self.agent_colors)]
                
                # Calculate horizontal offset for this agent
                offset_x = get_agent_offset_x(agent_id, i)
                
                # Draw agent as a circle
                circle = plt.Circle((col + 0.5 + offset_x, height - row - 0.5), 0.3, 
                                  facecolor=color, edgecolor='black', linewidth=2)
                ax.add_patch(circle)
                
                # Add agent label
                if agent_labels:
                    label = agent_id[:3] if len(agent_id) > 3 else agent_id
                    ax.text(col + 0.5 + offset_x, height - row - 0.5, label, 
                           ha='center', va='center', fontsize=8, fontweight='bold', color='white')
        
        # Set axis properties
        ax.set_xlim(0, width)
        ax.set_ylim(0, height)
        ax.set_aspect('equal')
        ax.set_xticks(np.arange(0, width + 1, 1))
        ax.set_yticks(np.arange(0, height + 1, 1))
        ax.grid(True, alpha=0.3)
        
        # Create title
        agents_at_exit = len(self.get_agents_at_exit())
        title = f'Multi-Agent Maze Navigation ({len(self.agents)} agents, {agents_at_exit} at exit)'
        ax.set_title(title, fontsize=16, fontweight='bold')
        
        # Add statistics panel
        global_stats = self.get_global_statistics()
        stats_text = f"""Total Agents: {global_stats['total_agents']}
At Exit: {global_stats['agents_at_exit']} ({global_stats['completion_rate']:.1f}%)
Total Moves: {global_stats['total_moves']}
Failed Moves: {global_stats['total_failed_moves']}
Success Rate: {global_stats['success_rate']:.1f}%
Avg Moves/Agent: {global_stats['average_moves_per_agent']:.1f}"""
        
        ax.text(0.02, 0.98, stats_text, transform=ax.transAxes, 
                fontsize=10, verticalalignment='top', 
                bbox=dict(boxstyle='round,pad=0.5', facecolor='white', alpha=0.9, edgecolor='black'),
                family='monospace')
        
        # Create legend
        legend_elements = [
            plt.Rectangle((0, 0), 1, 1, facecolor='#000000', label='Wall'),
            plt.Rectangle((0, 0), 1, 1, facecolor='#FFFFFF', label='Open Path'),
            plt.Rectangle((0, 0), 1, 1, facecolor='#FFD700', label='Exit'),
            plt.Rectangle((0, 0), 1, 1, facecolor='#808080', label='Frame'),
        ]
        
        # Add agent legend
        for i, (agent_id, agent) in enumerate(self.agents.items()):
            color = self.agent_colors[i % len(self.agent_colors)]
            status = " (EXIT)" if agent.is_at_exit() else ""
            legend_elements.append(
                plt.Line2D([0], [0], marker='o', color='w', markerfacecolor=color, 
                          markersize=10, label=f'Agent {agent_id}{status}')
            )
        
        ax.legend(handles=legend_elements, loc='upper right', bbox_to_anchor=(0.98, 0.98))
        
        plt.tight_layout()
        return fig
    
    def save_combined_visualization(self, filename: str, show_paths: bool = True, 
                                   show_agents: bool = True, show_starts: bool = True, 
                                   show_failed_moves: bool = False, agent_labels: bool = True, 
                                   dpi: int = 150) -> None:
        """
        Save combined visualization to file.
        
        Args:
            filename: Output filename
            show_paths: Whether to show agent paths
            show_agents: Whether to show current agent positions
            show_starts: Whether to show starting positions
            show_failed_moves: Whether to show failed move attempts
            agent_labels: Whether to show agent ID labels
            dpi: Image resolution
        """
        # Create output directory
        output_dir = "maze_aggregator_outputs"
        os.makedirs(output_dir, exist_ok=True)
        full_path = os.path.join(output_dir, filename)
        
        fig = self.get_combined_maze_visualization(show_paths, show_agents, show_starts, 
                                                  show_failed_moves, agent_labels)
        fig.savefig(full_path, dpi=dpi, bbox_inches='tight')
        plt.close(fig)
        print(f"Multi-agent visualization saved as {full_path}")
    
    def show_combined_interactive(self, show_paths: bool = True, show_agents: bool = True, 
                                 show_starts: bool = True, show_failed_moves: bool = False, 
                                 agent_labels: bool = True) -> None:
        """
        Show combined visualization in interactive matplotlib window.
        
        Args:
            show_paths: Whether to show agent paths
            show_agents: Whether to show current agent positions
            show_starts: Whether to show starting positions
            show_failed_moves: Whether to show failed move attempts
            agent_labels: Whether to show agent ID labels
        """
        fig = self.get_combined_maze_visualization(show_paths, show_agents, show_starts, 
                                                   show_failed_moves, agent_labels)
        # Determine safe mode via config service (single source of truth)
        from orchestrator_maze_implementation.config.config_service import config as _cfg
        safe_mode = _cfg.get_visualization_safe_mode()
        if safe_mode:
            try:
                plt.show()
            except Exception:
                output_dir = "maze_aggregator_outputs"
                os.makedirs(output_dir, exist_ok=True)
                fallback_path = os.path.join(output_dir, "interactive_fallback.png")
                fig.savefig(fallback_path, dpi=120, bbox_inches='tight')
                print(f"⚠️  Interactive display unavailable. Saved to {fallback_path}")
            finally:
                plt.close(fig)
        else:
            # Original behavior
            plt.show()
            plt.close(fig)
    
    def _lighten_color(self, color: str, factor: float) -> str:
        """
        Lighten a hex color by a given factor.
        
        Args:
            color: Hex color string
            factor: Factor to lighten (0-1, where 1 is white)
            
        Returns:
            Lightened hex color string
        """
        import matplotlib.colors as mcolors
        rgb = mcolors.hex2color(color)
        lightened = [rgb[i] + (1 - rgb[i]) * factor for i in range(3)]
        return mcolors.rgb2hex(lightened)
    
    def _darken_color(self, color: str, factor: float) -> str:
        """
        Darken a hex color by a given factor.
        
        Args:
            color: Hex color string
            factor: Factor to darken (0-1, where 1 is black)
            
        Returns:
            Darkened hex color string
        """
        import matplotlib.colors as mcolors
        rgb = mcolors.hex2color(color)
        darkened = [rgb[i] * (1 - factor) for i in range(3)]
        return mcolors.rgb2hex(darkened)
    
    def _find_all_open_positions(self) -> List[Tuple[int, int]]:
        """
        Find all open positions in the maze where agents can be placed.
        
        Returns:
            List of (row, col) positions that are open ('O') or exit ('E')
        """
        open_positions = []
        for row in range(self.height):
            for col in range(self.width):
                if self.maze[row, col] in ['O', 'E']:
                    open_positions.append((row, col))
        return open_positions
    
    def _find_unoccupied_start_position(self, occupied_positions: set) -> Optional[Tuple[int, int]]:
        """
        Find an unoccupied starting position in the maze.
        
        Args:
            occupied_positions: Set of positions already occupied by other agents
            
        Returns:
            Unoccupied position or None if no positions available
        """
        open_positions = self._find_all_open_positions()
        
        # Filter out occupied positions
        available_positions = [pos for pos in open_positions if pos not in occupied_positions]
        
        if not available_positions:
            return None
        
        # Return the first available position (could be randomized if desired)
        return available_positions[0]
    
    def _find_alternative_start_position(self, requested_position: Tuple[int, int], 
                                       occupied_positions: set) -> Optional[Tuple[int, int]]:
        """
        Find an alternative starting position near the requested position.
        
        Args:
            requested_position: The originally requested position
            occupied_positions: Set of positions already occupied by other agents
            
        Returns:
            Alternative position or None if no suitable position found
        """
        req_row, req_col = requested_position
        
        # Try positions in expanding radius around the requested position
        for radius in range(1, max(self.height, self.width)):
            candidates = []
            
            # Generate positions in a square around the requested position
            for dr in range(-radius, radius + 1):
                for dc in range(-radius, radius + 1):
                    if abs(dr) == radius or abs(dc) == radius:  # Only edge of square
                        new_row, new_col = req_row + dr, req_col + dc
                        new_pos = (new_row, new_col)
                        
                        # Check if position is valid and unoccupied
                        if (0 <= new_row < self.height and 0 <= new_col < self.width and
                            self.maze[new_row, new_col] in ['O', 'E'] and
                            new_pos not in occupied_positions):
                            candidates.append(new_pos)
            
            if candidates:
                # Return the closest candidate (could be randomized if desired)
                return min(candidates, key=lambda pos: abs(pos[0] - req_row) + abs(pos[1] - req_col))
        
        # If no position found nearby, try any unoccupied position
        return self._find_unoccupied_start_position(occupied_positions)


def main():
    """
    Test function to demonstrate MazeAggregator capabilities.
    
    Creates multiple agents in the same maze and shows their combined progress.
    """
    print("="*80)
    print("MULTI-AGENT MAZE AGGREGATOR TEST")
    print("="*80)
    
    # Create a maze
    print("Creating a 15x15 maze...")
    maze_gen = MazeGenerator(maze_size=15)
    aggregator = MazeAggregator(maze_gen)
    
    # Add multiple agents with different starting locations
    print("\n1. Adding multiple agents with different starting locations...")
    
    # First, show available positions
    available_positions = aggregator.get_available_start_positions()
    print(f"   Available starting positions: {len(available_positions)} found")
    print(f"   First few positions: {available_positions[:8]}")
    
    # Method 1: Add agents with specific positions
    print("\n   Method 1: Adding agents with specific positions...")
    specific_agents = [
        ("Alpha", (1, 1)),  # Try specific position
        ("Beta", (1, 3)),   # Try another specific position
        ("Gamma", (3, 1)),  # And another
    ]
    
    for agent_id, start_pos in specific_agents:
        success = aggregator.add_agent(agent_id, start_pos)
        if success:
            print(f"   ✅ Added agent '{agent_id}' at {start_pos}")
    
    # Method 2: Add agents with automatic positioning (avoid occupied)
    print("\n   Method 2: Adding agents with automatic positioning...")
    aggregator.add_agent("Delta")  # Auto-find position, avoid occupied
    print(f"   ✅ Added agent 'Delta' (auto-positioned)")
    
    # Method 3: Use the bulk add method with distribution
    print("\n   Method 3: Adding more agents with even distribution...")
    bulk_agents = ["Echo", "Foxtrot", "Golf"]
    results = aggregator.add_multiple_agents(bulk_agents, distribute_evenly=True)
    for agent_id, success in results.items():
        if success:
            pos = aggregator.get_agent(agent_id).get_agent_position()
            print(f"   ✅ Added agent '{agent_id}' at {pos} (distributed)")
    
    # Show final positions
    print(f"\n   Final agent positions:")
    positions = aggregator.get_agent_positions()
    for agent_id, pos in positions.items():
        print(f"     {agent_id}: {pos}")
    
    # Show initial status
    print("\n2. Initial status:")
    aggregator.print_status()
    
    # Save initial state
    print("\n3. Saving initial visualization...")
    aggregator.save_combined_visualization("initial_multi_agent.png")
    
    # Simulate some moves
    print("\n4. Simulating agent movements...")
    import random
    
    for round_num in range(1, 21):  # 20 rounds of moves
        print(f"\nRound {round_num}:")
        
        # Get moves for each agent
        moves = {}
        for agent_id in aggregator.get_agent_ids():
            agent = aggregator.get_agent(agent_id)
            if agent and not agent.is_at_exit():
                possible_moves = agent.get_possible_moves()
                if possible_moves:
                    moves[agent_id] = random.choice(possible_moves)
        
        # Execute moves
        if moves:
            results = aggregator.move_all_agents(moves)
            for agent_id, success in results.items():
                status = "✅" if success else "❌"
                direction = moves[agent_id]
                position = aggregator.get_agent(agent_id).get_agent_position()
                print(f"   {status} Agent {agent_id}: {direction} -> {position}")
        
        # Check for agents at exit
        agents_at_exit = aggregator.get_agents_at_exit()
        if agents_at_exit:
            print(f"   🎉 Agents at exit: {', '.join(agents_at_exit)}")
        
        # Save state every 5 rounds
        if round_num % 5 == 0:
            filename = f"multi_agent_round_{round_num}.png"
            aggregator.save_combined_visualization(filename)
            print(f"   💾 Saved state as {filename}")
        
        # Stop if all agents reached exit
        if len(agents_at_exit) == len(aggregator.get_agent_ids()):
            print(f"\n🎉 ALL AGENTS REACHED THE EXIT IN {round_num} ROUNDS! 🎉")
            break
    
    # Final status
    print("\n5. Final status:")
    aggregator.print_status()
    
    # Save final state
    print("\n6. Saving final visualizations...")
    aggregator.save_combined_visualization("final_multi_agent.png")
    aggregator.save_combined_visualization("final_multi_agent_no_paths.png", show_paths=False)
    aggregator.save_combined_visualization("final_multi_agent_detailed.png", 
                                          show_failed_moves=True)
    
    # Show individual agent details
    print("\n7. Individual agent details:")
    agent_stats = aggregator.get_agent_statistics()
    for agent_id, stats in agent_stats.items():
        print(f"\nAgent '{agent_id}':")
        print(f"   Final Position: {stats['position']}")
        print(f"   Total Moves: {stats['moves']}")
        print(f"   Failed Moves: {stats['failed_moves']}")
        print(f"   Success Rate: {(stats['moves']/(stats['moves']+stats['failed_moves'])*100):.1f}%")
        print(f"   Path Efficiency: {(stats['unique_positions']/stats['path_length']*100):.1f}%")
        print(f"   Status: {'🎉 REACHED EXIT' if stats['at_exit'] else '🏃 Still navigating'}")
    
    # Try interactive visualization
    try:
        print("\n8. Showing interactive visualization...")
        print("   (Close the window to continue)")
        aggregator.show_combined_interactive()
    except Exception as e:
        print(f"   Could not show interactive plot: {e}")
        print("   This is normal if running in a non-interactive environment")
    
    print("\n" + "="*80)
    print("MULTI-AGENT AGGREGATOR TEST COMPLETE")
    print("="*80)
    print("Generated files in maze_aggregator_outputs/:")
    print("- initial_multi_agent.png")
    print("- multi_agent_round_*.png (every 5 rounds)")
    print("- final_multi_agent.png")
    print("- final_multi_agent_no_paths.png")
    print("- final_multi_agent_detailed.png (with failed moves)")
    print("="*80)


if __name__ == "__main__":
    main()
