import numpy as np
import matplotlib.pyplot as plt
import os
from typing import Tuple, Optional, List, Union, Set, Dict, Any
from copy import deepcopy
from orchestrator_maze_implementation.utils.maze_utils.maze_generation import MazeGenerator


class MazeWrapper:
    """
    Wrapper class for MazeGenerator that allows an agent to navigate through the maze.
    The agent can also mark positions as dead ends for analysis and visualization.
    
    The agent is represented by its current position and can move through open paths ('O')
    and reach the exit ('E'), but cannot move through walls ('W') or outer frame ('X').
    
    Features state history tracking for analyzing agent behavior patterns.
    """
    
    def __init__(self, maze_generator: MazeGenerator, start_position: Optional[Tuple[int, int]] = None):
        """
        Initialize the maze wrapper with a maze generator.
        
        Args:
            maze_generator: MazeGenerator instance containing the maze
            start_position: Optional starting position (row, col). If None, finds first open cell.
        """
        self.maze_generator = maze_generator
        self.maze = maze_generator.maze  # numpy array
        self.height, self.width = self.maze.shape
        
        # Set agent starting position
        if start_position is not None:
            if self._is_valid_position(start_position):
                self.agent_position = start_position
            else:
                raise ValueError(f"Invalid start position {start_position}")
        else:
            self.agent_position = self._find_starting_position()
        
        # Track movement history for debugging/analysis
        self.move_history: List[Tuple[int, int]] = [self.agent_position]
        
        # Track failed move attempts (when agent tries to move into walls)
        self.failed_moves: List[Tuple[int, int]] = []
        
        # Track marked dead ends (agent action)
        self.marked_dead_ends: Set[Tuple[int, int]] = set()
        
        # Track marking history for analysis
        self.marking_history: List[Tuple[str, Tuple[int, int]]] = []  # ('mark'/'unmark', position)
        
        # State history tracking for behavioral analysis
        self.state_history: List[Dict[str, Any]] = []
        
        # Color mapping for maze visualization
        self.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
            'A': '#00FFFF',  # Cyan for agent (matching image)
            'P': '#FFB6C1',  # Light pink for visited path (matching image)
            'S': '#0066FF',  # Blue for starting point
            'F': '#FF6600',  # Orange for failed move attempts
            'D': '#FF0000',  # Red for marked dead ends (matching image)
            'DE': '#FF8C00'  # Dark orange for exit marked as dead end
        }
        
        # Save initial state
        self.turn_complete()
    
    @classmethod
    def from_uuid(cls, maze_uuid: str, mazes_directory: str = "mazes", 
                  start_position: Optional[Tuple[int, int]] = None) -> 'MazeWrapper':
        """
        Create a MazeWrapper instance by loading a maze from UUID.
        
        Args:
            maze_uuid: UUID string of the maze to load
            mazes_directory: Directory containing maze files (default: "mazes")
            start_position: Optional starting position for the agent
            
        Returns:
            MazeWrapper instance with loaded maze
            
        Raises:
            FileNotFoundError: If maze with given UUID is not found
            ValueError: If maze data is invalid
        """
        # Load maze generator from UUID
        maze_generator = MazeGenerator.find_maze_by_uuid(maze_uuid, mazes_directory)
        
        if maze_generator is None:
            raise FileNotFoundError(f"Maze with UUID '{maze_uuid}' not found in directory '{mazes_directory}'")
        
        # Create wrapper instance
        return cls(maze_generator, start_position)
    
    @classmethod
    def list_available_mazes(cls, mazes_directory: str = "mazes") -> List[str]:
        """
        List all available maze UUIDs in the specified directory.
        
        Args:
            mazes_directory: Directory containing maze files
            
        Returns:
            List of available maze UUIDs
        """
        available_uuids = []
        
        if not os.path.exists(mazes_directory):
            return available_uuids
        
        # Look for JSON files with maze_ prefix
        for filename in os.listdir(mazes_directory):
            if filename.startswith("maze_") and filename.endswith(".json"):
                # Extract UUID from filename
                uuid_part = filename[5:-5]  # Remove "maze_" prefix and ".json" suffix
                available_uuids.append(uuid_part)
        
        return sorted(available_uuids)
    
    def get_maze_uuid(self) -> str:
        """
        Get the UUID of the current maze.
        
        Returns:
            UUID string of the current maze
        """
        return self.maze_generator.get_maze_uuid()
        
    def _find_starting_position(self) -> Tuple[int, int]:
        """
        Find the first available open position in the maze.
        
        Returns:
            Tuple of (row, col) for the starting position
        """
        # Look for open cells ('O'), preferring positions closer to top-left
        for r in range(self.height):
            for c in range(self.width):
                if self.maze[r, c] == 'O':
                    return (r, c)
        
        # If no open cell found, raise an error
        raise ValueError("No valid starting position found in maze")
    
    def _is_valid_position(self, position: Tuple[int, int]) -> bool:
        """
        Check if a position is valid for the agent.
        
        Args:
            position: Tuple of (row, col)
            
        Returns:
            True if position is valid (within bounds and not a wall/frame)
        """
        row, col = position
        
        # Check bounds
        if not (0 <= row < self.height and 0 <= col < self.width):
            return False
            
        # Check if cell is passable (open path or exit)
        cell_type = self.maze[row, col]
        return cell_type in ['O', 'E']
    
    def try_move(self, direction: str) -> bool:
        """
        Try to move the agent in the specified direction.
        
        Args:
            direction: One of 'up', 'down', 'left', 'right'
            
        Returns:
            True if move was successful, False if blocked
        """
        # Direction mappings (row, col) deltas
        direction_map = {
            'north': (-1, 0),
            'south': (1, 0),
            'west': (0, -1),
            'east': (0, 1)
        }
        
        if direction not in direction_map:
            raise ValueError(f"Invalid direction: {direction}. Must be one of {list(direction_map.keys())}")
        
        # Calculate new position
        delta_row, delta_col = direction_map[direction]
        new_row = self.agent_position[0] + delta_row
        new_col = self.agent_position[1] + delta_col
        new_position = (new_row, new_col)
        
        # Check if move is valid
        if self._is_valid_position(new_position):
            self.agent_position = new_position
            self.move_history.append(new_position)
            # Save state after successful move
            self.turn_complete()
            return True
        else:
            # Record the failed move attempt
            self.failed_moves.append(new_position)
            return False
    
    def move_to_position(self, position: Tuple[int, int]) -> bool:
        """
        Try to move the agent directly to a specific position.
        
        Args:
            position: Target position (row, col)
            
        Returns:
            True if move was successful, False if invalid position
        """
        if self._is_valid_position(position):
            self.agent_position = position
            self.move_history.append(position)
            # Save state after successful position move
            self.turn_complete()
            return True
        else:
            return False
    
    def get_agent_position(self) -> Tuple[int, int]:
        """
        Get the current position of the agent.
        
        Returns:
            Tuple of (row, col) representing agent's current position
        """
        return self.agent_position
    
    def get_possible_moves(self) -> List[str]:
        """
        Get list of valid moves from current position.
        
        Returns:
            List of valid direction strings
        """
        valid_moves = []
        directions = ['north', 'south', 'west', 'east']
        direction_map = {
            'north': (-1, 0),
            'south': (1, 0),
            'west': (0, -1),
            'east': (0, 1)
        }
        
        current_row, current_col = self.agent_position
        
        for direction in directions:
            delta_row, delta_col = direction_map[direction]
            new_position = (current_row + delta_row, current_col + delta_col)
            
            if self._is_valid_position(new_position):
                valid_moves.append(direction)
        
        return valid_moves
    
    def is_at_exit(self) -> bool:
        """
        Check if the agent is currently at the exit.
        
        Returns:
            True if agent is at the exit position
        """
        row, col = self.agent_position
        return self.maze[row, col] == 'E'
    
    def get_cell_type(self, position: Optional[Tuple[int, int]] = None) -> str:
        """
        Get the type of cell at the specified position.
        
        Args:
            position: Position to check. If None, uses agent's current position.
            
        Returns:
            Cell type character ('O', 'W', 'E', 'X')
        """
        if position is None:
            position = self.agent_position
            
        row, col = position
        if 0 <= row < self.height and 0 <= col < self.width:
            return self.maze[row, col]
        else:
            return 'X'  # Out of bounds treated as outer frame
    
    def reset_agent_position(self, position: Optional[Tuple[int, int]] = None) -> None:
        """
        Reset the agent to a specific position or back to the starting position.
        
        Args:
            position: New position for the agent. If None, resets to original start.
        """
        if position is not None:
            if self._is_valid_position(position):
                self.agent_position = position
            else:
                raise ValueError(f"Invalid reset position {position}")
        else:
            # Reset to the first position in move history (original start)
            self.agent_position = self.move_history[0]
        
        # Clear move history and failed moves, then add current position
        self.move_history = [self.agent_position]
        self.failed_moves = []
    
    def get_move_count(self) -> int:
        """
        Get the number of moves made by the agent.
        
        Returns:
            Number of moves (length of move history - 1)
        """
        return len(self.move_history) - 1
    
    def get_failed_move_count(self) -> int:
        """
        Get the number of failed move attempts.
        
        Returns:
            Number of times the agent tried to move into a wall
        """
        return len(self.failed_moves)
    
    def get_failed_moves(self) -> List[Tuple[int, int]]:
        """
        Get list of positions where the agent tried to move but failed.
        
        Returns:
            List of (row, col) positions of failed move attempts
        """
        return self.failed_moves.copy()
    
    def get_maze_with_agent(self) -> List[List[str]]:
        """
        Get the maze as a list of lists with the agent position marked.
        
        Returns:
            Maze representation with 'A' marking the agent's current position
        """
        # Get maze as list of lists
        maze_copy = [row[:] for row in self.maze_generator.get_maze()]
        
        # Mark agent position
        agent_row, agent_col = self.agent_position
        maze_copy[agent_row][agent_col] = 'A'
        
        return maze_copy
    
    def print_maze_with_agent(self) -> None:
        """
        Print the maze with the agent's current position marked as 'A'.
        """
        maze_with_agent = self.get_maze_with_agent()
        
        print(f"\nMaze with Agent (A) at position {self.agent_position}:")
        print(f"Moves made: {self.get_move_count()}")
        print("-" * (len(maze_with_agent[0]) * 2 + 1))
        
        for row in maze_with_agent:
            print("|" + " ".join(row) + "|")
        
        print("-" * (len(maze_with_agent[0]) * 2 + 1))
        
        # Print legend
        print("Legend: W=Wall, O=Open, E=Exit, X=Frame, A=Agent")
        
        if self.is_at_exit():
            print("🎉 CONGRATULATIONS! Agent reached the exit! 🎉")
    
    def get_maze_with_path_as_image(self, cell_size: int = 20, show_grid: bool = True, 
                                   show_path: bool = True, show_agent: bool = True, 
                                   show_start: bool = True, show_failed_moves: bool = True,
                                   show_marked_dead_ends: bool = True) -> plt.Figure:
        """
        Convert the maze with agent and path into a matplotlib figure.
        
        Args:
            cell_size: Size of each cell in pixels
            show_grid: Whether to show grid lines between cells
            show_path: Whether to show the path taken by the agent
            show_agent: Whether to show the agent's current position
            show_start: Whether to show the starting position
            show_failed_moves: Whether to show failed move attempts
            show_marked_dead_ends: Whether to show marked dead ends
            
        Returns:
            matplotlib Figure object representing the maze with agent and path
        """
        # Create figure and axis with more space for text
        fig, ax = plt.subplots(figsize=(14, 10))
        
        # Get maze dimensions
        height, width = self.maze.shape
        
        # Create a copy of the maze for visualization
        visual_maze = self.maze.copy()
        
        # Mark the path taken (excluding current position and starting position)
        if show_path and len(self.move_history) > 1:
            for i, (row, col) in enumerate(self.move_history[1:-1], 1):  # Exclude start and current position
                if visual_maze[row, col] in ['O', 'E']:  # Only mark on open cells or exit
                    visual_maze[row, col] = 'P'  # Mark as path
        
        # Mark starting position
        if show_start and len(self.move_history) > 0:
            start_row, start_col = self.move_history[0]  # First position is the start
            # Only mark as start if it's not the current agent position
            if (start_row, start_col) != self.agent_position and visual_maze[start_row, start_col] in ['O', 'E', 'P']:
                visual_maze[start_row, start_col] = 'S'  # Mark as starting point
        
        # Mark failed move attempts (on walls)
        if show_failed_moves:
            for fail_row, fail_col in self.failed_moves:
                # Only mark failed moves on walls (W) or outer frame (X)
                if (0 <= fail_row < self.height and 0 <= fail_col < self.width and
                    visual_maze[fail_row, fail_col] in ['W', 'X']):
                    visual_maze[fail_row, fail_col] = 'F'  # Mark as failed move
        
        # Mark dead ends (before agent position to allow agent to override)
        if show_marked_dead_ends:
            for dead_row, dead_col in self.marked_dead_ends:
                if 0 <= dead_row < self.height and 0 <= dead_col < self.width:
                    if visual_maze[dead_row, dead_col] == 'E':
                        visual_maze[dead_row, dead_col] = 'DE'  # Dead end exit
                    elif visual_maze[dead_row, dead_col] in ['O', 'P', 'S']:
                        visual_maze[dead_row, dead_col] = 'D'  # Dead end
        
        # Mark agent's current position (this should be last to take priority)
        if show_agent:
            agent_row, agent_col = self.agent_position
            visual_maze[agent_row, agent_col] = 'A'
        
        # Create a grid of colored squares
        for y in range(height):
            for x in range(width):
                cell = visual_maze[y, x]
                color = self.colors.get(cell, '#FFFFFF')  # Default to white if unknown
                
                # Special handling for marked dead ends (thicker borders)
                edge_color = 'black' if show_grid else 'none'
                linewidth = 0.5
                
                if show_marked_dead_ends and (y, x) in self.marked_dead_ends:
                    edge_color = '#8B0000'  # Dark red border for marked dead ends
                    linewidth = 2
                
                # Draw rectangle for each cell
                rect = plt.Rectangle((x, height-y-1), 1, 1, 
                                   facecolor=color, 
                                   edgecolor=edge_color,
                                   linewidth=linewidth)
                ax.add_patch(rect)
        
        # 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))
        
        if not show_grid:
            ax.set_xticks([])
            ax.set_yticks([])
        
        # Create title with agent info
        title = f'Maze Navigation - Agent at {self.agent_position}'
        ax.set_title(title, fontsize=14, fontweight='bold')
        
        # Get agent's marking analysis
        marking_analysis = self.get_dead_end_analysis()
        
        # Add statistics text box
        stats_text = f"""Agent Position: {self.agent_position}
Moves Made: {self.get_move_count()}
Failed Moves: {self.get_failed_move_count()}
At Exit: {'Yes' if self.is_at_exit() else 'No'}
Path Length: {len(self.move_history)}
Agent Marked Dead Ends: {marking_analysis['marked_dead_ends_count']}
Marking Actions: {marking_analysis['total_marking_actions']}
Mark Actions: {marking_analysis['mark_actions']}
Unmark Actions: {marking_analysis['unmark_actions']}
Possible Moves: {', '.join(self.get_possible_moves())}"""
        
        # Position text box in the upper right corner
        ax.text(0.98, 0.98, stats_text, 
                transform=ax.transAxes, 
                fontsize=9,
                verticalalignment='top',
                horizontalalignment='right',
                bbox=dict(boxstyle='round,pad=0.5', facecolor='white', alpha=0.8, edgecolor='black'),
                family='monospace')
        
        # Add 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='#00FF00', label='Exit'),
            plt.Rectangle((0, 0), 1, 1, facecolor='#808080', label='Frame'),
        ]
        
        if show_start:
            legend_elements.append(plt.Rectangle((0, 0), 1, 1, facecolor='#0066FF', label='Start'))
        
        if show_agent:
            legend_elements.append(plt.Rectangle((0, 0), 1, 1, facecolor='#00FFFF', label='Agent'))
        
        if show_path:
            legend_elements.append(plt.Rectangle((0, 0), 1, 1, facecolor='#FFB6C1', label='Path Taken'))
        
        if show_failed_moves:
            legend_elements.append(plt.Rectangle((0, 0), 1, 1, facecolor='#FF6600', label='Failed Moves'))
        
        if show_marked_dead_ends and len(self.marked_dead_ends) > 0:
            legend_elements.append(plt.Rectangle((0, 0), 1, 1, facecolor='#FF0000', 
                                               edgecolor='#8B0000', linewidth=2, label='Agent Marked Dead End'))
            legend_elements.append(plt.Rectangle((0, 0), 1, 1, facecolor='#FF8C00', 
                                               edgecolor='#8B0000', linewidth=2, label='Agent Marked Exit'))
        
        ax.legend(handles=legend_elements, loc='upper left', bbox_to_anchor=(0.02, 0.98), fontsize=8)
        
        # Add status indicator
        if self.is_at_exit():
            status_color = '#00FF00'  # Green - Success
            status_text = 'SUCCESS!\nREACHED EXIT'
        else:
            status_color = '#FFA500'  # Orange - In Progress
            status_text = 'IN PROGRESS'
        
        # Add status indicator in bottom right
        ax.text(0.98, 0.02, status_text, 
                transform=ax.transAxes,
                fontsize=12,
                fontweight='bold',
                verticalalignment='bottom',
                horizontalalignment='right',
                bbox=dict(boxstyle='round,pad=0.3', facecolor=status_color, alpha=0.7, edgecolor='black'),
                color='black')
        
        plt.tight_layout()
        return fig
    
    def save_maze_with_path_image(self, filename: str, cell_size: int = 20, show_grid: bool = True, 
                                 show_path: bool = True, show_agent: bool = True, 
                                 show_start: bool = True, show_failed_moves: bool = True, 
                                 show_marked_dead_ends: bool = True, dpi: int = 100) -> None:
        """
        Save the maze with agent and path as an image file.
        
        Args:
            filename: Name of the output file (e.g., 'maze_with_agent.png')
            cell_size: Size of each cell in pixels
            show_grid: Whether to show grid lines between cells
            show_path: Whether to show the path taken by the agent
            show_agent: Whether to show the agent's current position
            show_start: Whether to show the starting position
            show_failed_moves: Whether to show failed move attempts
            show_marked_dead_ends: Whether to show marked dead ends
            dpi: Dots per inch for the output image
        """
        # Create output directory if it doesn't exist
        output_dir = "maze_wrapper_outputs"
        os.makedirs(output_dir, exist_ok=True)
        
        # Create full path for the file
        full_path = os.path.join(output_dir, filename)
        
        fig = self.get_maze_with_path_as_image(cell_size, show_grid, show_path, show_agent, 
                                              show_start, show_failed_moves, show_marked_dead_ends)
        fig.savefig(full_path, dpi=dpi, bbox_inches='tight')
        plt.close(fig)
        print(f"Maze visualization saved as {full_path}")

    def show_maze_interactive(self, show_path: bool = True, show_agent: bool = True, 
                             show_start: bool = True, show_failed_moves: bool = True,
                             show_marked_dead_ends: bool = True) -> None:
        """
        Display the maze with agent and path in an interactive matplotlib window.
        
        Args:
            show_path: Whether to show the path taken by the agent
            show_agent: Whether to show the agent's current position
            show_start: Whether to show the starting position  
            show_failed_moves: Whether to show failed move attempts
            show_marked_dead_ends: Whether to show marked dead ends
        """
        fig = self.get_maze_with_path_as_image(show_path=show_path, show_agent=show_agent, 
                                              show_start=show_start, show_failed_moves=show_failed_moves,
                                              show_marked_dead_ends=show_marked_dead_ends)
        # Try to show interactively; if backend is headless, save instead (controlled by env flag)
        # 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_wrapper_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 print_dead_end_analysis(self) -> None:
        """
        Print analysis of agent's marked dead ends (no automatic detection).
        """
        analysis = self.get_dead_end_analysis()
        
        print(f"\n{'='*60}")
        print(f"AGENT DEAD END MARKING ANALYSIS")
        print(f"{'='*60}")
        print(f"Agent marked dead ends: {analysis['marked_dead_ends_count']}")
        print(f"Total marking actions: {analysis['total_marking_actions']}")
        print(f"  - Mark actions: {analysis['mark_actions']}")
        print(f"  - Unmark actions: {analysis['unmark_actions']}")
        
        if analysis['marked_dead_ends_positions']:
            print(f"Marked positions: {sorted(analysis['marked_dead_ends_positions'])}")
        
        print(f"Agent's marking strategy: {self._analyze_marking_strategy()}")
        print(f"{'='*60}")

    def _analyze_marking_strategy(self) -> str:
        """
        Analyze the agent's marking strategy based on marking history.
        
        Returns:
            String describing the agent's marking behavior
        """
        total_actions = len(self.marking_history)
        if total_actions == 0:
            return "No marking actions taken"
        
        mark_count = len([action for action, pos in self.marking_history if action == 'mark'])
        unmark_count = len([action for action, pos in self.marking_history if action == 'unmark'])
        
        if unmark_count == 0:
            return "Conservative (only marks, never unmarks)"
        elif unmark_count > mark_count * 0.3:
            return "Adaptive (frequently reconsiders markings)"
        else:
            return "Strategic (marks with occasional corrections)"

    def mark_dead_end(self, position: Optional[Tuple[int, int]] = None) -> bool:
        """
        Mark a position as a dead end. This is an action the agent can take.
        
        Args:
            position: Position to mark (row, col). If None, marks current position.
            
        Returns:
            True if successfully marked, False if invalid position
        """
        if position is None:
            position = self.agent_position
        
        row, col = position
        
        # Check if position is valid and is an open path or exit
        if (0 <= row < self.height and 
            0 <= col < self.width and 
            self.maze[row, col] in ['O', 'E']):
            
            if position not in self.marked_dead_ends:
                self.marked_dead_ends.add(position)
                self.marking_history.append(('mark', position))
                # Save state after marking dead end
                self.turn_complete()
                return True
            else:
                # Already marked
                return False
        return False

    def unmark_dead_end(self, position: Optional[Tuple[int, int]] = None) -> bool:
        """
        Remove dead end marking from a position. This is an action the agent can take.
        
        Args:
            position: Position to unmark (row, col). If None, unmarks current position.
            
        Returns:
            True if successfully unmarked, False if position wasn't marked
        """
        if position is None:
            position = self.agent_position
        
        if position in self.marked_dead_ends:
            self.marked_dead_ends.remove(position)
            self.marking_history.append(('unmark', position))
            # Save state after unmarking dead end
            self.turn_complete()
            return True
        return False

    def is_position_marked_dead_end(self, position: Optional[Tuple[int, int]] = None) -> bool:
        """
        Check if a position is marked as a dead end.
        
        Args:
            position: Position to check (row, col). If None, checks current position.
            
        Returns:
            True if position is marked as dead end, False otherwise
        """
        if position is None:
            position = self.agent_position
        return position in self.marked_dead_ends

    def get_marked_dead_ends(self) -> List[Tuple[int, int]]:
        """
        Get list of all positions marked as dead ends.
        
        Returns:
            List of (row, col) tuples for marked dead ends
        """
        return list(self.marked_dead_ends)

    def get_marking_count(self) -> int:
        """
        Get the total number of marking actions (mark + unmark) performed.
        
        Returns:
            Number of marking actions performed
        """
        return len(self.marking_history)

    def clear_all_marked_dead_ends(self) -> int:
        """
        Clear all marked dead ends. This is an action the agent can take.
        
        Returns:
            Number of dead ends cleared
        """
        cleared_count = len(self.marked_dead_ends)
        
        # Record all clearing actions
        for position in self.marked_dead_ends:
            self.marking_history.append(('unmark', position))
        
        self.marked_dead_ends.clear()
        
        # Save state after clearing if any dead ends were cleared
        if cleared_count > 0:
            self.turn_complete()
            
        return cleared_count

    def get_dead_end_analysis(self) -> dict:
        """
        Get analysis of agent's marked dead ends (no automatic detection).
        
        Returns:
            Dictionary containing agent's marking information
        """
        analysis = {
            'marked_dead_ends_count': len(self.marked_dead_ends),
            'marked_dead_ends_positions': list(self.marked_dead_ends),
            'total_marking_actions': len(self.marking_history),
            'mark_actions': len([action for action, pos in self.marking_history if action == 'mark']),
            'unmark_actions': len([action for action, pos in self.marking_history if action == 'unmark'])
        }
        
        return analysis

    def get_current_view(self) -> str:
        """
        Get the current view from the current position.
        
        Returns:
            String representation of the view
        """
        """
        Get the current view from the current position.
        
        Returns:
            String representation showing:
            - Current position (A)
            - Walls (W)
            - Open paths (O) 
            - Exit (E)
            - Dead ends (D)
            - Outer frame (X)
        """
        # Get current position
        row, col = self.agent_position
        
        # Build view string
        view = "Current view:\n"
        view += "Legend:\n"
        view += "  A = Agent position\n"
        view += "  W = Wall\n" 
        view += "  O = Open path\n"
        view += "  E = Exit\n"
        view += "  D = Marked dead end\n"
        view += "  X = Outer frame\n\n"
        
        # Add 3x3 grid view centered on agent
        for r in range(max(0, row-1), min(self.height, row+2)):
            for c in range(max(0, col-1), min(self.width, col+2)):
                if (r,c) == self.agent_position:
                    view += "A"
                elif (r,c) in self.marked_dead_ends:
                    view += "D"
                else:
                    view += self.maze[r,c]
            view += "\n"
            
        return view

    def turn_complete(self) -> None:
        """
        Save a snapshot of the current maze state, including agent position,
        marked dead ends, and move history.
        """
        current_state = {
            'agent_position': self.agent_position,
            'marked_dead_ends': deepcopy(self.marked_dead_ends),
            'move_history': deepcopy(self.move_history),
            'failed_moves': deepcopy(self.failed_moves),
            'maze_uuid': self.get_maze_uuid()
        }
        self.state_history.append(current_state)

    def _create_state_matrix(self, agent_position: Tuple[int, int], 
                            marked_dead_ends: Set[Tuple[int, int]], 
                            view_radius: int = 3) -> np.ndarray:
        """
        Create a state matrix representing what the agent can see from their position.
        
        Args:
            agent_position: Position of the agent (row, col)
            marked_dead_ends: Set of positions marked as dead ends
            view_radius: Radius of vision around agent position
            
        Returns:
            numpy array representing the visible state matrix
        """
        row, col = agent_position
        
        # Create view matrix centered on agent
        view_size = 2 * view_radius + 1
        state_matrix = np.full((view_size, view_size), 'X', dtype='<U2')  # Default to out-of-bounds
        
        # Fill in visible cells
        for i in range(view_size):
            for j in range(view_size):
                maze_row = row - view_radius + i
                maze_col = col - view_radius + j
                
                # Check if position is within maze bounds
                if 0 <= maze_row < self.height and 0 <= maze_col < self.width:
                    cell_value = self.maze[maze_row, maze_col]
                    
                    # Check if position is marked as dead end
                    if (maze_row, maze_col) in marked_dead_ends:
                        if cell_value == 'E':
                            state_matrix[i, j] = 'DE'  # Dead end exit
                        else:
                            state_matrix[i, j] = 'D'   # Dead end
                    # Mark agent position
                    elif (maze_row, maze_col) == agent_position:
                        state_matrix[i, j] = 'A'
                    else:
                        state_matrix[i, j] = cell_value
        
        return state_matrix
    
    def compute_state_similarity(self, other_state: Dict[str, Any], view_radius: int = 3) -> float:
        """
        Compute similarity between current state and another state based on visible state matrices.
        
        Args:
            other_state: Previous state dictionary to compare with
            view_radius: Radius of vision for state matrix comparison
            
        Returns:
            Similarity score between 0.0 (no similarity) and 1.0 (identical)
        """
        # Create state matrices
        current_matrix = self._create_state_matrix(
            self.agent_position, 
            self.marked_dead_ends, 
            view_radius
        )
        
        other_matrix = self._create_state_matrix(
            other_state['agent_position'],
            other_state['marked_dead_ends'],
            view_radius
        )
        
        # Compare matrices element by element
        total_cells = current_matrix.size
        matching_cells = np.sum(current_matrix == other_matrix)
        
        similarity = matching_cells / total_cells
        return float(similarity)
    
    def compute_recent_state_similarity(self, n: int = 5, view_radius: int = 3) -> List[float]:
        """
        Compute similarity scores for the last n states compared to the current state.
        
        Args:
            n: Number of previous states to compare (default: 5)
            view_radius: Radius of vision for state matrix comparison (default: 3)
            
        Returns:
            List of similarity scores, where index 0 is most recent state similarity
        """
        if len(self.state_history) <= 1:
            return []
        
        # Get the last n states (excluding the current one which is the last in history)
        recent_states = self.state_history[-(n+1):-1]  # Exclude current state
        
        similarities = []
        for state in reversed(recent_states):  # Most recent first
            similarity = self.compute_state_similarity(state, view_radius)
            similarities.append(similarity)
        
        return similarities
    
    def get_state_history_length(self) -> int:
        """
        Get the number of states stored in history.
        
        Returns:
            Number of state snapshots in history
        """
        return len(self.state_history)
    
    def clear_state_history(self) -> int:
        """
        Clear all state history except the current state.
        
        Returns:
            Number of states that were cleared
        """
        cleared_count = len(self.state_history) - 1  # Keep current state
        if cleared_count > 0:
            current_state = self.state_history[-1]  # Save current state
            self.state_history = [current_state]
        return max(0, cleared_count)
    
    def get_state_similarity_analysis(self, n: int = 5, view_radius: int = 3) -> Dict[str, Any]:
        """
        Get comprehensive analysis of recent state similarities.
        
        Args:
            n: Number of recent states to analyze
            view_radius: Radius of vision for state comparison
            
        Returns:
            Dictionary containing similarity analysis metrics
        """
        similarities = self.compute_recent_state_similarity(n, view_radius)
        
        if not similarities:
            return {
                'similarities': [],
                'average_similarity': 0.0,
                'max_similarity': 0.0,
                'min_similarity': 0.0,
                'perfect_matches': 0,
                'total_states_compared': 0
            }
        
        return {
            'similarities': similarities,
            'average_similarity': float(np.mean(similarities)),
            'max_similarity': float(np.max(similarities)),
            'min_similarity': float(np.min(similarities)),
            'perfect_matches': sum(1 for s in similarities if s == 1.0),
            'total_states_compared': len(similarities)
        }

def main():
    """
    Test function to demonstrate MazeWrapper visualization capabilities including dead end marking and UUID loading.
    
    Creates a maze, moves the agent around, demonstrates dead end marking, and shows various visualizations.
    Also demonstrates loading mazes by UUID from the mazes folder.
    """
    print("="*70)
    print("MAZE WRAPPER WITH DEAD END MARKING AND UUID LOADING TEST")
    print("="*70)
    
    # First, try to load a maze from UUID if available
    print("1. Checking for available mazes in 'mazes' folder...")
    available_mazes = MazeWrapper.list_available_mazes()
    
    if available_mazes:
        print(f"   Found {len(available_mazes)} available mazes:")
        for i, uuid in enumerate(available_mazes[:5]):  # Show first 5
            print(f"     {i+1}. {uuid}")
        if len(available_mazes) > 5:
            print(f"     ... and {len(available_mazes) - 5} more")
        
        # Load the first available maze
        print(f"\n2. Loading maze with UUID: {available_mazes[0]}")
        try:
            maze_wrapper = MazeWrapper.from_uuid(available_mazes[0])
            print(f"   ✅ Successfully loaded maze {maze_wrapper.get_maze_uuid()}")
        except Exception as e:
            print(f"   ❌ Failed to load maze: {e}")
            print("   Creating a new maze instead...")
            maze_gen = MazeGenerator(maze_size=12)
            maze_wrapper = MazeWrapper(maze_gen)
    else:
        print("   No mazes found in 'mazes' folder")
        print("   Creating a new 12x12 maze...")
        maze_gen = MazeGenerator(maze_size=12)
        maze_wrapper = MazeWrapper(maze_gen)
    
    print(f"Agent starting position: {maze_wrapper.get_agent_position()}")
    
    # Show initial state
    print("\n3. Initial maze state:")
    maze_wrapper.print_maze_with_agent()
    
    # Save initial state as image
    maze_wrapper.save_maze_with_path_image("01_initial_maze.png")
    
    # Demonstrate dead end marking at current position
    print("\n4. Testing dead end marking actions...")
    print(f"Current position: {maze_wrapper.get_agent_position()}")
    
    # Try to mark current position as dead end
    mark_success = maze_wrapper.mark_dead_end()
    print(f"✅ Marked current position as dead end: {mark_success}")
    
    # Try to mark a specific position
    test_position = (5, 7)  # Arbitrary position
    mark_success = maze_wrapper.mark_dead_end(test_position)
    print(f"✅ Marked position {test_position} as dead end: {mark_success}")
    
    # Save state with marked dead ends
    maze_wrapper.save_maze_with_path_image("02_with_marked_dead_ends.png")
    
    # Make some moves to create a path (including some that will fail)
    print("\n5. Making some moves to create a path...")
    moves_to_try = ['east', 'right', 'up', 'up', 'down', 'down', 'left', 'left', 'right', 'down', 'right', 'down']
    
    successful_moves = 0
    for i, move in enumerate(moves_to_try):
        success = maze_wrapper.try_move(move)
        if success:
            successful_moves += 1
            print(f"✅ Move {i+1}: {move} successful - Position: {maze_wrapper.get_agent_position()}")
            
            # Mark some positions as dead ends during navigation
            if (i + 1) % 4 == 0:  # Mark every 4th position
                mark_success = maze_wrapper.mark_dead_end()
                print(f"   📍 Marked position {maze_wrapper.get_agent_position()} as dead end: {mark_success}")
            
            # Save intermediate state for some moves
            if (i + 1) % 3 == 0:  # Save every 3rd move
                filename = f"03_maze_step_{i+1}.png"
                maze_wrapper.save_maze_with_path_image(filename)
                print(f"   Saved state as {filename}")
                
            if maze_wrapper.is_at_exit():
                print("🎉 Reached the exit!")
                break
        else:
            print(f"❌ Move {i+1}: {move} failed - blocked")
    
    print(f"\nNavigation results:")
    print(f"Total successful moves: {successful_moves}")
    print(f"Total failed moves: {maze_wrapper.get_failed_move_count()}")
    print(f"Final position: {maze_wrapper.get_agent_position()}")
    print(f"At exit: {maze_wrapper.is_at_exit()}")
    
    # Show final state
    print("\n6. Final maze state:")
    maze_wrapper.print_maze_with_agent()
    
    # Save visualization variants
    print("\n7. Saving different visualization variants...")
    maze_wrapper.save_maze_with_path_image("05_final_with_path_and_dead_ends.png", 
                                          show_path=True, show_marked_dead_ends=True)
    maze_wrapper.save_maze_with_path_image("06_final_no_path.png", 
                                          show_path=False, show_marked_dead_ends=True)
    maze_wrapper.save_maze_with_path_image("07_final_no_dead_ends.png", 
                                          show_path=True, show_marked_dead_ends=False)
    
    # Demonstrate unmarking
    print("\n8. Testing unmarking functionality...")
    marked_positions = maze_wrapper.get_marked_dead_ends()
    if len(marked_positions) >= 2:
        pos_to_unmark = marked_positions[0]
        unmark_success = maze_wrapper.unmark_dead_end(pos_to_unmark)
        print(f"✅ Unmarked position {pos_to_unmark}: {unmark_success}")
        
        # Save state after unmarking
        maze_wrapper.save_maze_with_path_image("08_after_unmarking.png")
    
    # Try to find and reach the exit
    print("\n9. Attempting to reach the exit...")
    exit_found = False
    max_attempts = 50
    
    for attempt in range(max_attempts):
        if maze_wrapper.is_at_exit():
            exit_found = True
            break
            
        # Get possible moves and choose randomly
        possible_moves = maze_wrapper.get_possible_moves()
        if possible_moves:
            import random
            move = random.choice(possible_moves)
            success = maze_wrapper.try_move(move)
            if success and (attempt + 1) % 10 == 0:  # Print every 10th move
                print(f"   Attempt {attempt + 1}: Moved {move} to {maze_wrapper.get_agent_position()}")
        else:
            print("   No possible moves - agent is stuck!")
            break
    
    if exit_found:
        print(f"🎉 SUCCESS! Agent reached the exit in {maze_wrapper.get_move_count()} total moves!")
        maze_wrapper.save_maze_with_path_image("09_victory_maze.png")
    else:
        print(f"❌ Could not reach exit after {max_attempts} attempts")
    
    # Final statistics
    print(f"\n10. Final Statistics:")
    print(f"   Total moves made: {maze_wrapper.get_move_count()}")
    print(f"   Total failed moves: {maze_wrapper.get_failed_move_count()}")
    total_attempts = maze_wrapper.get_move_count() + maze_wrapper.get_failed_move_count()
    success_rate = (maze_wrapper.get_move_count() / total_attempts * 100) if total_attempts > 0 else 0
    print(f"   Success rate: {success_rate:.1f}%")
    print(f"   Path length: {len(maze_wrapper.move_history)}")
    print(f"   Unique positions visited: {len(set(maze_wrapper.move_history))}")
    print(f"   Unique failed positions: {len(set(maze_wrapper.failed_moves))}")
    print(f"   Dead ends marked: {len(maze_wrapper.get_marked_dead_ends())}")
    print(f"   Total marking actions: {maze_wrapper.get_marking_count()}")
    
    # Final dead end analysis
    print("\n11. Final dead end analysis:")
    maze_wrapper.print_dead_end_analysis()
    
    # Test clearing all markings
    print("\n12. Testing clear all markings...")
    cleared_count = maze_wrapper.clear_all_marked_dead_ends()
    print(f"✅ Cleared {cleared_count} marked dead ends")
    maze_wrapper.save_maze_with_path_image("10_cleared_markings.png")
    
    # Show interactive visualization (if running in interactive environment)
    try:
        print("\n13. Showing interactive visualization...")
        print("   (Close the window to continue)")
        maze_wrapper.show_maze_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" + "="*70)
    print("MAZE WRAPPER WITH DEAD END MARKING AND UUID LOADING TEST COMPLETE")
    print("="*70)
    print("Generated files in maze_wrapper_outputs/:")
    print("- 01_initial_maze.png")
    print("- 02_with_marked_dead_ends.png")
    print("- 03_maze_step_*.png (intermediate states)")

    print("- 05_final_with_path_and_dead_ends.png")
    print("- 06_final_no_path.png")
    print("- 07_final_no_dead_ends.png")
    print("- 08_after_unmarking.png")
    if exit_found:
        print("- 09_victory_maze.png")
    print("- 10_cleared_markings.png")
    print("\nKey features demonstrated:")
    print("✅ UUID-based maze loading from 'mazes' folder")
    print("✅ Listing available mazes by UUID")
    print("✅ Agent movement and pathfinding")
    print("✅ Manual dead end marking at current/specific positions")
    print("✅ Dead end unmarking and clearing")
    print("✅ Visual highlighting of marked dead ends")
    print("✅ Comprehensive statistics and analysis")
    print("✅ Multiple visualization options")
    print("="*70)


if __name__ == "__main__":
    main()
