import numpy as np
import random
import uuid
import json
import math
from typing import List, Tuple, Dict, Optional, Set
from collections import deque, Counter
import matplotlib.pyplot as plt
import os

class MazeGenerator:
    """
    Generates mazes using recursive backtracking algorithm.
    Uses numpy arrays for efficient maze generation.
    Each maze gets a unique 16-character UUID for identification.
    
    Enhanced with Shannon entropy-based complexity evaluation following
    the AMaze framework for scientific maze evaluation.
    """
    def __init__(self, maze_size: int, difficulty_level: str = "medium"):
        """
        Initialize maze generator with given size and difficulty level.
        
        Args:
            maze_size: Size of the maze (will be increased by 2 for outer frame)
            difficulty_level: Target difficulty ("easy", "medium", "hard", "very_hard")
        """
        self.maze_size = maze_size
        self.maze_uuid = self._generate_uuid()
        self.difficulty_level = difficulty_level  # Set difficulty level first
        self.maze = self.generate_maze(difficulty_level=difficulty_level)
        
        # Color mapping for maze elements
        self.colors = {
            'W': '#000000',  # Black for walls
            'O': '#FFFFFF',  # White for open paths
            'E': '#00FF00',  # Green for exit
            'X': '#808080'   # Gray for outer frame
        }

    def _generate_uuid(self) -> str:
        """Generate a 16-character UUID for the maze."""
        return str(uuid.uuid4()).replace('-', '')[:16]

    def get_maze_uuid(self) -> str:
        """Get the UUID of the current maze."""
        return self.maze_uuid

    def save_maze_data(self, directory: str = "maze_outputs") -> str:
        """
        Save maze data including UUID, size, and array to JSON file.
        
        Args:
            directory: Directory to save the maze data
            
        Returns:
            Path to the saved file
        """
        os.makedirs(directory, exist_ok=True)
        
        # Prepare maze data
        maze_data = {
            'uuid': self.maze_uuid,
            'size': self.maze_size,
            'maze_array': self.maze.tolist(),
            'dead_end_analysis': self.analyze_dead_ends(),
            'shannon_complexity': self.calculate_shannon_complexity(),
            'difficulty_score': self.calculate_difficulty_score()
        }
        
        # Save to JSON file named with UUID
        filename = f"maze_{self.maze_uuid}.json"
        filepath = os.path.join(directory, filename)
        
        with open(filepath, 'w') as f:
            json.dump(maze_data, f, indent=2)
        
        return filepath

    @classmethod
    def load_maze_data(cls, filepath: str) -> 'MazeGenerator':
        """
        Load maze from JSON file.
        
        Args:
            filepath: Path to the JSON file
            
        Returns:
            MazeGenerator instance with loaded maze
        """
        with open(filepath, 'r') as f:
            maze_data = json.load(f)
        
        # Create new instance
        generator = cls.__new__(cls)
        generator.maze_size = maze_data['size']
        generator.maze_uuid = maze_data['uuid']
        generator.maze = np.array(maze_data['maze_array'])
        
        # Color mapping for maze elements
        generator.colors = {
            'W': '#000000',  # Black for walls
            'O': '#FFFFFF',  # White for open paths
            'E': '#00FF00',  # Green for exit
            'X': '#808080'   # Gray for outer frame
        }
        
        return generator

    @classmethod
    def find_maze_by_uuid(cls, uuid_str: str, directory: str = "mazes") -> Optional['MazeGenerator']:
        """
        Find and load a maze by its UUID.
        
        Args:
            uuid_str: UUID string to search for
            directory: Directory to search in
            
        Returns:
            MazeGenerator instance if found, None otherwise
        """ 
        filename = f"maze_{uuid_str}.json"
        filepath = os.path.join(directory, filename)
        
        if os.path.exists(filepath):
            return cls.load_maze_data(filepath)
        return None

    def get_maze_info(self) -> Dict:
        """
        Get comprehensive information about the maze including UUID.
        
        Returns:
            Dictionary with maze information
        """
        analysis = self.analyze_dead_ends()
        
        return {
            'uuid': self.maze_uuid,
            'size': self.maze_size,
            'dimensions': f"{self.maze.shape[0]}x{self.maze.shape[1]}",
            'difficulty_score': self.calculate_difficulty_score(),
            'dead_end_analysis': analysis,
            'total_cells': self.maze.shape[0] * self.maze.shape[1],
            'open_cells': int(np.sum((self.maze == 'O') | (self.maze == 'E'))),
            'wall_cells': int(np.sum(self.maze == 'W')),
            'outer_frame_cells': int(np.sum(self.maze == 'X'))
        }

    def _get_unvisited_neighbors(self, grid: np.ndarray, r: int, c: int) -> List[Tuple[int, int]]:
        """
        Get unvisited neighbors that are two steps away.
        
        Args:
            grid: The maze grid
            r, c: Current position
            
        Returns:
            List of (row, col) tuples for valid unvisited neighbors
        """
        neighbors = []
        for dr, dc in [(0, 2), (0, -2), (2, 0), (-2, 0)]:
            nr, nc = r + dr, c + dc
            if (0 <= nr < self.maze_size and 
                0 <= nc < self.maze_size and 
                grid[nr, nc] == 'W'):
                neighbors.append((nr, nc))
        return neighbors

    def generate_maze(self, dead_end_factor: float = 0.25, difficulty_level: str = "medium") -> np.ndarray:
        """
        Generate a maze using recursive backtracking with dead ends.
        
        Args:
            dead_end_factor: Probability of creating dead ends (0.0 to 1.0)
            difficulty_level: Target difficulty ("easy", "medium", "hard", "very_hard")
            
        Returns:
            numpy array representing the maze with symbols:
            'W' - Wall
            'O' - Open path
            'E' - Exit
            'X' - Outer frame
        """
        if not (0.0 <= dead_end_factor <= 1.0):
            raise ValueError("dead_end_factor must be between 0.0 and 1.0")

        # Store difficulty level for use in scoring
        self.difficulty_level = difficulty_level

        # Adjust dead end factor based on difficulty level
        if difficulty_level == "easy":
            dead_end_factor *= 0.3  # Significantly reduce dead ends for easy mazes
        elif difficulty_level == "medium":
            dead_end_factor *= 0.8  # Less reduction to allow more complexity for medium mazes
        elif difficulty_level == "hard":
            dead_end_factor *= 1.0  # Keep original factor for hard mazes
        elif difficulty_level == "very_hard":
            dead_end_factor *= 1.3  # Increase dead ends for very hard mazes

        # Handle special cases
        if self.maze_size < 1:
            return np.array([])
        if self.maze_size == 1:
            return np.array([['X', 'X', 'X'],
                           ['X', 'E', 'X'],
                           ['X', 'X', 'X']])
        if self.maze_size == 2:
            return np.array([['X', 'X', 'X', 'X'],
                           ['X', 'O', 'E', 'X'],
                           ['X', 'W', 'W', 'X'],
                           ['X', 'X', 'X', 'X']])

        # Initialize maze with walls
        grid = np.full((self.maze_size, self.maze_size), 'W', dtype=str)
        
        # Use multiple starting points to ensure better distribution
        starting_points = self._get_distributed_starting_points()
        all_visited = set()
        
        for start_r, start_c in starting_points:
            if (start_r, start_c) in all_visited:
                continue
                
            grid[start_r, start_c] = 'O'
            stack = [(start_r, start_c)]
            visited = {(start_r, start_c)}
            all_visited.add((start_r, start_c))

            while stack:
                r, c = stack[-1]
                neighbors = self._get_unvisited_neighbors(grid, r, c)
                
                if neighbors:
                    # Reduce dead end factor for better connectivity
                    if random.random() < dead_end_factor * 0.7:  # Reduced dead end probability
                        stack.pop()
                    else:
                        # Choose a random unvisited neighbor
                        nr, nc = random.choice(neighbors)
                        
                        # Carve path by marking the wall and the neighbor as open
                        wall_r, wall_c = (r + nr) // 2, (c + nc) // 2
                        grid[wall_r, wall_c] = 'O'
                        grid[nr, nc] = 'O'
                        
                        visited.add((nr, nc))
                        all_visited.add((nr, nc))
                        stack.append((nr, nc))
                else:
                    stack.pop()

        # Ensure connectivity between different regions
        self._ensure_connectivity(grid)
        
        # Skip distributed paths for easy mazes, apply for others based on difficulty
        if difficulty_level == "easy":
            # Skip distributed paths to reduce complexity
            pass
        elif difficulty_level in ["medium", "hard"]:
            # Apply normal distributed paths
            self._ensure_distributed_paths(grid)
        elif difficulty_level == "very_hard":
            # Apply more aggressive distributed paths for very hard mazes
            self._ensure_distributed_paths(grid)
            # Add extra complexity for very hard mazes
            self._add_extra_complexity(grid)

        # Place exit with bias for difficulty level
        open_cells = list(zip(*np.where(grid == 'O')))
        if not open_cells:
            grid[1, 1] = 'E'
        else:
            if difficulty_level == "easy":
                # For easy mazes, place exit in corner areas for accessibility
                corner_cells = self._get_corner_open_cells(open_cells)
                if corner_cells:
                    exit_r, exit_c = random.choice(corner_cells)
                else:
                    exit_r, exit_c = random.choice(open_cells)
            elif difficulty_level == "medium":
                # For medium mazes, prefer corner or edge areas but allow some randomness
                corner_cells = self._get_corner_open_cells(open_cells)
                if corner_cells and random.random() < 0.8:  # 80% chance to use corner
                    exit_r, exit_c = random.choice(corner_cells)
                else:
                    exit_r, exit_c = random.choice(open_cells)
            elif difficulty_level == "hard":
                # For hard mazes, use corner/edge placement with some central bias
                corner_cells = self._get_corner_open_cells(open_cells)
                central_cells = self._get_central_open_cells(open_cells)
                if corner_cells and central_cells and random.random() < 0.6:
                    # Mix corner and central placement
                    candidates = corner_cells + central_cells[:len(corner_cells)//2]
                    exit_r, exit_c = random.choice(candidates)
                elif corner_cells:
                    exit_r, exit_c = random.choice(corner_cells)
                else:
                    exit_r, exit_c = random.choice(open_cells)
            elif difficulty_level == "very_hard":
                # For very hard mazes, prefer central or hard-to-reach locations
                central_cells = self._get_central_open_cells(open_cells)
                if central_cells and random.random() < 0.7:  # 70% chance for central placement
                    exit_r, exit_c = random.choice(central_cells)
                else:
                    # Use corner cells as fallback
                    corner_cells = self._get_corner_open_cells(open_cells)
                    if corner_cells:
                        exit_r, exit_c = random.choice(corner_cells)
                    else:
                        exit_r, exit_c = random.choice(open_cells)
            else:
                # Default fallback
                exit_r, exit_c = random.choice(open_cells)
            grid[exit_r, exit_c] = 'E'

        # Add outer frame
        framed_size = self.maze_size + 2
        framed_grid = np.full((framed_size, framed_size), 'X', dtype=str)
        framed_grid[1:-1, 1:-1] = grid

        return framed_grid

    def _get_distributed_starting_points(self) -> List[Tuple[int, int]]:
        """
        Get multiple starting points distributed across the maze to ensure better coverage.
        """
        points = []
        
        # Always start with a point near the center
        center_r, center_c = self.maze_size // 2, self.maze_size // 2
        if center_r % 2 == 0:
            center_r += 1
        if center_c % 2 == 0:
            center_c += 1
        points.append((center_r, center_c))
        
        # Add corner points for larger mazes
        if self.maze_size >= 15:
            corners = [
                (1, 1),  # Top-left
                (1, self.maze_size - 2),  # Top-right
                (self.maze_size - 2, 1),  # Bottom-left
                (self.maze_size - 2, self.maze_size - 2)  # Bottom-right
            ]
            # Ensure odd coordinates for proper maze generation
            corners = [(r if r % 2 == 1 else r + 1, c if c % 2 == 1 else c + 1) for r, c in corners]
            points.extend(corners)
        
        # Add mid-edge points for very large mazes
        if self.maze_size >= 25:
            mid = self.maze_size // 2
            if mid % 2 == 0:
                mid += 1
            edges = [
                (1, mid),  # Top-middle
                (mid, 1),  # Left-middle
                (mid, self.maze_size - 2),  # Right-middle
                (self.maze_size - 2, mid)  # Bottom-middle
            ]
            points.extend(edges)
        
        # Filter points to ensure they're within bounds and odd coordinates
        valid_points = []
        for r, c in points:
            if (1 <= r < self.maze_size - 1 and 1 <= c < self.maze_size - 1 and
                r % 2 == 1 and c % 2 == 1):
                valid_points.append((r, c))
        
        return valid_points if valid_points else [(1, 1)]

    def _ensure_connectivity(self, grid: np.ndarray) -> None:
        """
        Ensure different regions of the maze are connected.
        """
        # Find all open cells
        open_cells = list(zip(*np.where(grid == 'O')))
        if len(open_cells) < 2:
            return
        
        # Use flood fill to find connected components
        visited = set()
        components = []
        
        for cell in open_cells:
            if cell not in visited:
                component = self._flood_fill(grid, cell, visited)
                components.append(component)
        
        # If we have multiple components, connect them
        if len(components) > 1:
            main_component = max(components, key=len)
            
            for component in components:
                if component != main_component:
                    # Find the closest cells between components
                    min_dist = float('inf')
                    best_connection = None
                    
                    for cell1 in component:
                        for cell2 in main_component:
                            dist = abs(cell1[0] - cell2[0]) + abs(cell1[1] - cell2[1])
                            if dist < min_dist:
                                min_dist = dist
                                best_connection = (cell1, cell2)
                    
                    if best_connection:
                        self._connect_cells(grid, best_connection[0], best_connection[1])

    def _flood_fill(self, grid: np.ndarray, start: Tuple[int, int], visited: set) -> List[Tuple[int, int]]:
        """
        Flood fill to find connected open cells.
        """
        component = []
        stack = [start]
        
        while stack:
            r, c = stack.pop()
            if (r, c) in visited:
                continue
                
            visited.add((r, c))
            component.append((r, c))
            
            # Check 4-connected neighbors
            for dr, dc in [(0, 1), (0, -1), (1, 0), (-1, 0)]:
                nr, nc = r + dr, c + dc
                if (0 <= nr < self.maze_size and 0 <= nc < self.maze_size and
                    grid[nr, nc] == 'O' and (nr, nc) not in visited):
                    stack.append((nr, nc))
        
        return component

    def _connect_cells(self, grid: np.ndarray, cell1: Tuple[int, int], cell2: Tuple[int, int]) -> None:
        """
        Connect two cells by carving a path between them.
        """
        r1, c1 = cell1
        r2, c2 = cell2
        
        # Simple path: go horizontal first, then vertical
        current_r, current_c = r1, c1
        
        # Move horizontally
        while current_c != c2:
            if current_c < c2:
                current_c += 1
            else:
                current_c -= 1
            grid[current_r, current_c] = 'O'
        
        # Move vertically
        while current_r != r2:
            if current_r < r2:
                current_r += 1
            else:
                current_r -= 1
            grid[current_r, current_c] = 'O'

    def _ensure_distributed_paths(self, grid: np.ndarray) -> None:
        """
        Ensure paths are distributed across different quadrants of the maze.
        """
        mid_r, mid_c = self.maze_size // 2, self.maze_size // 2
        
        # Check each quadrant
        quadrants = [
            (0, mid_r, 0, mid_c),  # Top-left
            (0, mid_r, mid_c, self.maze_size),  # Top-right
            (mid_r, self.maze_size, 0, mid_c),  # Bottom-left
            (mid_r, self.maze_size, mid_c, self.maze_size)  # Bottom-right
        ]
        
        for r_start, r_end, c_start, c_end in quadrants:
            # Count open cells in this quadrant
            open_count = 0
            for r in range(r_start, r_end):
                for c in range(c_start, c_end):
                    if grid[r, c] == 'O':
                        open_count += 1
            
            # If quadrant has too few open cells, add some
            quadrant_size = (r_end - r_start) * (c_end - c_start)
            min_open = max(1, quadrant_size // 8)  # At least 12.5% open
            
            if open_count < min_open:
                self._add_paths_to_quadrant(grid, r_start, r_end, c_start, c_end, min_open - open_count)

    def _add_paths_to_quadrant(self, grid: np.ndarray, r_start: int, r_end: int, 
                              c_start: int, c_end: int, paths_needed: int) -> None:
        """
        Add paths to a specific quadrant of the maze.
        """
        added = 0
        attempts = 0
        max_attempts = paths_needed * 10
        
        while added < paths_needed and attempts < max_attempts:
            attempts += 1
            r = random.randint(r_start + 1, r_end - 2)
            c = random.randint(c_start + 1, c_end - 2)
            
            if grid[r, c] == 'W':
                # Check if adjacent to an open cell
                adjacent_open = False
                for dr, dc in [(0, 1), (0, -1), (1, 0), (-1, 0)]:
                    nr, nc = r + dr, c + dc
                    if (0 <= nr < self.maze_size and 0 <= nc < self.maze_size and
                        grid[nr, nc] == 'O'):
                        adjacent_open = True
                        break
                
                if adjacent_open:
                    grid[r, c] = 'O'
                    added += 1

    def _get_corner_open_cells(self, open_cells: List[Tuple[int, int]]) -> List[Tuple[int, int]]:
        """
        Get open cells that are in corner or edge areas of the maze.
        """
        corner_cells = []
        edge_threshold = self.maze_size // 4
        
        for r, c in open_cells:
            # Check if cell is near edges
            near_edge = (r < edge_threshold or r >= self.maze_size - edge_threshold or
                        c < edge_threshold or c >= self.maze_size - edge_threshold)
            if near_edge:
                corner_cells.append((r, c))
        
        return corner_cells

    def _get_central_open_cells(self, open_cells: List[Tuple[int, int]]) -> List[Tuple[int, int]]:
        """
        Get open cells that are in the central area of the maze.
        """
        central_cells = []
        center_threshold = self.maze_size // 3
        center_r, center_c = self.maze_size // 2, self.maze_size // 2
        
        for r, c in open_cells:
            # Check if cell is within central area
            in_central = (abs(r - center_r) <= center_threshold and 
                         abs(c - center_c) <= center_threshold)
            if in_central:
                central_cells.append((r, c))
        
        return central_cells

    def _add_extra_complexity(self, grid: np.ndarray) -> None:
        """
        Add extra complexity for very hard mazes by creating more branching paths.
        """
        # Add some extra random paths in the maze for very hard difficulty
        attempts = self.maze_size // 3  # Scale with maze size
        added = 0
        max_attempts = attempts * 3
        
        for _ in range(max_attempts):
            if added >= attempts:
                break
                
            r = random.randint(2, self.maze_size - 3)
            c = random.randint(2, self.maze_size - 3)
            
            if grid[r, c] == 'W':
                # Check if this wall is between two open areas
                neighbors = [(r-1, c), (r+1, c), (r, c-1), (r, c+1)]
                open_neighbors = []
                
                for nr, nc in neighbors:
                    if (0 <= nr < self.maze_size and 0 <= nc < self.maze_size and
                        grid[nr, nc] == 'O'):
                        open_neighbors.append((nr, nc))
                
                # If we have at least 2 open neighbors, this could create a good connection
                if len(open_neighbors) >= 2:
                    grid[r, c] = 'O'
                    added += 1

    def get_maze(self) -> List[List[str]]:
        """
        Get the generated maze as a list of lists.
        
        Returns:
            List of lists containing maze symbols
        """
        return self.maze.tolist()

    def calculate_difficulty_score(self) -> float:
        """
        Calculate a difficulty score for the maze based on various factors.
        Enhanced with Shannon entropy-based complexity metrics.
        
        Returns:
            Difficulty score (0-100, where higher is more difficult)
        """
        # Get traditional analysis
        analysis = self.analyze_dead_ends()
        
        # Get Shannon entropy-based analysis
        shannon_analysis = self.calculate_shannon_complexity()
        
        # Get difficulty level (default to medium if not set)
        difficulty_level = getattr(self, 'difficulty_level', 'medium')
        
        # Base factors for difficulty (reduced weights to make room for Shannon metrics)
        dead_end_factor = min(analysis['dead_end_density'] * 100, 20)  
        avg_length_factor = min(analysis['average_dead_end_length'] * 1.5, 15)  
        total_dead_ends_factor = min(analysis['dead_end_count'] / self.maze_size * 8, 15) 
        longest_dead_end_factor = min(analysis['longest_dead_end'] * 1.0, 10) 
        
        # Size factor (reduced)
        size_factor = min(self.maze_size / 15, 5)  # Reduced from 10
        
        # Shannon entropy factors (new - up to 35 points total)
        surprisingness_factor = min(shannon_analysis['surprisingness'] * 10, 15)  
        deceptiveness_factor = min(shannon_analysis['deceptiveness'] * 8, 12)  
        trap_complexity_factor = min(shannon_analysis['trap_complexity_score'], 8) 
        
        # Apply difficulty level adjustments
        hard_bonus = 0  # Initialize hard bonus
        
        if difficulty_level == "easy":
            # Significantly reduce complexity factors for easy mazes
            trap_complexity_factor = min(trap_complexity_factor * 0.3, 2)  # Much more aggressive reduction
            surprisingness_factor *= 0.2  # Much more aggressive (was 0.4)
            deceptiveness_factor *= 0.2   # Much more aggressive (was 0.4)
            # Reduce traditional factors too
            dead_end_factor *= 0.4        # More aggressive (was 0.6)
            avg_length_factor *= 0.4      # More aggressive (was 0.6)
            total_dead_ends_factor *= 0.4 # More aggressive (was 0.6)
            # Add a negative bonus for easy mazes to push scores lower
            hard_bonus = -5
        elif difficulty_level == "medium":
            # Adjustments for medium difficulty to target ~55 average
            surprisingness_factor *= 0.8   # More reduction (was 0.85)
            deceptiveness_factor *= 0.8    # More reduction (was 0.85)
            trap_complexity_factor *= 0.85 # More reduction (was 0.9)
            # More reduction on traditional factors
            dead_end_factor *= 0.85        # More reduction (was 0.9)
            avg_length_factor *= 0.85      # More reduction (was 0.9)
            total_dead_ends_factor *= 0.85 # More reduction (was 0.9)
            # Smaller bonus to target ~55
            hard_bonus = 0  # No bonus (was 5)
        elif difficulty_level == "hard":
            # More conservative adjustments for hard difficulty to target ~63 average
            surprisingness_factor *= 0.95  # Slight reduction (was 1.02)
            deceptiveness_factor *= 0.95   # Slight reduction (was 1.02)
            trap_complexity_factor *= 0.98 # Very slight reduction (was 1.05)
            # Reduce the bonus significantly
            hard_bonus = -5  # Negative bonus to bring scores down (was 8 * size_adjustment)

        elif difficulty_level == "very_hard":
            # Moderate boost for very hard mazes to target ~82 average
            surprisingness_factor *= 1.15  # Reduced multiplier (was 1.3)
            deceptiveness_factor *= 1.15   # Reduced multiplier (was 1.3)
            trap_complexity_factor *= 1.1  # Reduced multiplier (was 1.25)
            # Apply smaller multipliers to traditional factors
            dead_end_factor *= 1.05        # Smaller boost (was 1.15)
            avg_length_factor *= 1.05      # Smaller boost (was 1.15)
            total_dead_ends_factor *= 1.05 # Smaller boost (was 1.15)
            # Smaller bonus for very hard mazes
            hard_bonus = 5  # Much smaller bonus (was max(18, 12 * size_adjustment))
        
        # Combine all factors
        base_difficulty = (dead_end_factor + avg_length_factor + total_dead_ends_factor + 
                          longest_dead_end_factor + size_factor + surprisingness_factor + 
                          deceptiveness_factor + trap_complexity_factor)
        
        # Add hard bonus
        difficulty = base_difficulty + hard_bonus
        
        return min(difficulty, 100)  # Cap at 100

    def calculate_shannon_complexity(self) -> Dict[str, float]:
        """
        Calculate Shannon entropy-based complexity metrics following AMaze framework.
        
        Returns:
            Dictionary containing:
            - surprisingness: S(M) based on input frequency during optimal trajectory
            - deceptiveness: D(M) based on cell-to-trap transitions
            - trap_complexity_score: Weighted trap complexity
            - trap_positions: List of trap positions
            - optimal_path: The optimal path from start to exit
        """
        try:
            # Find optimal path from start to exit
            start_pos = self._find_starting_position_internal()
            exit_pos = self._find_exit_position()
            
            if not start_pos or not exit_pos:
                return self._get_default_shannon_analysis()
            
            optimal_path = self._find_optimal_path(start_pos, exit_pos)
            if not optimal_path:
                return self._get_default_shannon_analysis()
            
            # Calculate surprisingness S(M)
            surprisingness = self._calculate_surprisingness(optimal_path)
            
            # Detect traps and calculate deceptiveness D(M)
            traps = self._detect_traps(optimal_path)
            deceptiveness = self._calculate_deceptiveness(traps)
            
            # Calculate trap complexity score
            trap_complexity = self._calculate_trap_complexity(traps, optimal_path)
            
            # Convert trap analysis to JSON-serializable format
            trap_analysis_serializable = {}
            for trap_pos, trap_info in traps.items():
                # Convert tuple position to string key for JSON serialization
                pos_key = f"{trap_pos[0]},{trap_pos[1]}"
                trap_analysis_serializable[pos_key] = trap_info
            
            return {
                'surprisingness': surprisingness,
                'deceptiveness': deceptiveness,
                'trap_complexity_score': trap_complexity,
                'trap_positions': list(traps.keys()) if traps else [],
                'optimal_path': optimal_path,
                'trap_analysis': trap_analysis_serializable
            }
            
        except Exception as e:
            print(f"Warning: Shannon complexity calculation failed: {e}")
            return self._get_default_shannon_analysis()
    
    def _get_default_shannon_analysis(self) -> Dict[str, float]:
        """Return default Shannon analysis when calculation fails."""
        return {
            'surprisingness': 0.0,
            'deceptiveness': 0.0,
            'trap_complexity_score': 0.0,
            'trap_positions': [],
            'optimal_path': [],
            'trap_analysis': {}
        }
    
    def _find_starting_position_internal(self) -> Optional[Tuple[int, int]]:
        """Find a suitable starting position in the maze."""
        # Look for open cells ('O'), preferring positions closer to corners
        candidates = []
        for r in range(1, self.maze.shape[0] - 1):
            for c in range(1, self.maze.shape[1] - 1):
                if self.maze[r, c] == 'O':
                    candidates.append((r, c))
        
        if not candidates:
            return None
        
        # Prefer corner positions for starting
        corner_candidates = []
        for r, c in candidates:
            if r < self.maze.shape[0] // 3 or r > 2 * self.maze.shape[0] // 3:
                if c < self.maze.shape[1] // 3 or c > 2 * self.maze.shape[1] // 3:
                    corner_candidates.append((r, c))
        
        return corner_candidates[0] if corner_candidates else candidates[0]
    
    def _find_exit_position(self) -> Optional[Tuple[int, int]]:
        """Find the exit position in the maze."""
        for r in range(self.maze.shape[0]):
            for c in range(self.maze.shape[1]):
                if self.maze[r, c] == 'E':
                    return (r, c)
        return None
    
    def _find_optimal_path(self, start: Tuple[int, int], exit: Tuple[int, int]) -> List[Tuple[int, int]]:
        """Find the shortest path from start to exit using BFS."""
        if start == exit:
            return [start]
        
        queue = deque([(start, [start])])
        visited = {start}
        
        directions = [(0, 1), (0, -1), (1, 0), (-1, 0)]
        
        while queue:
            (r, c), path = queue.popleft()
            
            for dr, dc in directions:
                nr, nc = r + dr, c + dc
                
                if (0 <= nr < self.maze.shape[0] and 0 <= nc < self.maze.shape[1] and
                    (nr, nc) not in visited and self.maze[nr, nc] in ['O', 'E']):
                    
                    new_path = path + [(nr, nc)]
                    
                    if (nr, nc) == exit:
                        return new_path
                    
                    queue.append(((nr, nc), new_path))
                    visited.add((nr, nc))
        
        return []  # No path found
    
    def _calculate_surprisingness(self, optimal_path: List[Tuple[int, int]]) -> float:
        """
        Calculate surprisingness S(M) based on Shannon's entropy of input frequencies.
        
        S(M) = -Σ p(i) * log₂(p(i))
        where p(i) is the observed frequency of input i during optimal trajectory.
        """
        if len(optimal_path) < 2:
            return 0.0
        
        # Calculate direction frequencies along optimal path
        direction_counts = Counter()
        
        for i in range(len(optimal_path) - 1):
            curr_pos = optimal_path[i]
            next_pos = optimal_path[i + 1]
            
            # Determine direction
            dr = next_pos[0] - curr_pos[0]
            dc = next_pos[1] - curr_pos[1]
            
            if dr == -1 and dc == 0:
                direction = 'north'
            elif dr == 1 and dc == 0:
                direction = 'south'
            elif dr == 0 and dc == 1:
                direction = 'east'
            elif dr == 0 and dc == -1:
                direction = 'west'
            else:
                continue  # Invalid move
            
            direction_counts[direction] += 1
        
        # Calculate Shannon entropy
        total_moves = sum(direction_counts.values())
        if total_moves == 0:
            return 0.0
        
        entropy = 0.0
        for count in direction_counts.values():
            if count > 0:
                p = count / total_moves
                entropy -= p * math.log2(p)
        
        # Normalize entropy (maximum entropy is log₂(4) = 2 for 4 directions)
        max_entropy = math.log2(min(4, len(direction_counts))) if direction_counts else 1
        return entropy / max_entropy if max_entropy > 0 else 0.0
    
    def _detect_traps(self, optimal_path: List[Tuple[int, int]]) -> Dict[Tuple[int, int], Dict[str, any]]:
        """
        Detect traps: positions that deviate from optimal path and lead to dead ends.
        
        Returns:
            Dictionary mapping trap positions to their analysis:
            - 'depth': How deep the trap is (path length to dead end)
            - 'branches': Number of sub-branches in the trap
            - 'weight': Calculated weight based on complexity
        """
        if not optimal_path:
            return {}
        
        optimal_path_set = set(optimal_path)
        traps = {}
        
        # Find all positions adjacent to optimal path that are not on the path
        for pos in optimal_path:
            r, c = pos
            directions = [(0, 1), (0, -1), (1, 0), (-1, 0)]
            
            for dr, dc in directions:
                nr, nc = r + dr, c + dc
                trap_pos = (nr, nc)
                
                # Skip if out of bounds, wall, or already on optimal path
                if (not (0 <= nr < self.maze.shape[0] and 0 <= nc < self.maze.shape[1]) or
                    self.maze[nr, nc] not in ['O', 'E'] or
                    trap_pos in optimal_path_set or
                    trap_pos in traps):
                    continue
                
                # Analyze this potential trap
                trap_analysis = self._analyze_trap_branch(trap_pos, optimal_path_set)
                if trap_analysis['is_trap']:
                    traps[trap_pos] = trap_analysis
        
        return traps
    
    def _analyze_trap_branch(self, start_pos: Tuple[int, int], 
                           optimal_path_set: Set[Tuple[int, int]]) -> Dict[str, any]:
        """
        Analyze a branch starting from start_pos to determine if it's a trap.
        
        Returns:
            Dictionary with trap analysis:
            - 'is_trap': Boolean indicating if this is a trap
            - 'depth': Maximum depth of the trap
            - 'branches': Number of branching points
            - 'dead_ends': Number of dead ends in this trap
            - 'weight': Calculated weight
        """
        visited = set()
        max_depth = 0
        branches = 0
        dead_ends = 0
        
        def dfs(pos: Tuple[int, int], depth: int) -> bool:
            nonlocal max_depth, branches, dead_ends
            
            if pos in visited or pos in optimal_path_set:
                return False
            
            r, c = pos
            if (not (0 <= r < self.maze.shape[0] and 0 <= c < self.maze.shape[1]) or
                self.maze[r, c] not in ['O', 'E']):
                return False
            
            visited.add(pos)
            max_depth = max(max_depth, depth)
            
            # Count valid neighbors (excluding visited and optimal path)
            valid_neighbors = []
            for dr, dc in [(0, 1), (0, -1), (1, 0), (-1, 0)]:
                nr, nc = r + dr, c + dc
                neighbor_pos = (nr, nc)
                
                if (0 <= nr < self.maze.shape[0] and 0 <= nc < self.maze.shape[1] and
                    self.maze[nr, nc] in ['O', 'E'] and
                    neighbor_pos not in visited and
                    neighbor_pos not in optimal_path_set):
                    valid_neighbors.append(neighbor_pos)
            
            if len(valid_neighbors) == 0:
                # This is a dead end
                dead_ends += 1
                return True
            elif len(valid_neighbors) > 1:
                # This is a branching point
                branches += 1
            
            # Continue exploring
            has_dead_end = False
            for neighbor in valid_neighbors:
                if dfs(neighbor, depth + 1):
                    has_dead_end = True
            
            return has_dead_end
        
        has_dead_end = dfs(start_pos, 1)
        
        # Calculate weight based on complexity
        weight = 1.0
        if has_dead_end:
            weight += max_depth * 0.5  # Deeper traps are more deceptive
            weight += branches * 0.3   # More branches increase complexity
            weight += dead_ends * 0.2  # Multiple dead ends increase difficulty
        
        return {
            'is_trap': has_dead_end,
            'depth': max_depth,
            'branches': branches,
            'dead_ends': dead_ends,
            'weight': weight,
            'cells_count': len(visited)
        }
    
    def _calculate_deceptiveness(self, traps: Dict[Tuple[int, int], Dict[str, any]]) -> float:
        """
        Calculate deceptiveness D(M) based on entropy of state transitions to traps.
        
        D(M) = ΣΣ -p(s|c) * log₂(p(s|c))
        where c is a cell and s is a trap state.
        """
        if not traps:
            return 0.0
        
        total_entropy = 0.0
        cells_analyzed = 0
        
        # For each cell that can reach traps, calculate transition entropy
        for trap_pos, trap_info in traps.items():
            if not trap_info['is_trap']:
                continue
            
            # Find all cells that can transition to this trap
            accessible_cells = self._find_cells_leading_to_trap(trap_pos)
            
            if not accessible_cells:
                continue
            
            # Calculate transition probabilities
            trap_transitions = 0
            total_transitions = 0
            
            for cell_pos in accessible_cells:
                r, c = cell_pos
                # Count total possible transitions from this cell
                possible_moves = 0
                trap_moves = 0
                
                for dr, dc in [(0, 1), (0, -1), (1, 0), (-1, 0)]:
                    nr, nc = r + dr, c + dc
                    if (0 <= nr < self.maze.shape[0] and 0 <= nc < self.maze.shape[1] and
                        self.maze[nr, nc] in ['O', 'E']):
                        possible_moves += 1
                        if (nr, nc) == trap_pos:
                            trap_moves += 1
                
                if possible_moves > 0:
                    trap_transitions += trap_moves
                    total_transitions += possible_moves
                    cells_analyzed += 1
            
            # Calculate entropy for this trap
            if total_transitions > 0:
                p_trap = trap_transitions / total_transitions
                p_non_trap = 1 - p_trap
                
                entropy = 0.0
                if p_trap > 0:
                    entropy -= p_trap * math.log2(p_trap)
                if p_non_trap > 0:
                    entropy -= p_non_trap * math.log2(p_non_trap)
                
                # Weight by trap complexity
                weighted_entropy = entropy * trap_info['weight']
                total_entropy += weighted_entropy
        
        # Normalize by number of cells analyzed
        return total_entropy / max(cells_analyzed, 1)
    
    def _find_cells_leading_to_trap(self, trap_pos: Tuple[int, int]) -> List[Tuple[int, int]]:
        """Find all cells that can directly transition to a trap position."""
        leading_cells = []
        trap_r, trap_c = trap_pos
        
        for dr, dc in [(0, 1), (0, -1), (1, 0), (-1, 0)]:
            r, c = trap_r + dr, trap_c + dc
            if (0 <= r < self.maze.shape[0] and 0 <= c < self.maze.shape[1] and
                self.maze[r, c] in ['O', 'E']):
                leading_cells.append((r, c))
        
        return leading_cells
    
    def _calculate_trap_complexity(self, traps: Dict[Tuple[int, int], Dict[str, any]], 
                                 optimal_path: List[Tuple[int, int]]) -> float:
        """
        Calculate overall trap complexity score based on trap weights and distribution.
        """
        if not traps:
            return 0.0
        
        total_weight = 0.0
        total_cells_in_traps = 0
        max_trap_depth = 0
        
        for trap_info in traps.values():
            if trap_info['is_trap']:
                total_weight += trap_info['weight']
                total_cells_in_traps += trap_info['cells_count']
                max_trap_depth = max(max_trap_depth, trap_info['depth'])
        
        # Calculate complexity score
        if len(optimal_path) == 0:
            return 0.0
        
        # Normalize by maze size and optimal path length
        maze_cells = np.sum((self.maze == 'O') | (self.maze == 'E'))
        trap_density = total_cells_in_traps / max(maze_cells, 1)
        path_trap_ratio = total_weight / max(len(optimal_path), 1)
        depth_factor = max_trap_depth / max(self.maze_size, 1)
        
        complexity_score = (trap_density * 3 + path_trap_ratio * 4 + depth_factor * 3)
        return min(complexity_score, 10.0)  # Cap at reasonable maximum
    
    def get_maze_as_image(self, cell_size: int = 20, show_grid: bool = True, dead_end_factor: float = None) -> plt.Figure:
        """
        Convert the maze into a matplotlib figure with enhanced Shannon complexity information.
        
        Args:
            cell_size: Size of each cell in pixels
            show_grid: Whether to show grid lines between cells
            dead_end_factor: Dead end factor used for generation (for title display)
            
        Returns:
            matplotlib Figure object representing the maze
        """
        # Create figure and axis with more space for text
        fig, ax = plt.subplots(figsize=(12, 10))
        
        # Calculate dimensions
        height, width = self.maze.shape
        
        # Create a grid of colored squares
        for y in range(height):
            for x in range(width):
                cell = self.maze[y, x]
                color = self.colors.get(cell, '#FFFFFF')  # Default to white if unknown
                
                # Draw rectangle for each cell
                rect = plt.Rectangle((x, height-y-1), 1, 1, 
                                   facecolor=color, 
                                   edgecolor='black' if show_grid else 'none')
                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([])
        
        # Get dead end analysis and difficulty score
        analysis = self.analyze_dead_ends()
        shannon_analysis = self.calculate_shannon_complexity()
        difficulty = self.calculate_difficulty_score()
        
        # Create title with maze info
        title = f'Maze {self.maze_size}x{self.maze_size} - UUID: {self.maze_uuid}'
        if dead_end_factor is not None:
            title += f' - Dead End Factor: {dead_end_factor:.2f}'
        ax.set_title(title, fontsize=14, fontweight='bold')
        
        # Enhanced statistics text box with Shannon metrics
        stats_text = f"""Dead Ends: {analysis['dead_end_count']} ({analysis['dead_end_density']:.1%} density)
Avg Length: {analysis['average_dead_end_length']:.1f} | Max: {analysis['longest_dead_end']} | Min: {analysis['shortest_dead_end']}
Total Dead End Cells: {analysis['total_dead_end_cells']}

SHANNON COMPLEXITY:
Surprisingness: {shannon_analysis['surprisingness']:.3f}
Deceptiveness: {shannon_analysis['deceptiveness']:.3f}
Trap Complexity: {shannon_analysis['trap_complexity_score']:.2f}
Traps Detected: {len(shannon_analysis['trap_positions'])}

Difficulty Score: {difficulty:.1f}/100"""
        
        # 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 difficulty color indicator
        if difficulty < 30:
            diff_color = '#00FF00'  # Green - Easy
            diff_text = 'EASY'
        elif difficulty < 60:
            diff_color = '#FFFF00'  # Yellow - Medium
            diff_text = 'MEDIUM'
        elif difficulty < 80:
            diff_color = '#FF8000'  # Orange - Hard
            diff_text = 'HARD'
        else:
            diff_color = '#FF0000'  # Red - Very Hard
            diff_text = 'VERY HARD'
        
        # Add difficulty indicator in bottom right
        ax.text(0.98, 0.02, f'{diff_text}\n{difficulty:.1f}/100', 
                transform=ax.transAxes,
                fontsize=12,
                fontweight='bold',
                verticalalignment='bottom',
                horizontalalignment='right',
                bbox=dict(boxstyle='round,pad=0.3', facecolor=diff_color, alpha=0.7, edgecolor='black'),
                color='black')
        
        return fig

    def save_maze_image(self, filename: str, cell_size: int = 20, show_grid: bool = True, dpi: int = 100, dead_end_factor: float = None) -> None:
        """
        Save the maze as an image file.
        
        Args:
            filename: Name of the output file (e.g., 'maze.png')
            cell_size: Size of each cell in pixels
            show_grid: Whether to show grid lines between cells
            dpi: Dots per inch for the output image
            dead_end_factor: Dead end factor used for generation (for title display)
        """
        fig = self.get_maze_as_image(cell_size, show_grid, dead_end_factor)
        fig.savefig(filename, dpi=dpi, bbox_inches='tight')
        plt.close(fig)

    def analyze_dead_ends(self) -> dict:
        """
        Analyze the dead ends in the generated maze.
        
        Returns:
            Dictionary containing dead end analysis:
            - 'dead_end_count': Number of dead ends
            - 'dead_end_positions': List of (row, col) positions of dead ends
            - 'dead_end_lengths': List of lengths for each dead end path
            - 'average_dead_end_length': Average length of dead end paths
            - 'longest_dead_end': Length of the longest dead end
            - 'shortest_dead_end': Length of the shortest dead end
        """
        # Find all open cells (excluding exit)
        open_cells = []
        exit_pos = None
        
        for r in range(self.maze.shape[0]):
            for c in range(self.maze.shape[1]):
                if self.maze[r, c] == 'O':
                    open_cells.append((r, c))
                elif self.maze[r, c] == 'E':
                    exit_pos = (r, c)
                    open_cells.append((r, c))  # Include exit in pathfinding
        
        # Find dead ends (cells with only one open neighbor)
        dead_ends = []
        for r, c in open_cells:
            if (r, c) == exit_pos:
                continue  # Skip exit position
                
            open_neighbors = 0
            for dr, dc in [(0, 1), (0, -1), (1, 0), (-1, 0)]:
                nr, nc = r + dr, c + dc
                if (0 <= nr < self.maze.shape[0] and 0 <= nc < self.maze.shape[1] and
                    self.maze[nr, nc] in ['O', 'E']):
                    open_neighbors += 1
            
            if open_neighbors == 1:
                dead_ends.append((r, c))
        
        # Calculate dead end path lengths
        dead_end_lengths = []
        for dead_end in dead_ends:
            length = self._calculate_dead_end_length(dead_end, open_cells)
            dead_end_lengths.append(length)
        
        # Compile analysis results
        analysis = {
            'dead_end_count': len(dead_ends),
            'dead_end_positions': dead_ends,
            'dead_end_lengths': dead_end_lengths,
            'average_dead_end_length': sum(dead_end_lengths) / len(dead_end_lengths) if dead_end_lengths else 0,
            'longest_dead_end': max(dead_end_lengths) if dead_end_lengths else 0,
            'shortest_dead_end': min(dead_end_lengths) if dead_end_lengths else 0,
            'total_dead_end_cells': sum(dead_end_lengths),
            'dead_end_density': len(dead_ends) / len(open_cells) if open_cells else 0
        }
        
        return analysis

    def _calculate_dead_end_length(self, dead_end_start: Tuple[int, int], open_cells: List[Tuple[int, int]]) -> int:
        """
        Calculate the length of a dead end path from its tip to where it connects to the main path.
        
        Args:
            dead_end_start: Starting position of the dead end (the tip)
            open_cells: List of all open cells in the maze
            
        Returns:
            Length of the dead end path
        """
        visited = set()
        current = dead_end_start
        length = 1
        
        while True:
            visited.add(current)
            r, c = current
            
            # Find unvisited open neighbors
            unvisited_neighbors = []
            for dr, dc in [(0, 1), (0, -1), (1, 0), (-1, 0)]:
                nr, nc = r + dr, c + dc
                if ((nr, nc) in open_cells and (nr, nc) not in visited and
                    0 <= nr < self.maze.shape[0] and 0 <= nc < self.maze.shape[1] and
                    self.maze[nr, nc] in ['O', 'E']):
                    unvisited_neighbors.append((nr, nc))
            
            if len(unvisited_neighbors) == 0:
                # No more unvisited neighbors, we've reached the end
                break
            elif len(unvisited_neighbors) == 1:
                # Continue along the dead end path
                current = unvisited_neighbors[0]
                length += 1
            else:
                # Multiple unvisited neighbors means we've reached a junction
                break
        
        return length

    def print_dead_end_analysis(self) -> None:
        """
        Print a detailed analysis of the maze's dead ends.
        """
        analysis = self.analyze_dead_ends()
        
        print(f"\n{'='*50}")
        print(f"DEAD END ANALYSIS - Maze {self.maze_size}x{self.maze_size}")
        print(f"UUID: {self.maze_uuid}")
        print(f"{'='*50}")
        print(f"Total dead ends: {analysis['dead_end_count']}")
        print(f"Dead end density: {analysis['dead_end_density']:.2%}")
        print(f"Total cells in dead ends: {analysis['total_dead_end_cells']}")
        
        if analysis['dead_end_lengths']:
            print(f"\nDead end length statistics:")
            print(f"  Average length: {analysis['average_dead_end_length']:.1f} cells")
            print(f"  Longest dead end: {analysis['longest_dead_end']} cells")
            print(f"  Shortest dead end: {analysis['shortest_dead_end']} cells")
            
            # Length distribution
            length_counts = {}
            for length in analysis['dead_end_lengths']:
                length_counts[length] = length_counts.get(length, 0) + 1
            
            print(f"\nDead end length distribution:")
            for length in sorted(length_counts.keys()):
                count = length_counts[length]
                percentage = (count / analysis['dead_end_count']) * 100
                print(f"  Length {length}: {count} dead ends ({percentage:.1f}%)")
        
        print(f"{'='*50}")

    def print_shannon_complexity_analysis(self) -> None:
        """
        Print a detailed analysis of the maze's Shannon complexity metrics.
        """
        shannon_analysis = self.calculate_shannon_complexity()
        
        print(f"\n{'='*60}")
        print(f"SHANNON COMPLEXITY ANALYSIS - Maze {self.maze_size}x{self.maze_size}")
        print(f"UUID: {self.maze_uuid}")
        print(f"{'='*60}")
        
        print(f"Surprisingness S(M): {shannon_analysis['surprisingness']:.4f}")
        print(f"  - Based on direction frequency entropy along optimal path")
        print(f"  - Higher values indicate more varied/surprising navigation")
        
        print(f"\nDeceptiveness D(M): {shannon_analysis['deceptiveness']:.4f}")
        print(f"  - Based on entropy of transitions from cells to traps")
        print(f"  - Higher values indicate more deceptive trap placement")
        
        print(f"\nTrap Complexity Score: {shannon_analysis['trap_complexity_score']:.2f}")
        print(f"  - Overall trap complexity based on depth and branching")
        print(f"  - Accounts for trap weight and distribution")
        
        print(f"\nDetected Traps: {len(shannon_analysis['trap_positions'])}")
        if shannon_analysis['trap_positions']:
            print(f"Trap positions: {shannon_analysis['trap_positions'][:5]}{'...' if len(shannon_analysis['trap_positions']) > 5 else ''}")
        
        print(f"\nOptimal Path Length: {len(shannon_analysis['optimal_path'])}")
        
        # Detailed trap analysis - handle string keys
        if shannon_analysis.get('trap_analysis'):
            print(f"\nDETAILED TRAP ANALYSIS:")
            trap_items = list(shannon_analysis['trap_analysis'].items())
            for i, (pos_key, trap_info) in enumerate(trap_items):
                if i >= 3:  # Show only first 3 traps for brevity
                    print(f"  ... and {len(trap_items) - 3} more traps")
                    break
                # Convert string key back to position format for display
                row, col = pos_key.split(',')
                trap_pos_display = f"({row}, {col})"
                print(f"  Trap at {trap_pos_display}:")
                print(f"    - Depth: {trap_info['depth']} cells")
                print(f"    - Branches: {trap_info['branches']}")
                print(f"    - Dead ends: {trap_info['dead_ends']}")
                print(f"    - Weight: {trap_info['weight']:.2f}")
                print(f"    - Total cells: {trap_info['cells_count']}")
        
        print(f"{'='*60}")

if __name__ == "__main__":
    # Create output directory if it doesn't exist
    output_dir = "maze_outputs"
    os.makedirs(output_dir, exist_ok=True)
    
    # Test configurations with Shannon complexity evaluation
    test_sizes = [15, 25]
    dead_end_factors = [0.1, 0.25, 0.4]  # Low, medium, high dead end probability
    
    for size in test_sizes:
        for dead_end_factor in dead_end_factors:
            print(f"\nGenerating maze of size {size}x{size} with dead end factor {dead_end_factor}")
            
            # Create maze generator with specific dead end factor
            generator = MazeGenerator(maze_size=size)
            generator.maze = generator.generate_maze(dead_end_factor=dead_end_factor)
            
            # Print UUID information
            print(f"Maze UUID: {generator.get_maze_uuid()}")
            
            # Get maze as list of lists
            maze_list = generator.get_maze()
            
            # Print basic maze properties
            print(f"Maze dimensions: {len(maze_list)}x{len(maze_list[0])}")
            print(f"Number of walls: {sum(row.count('W') for row in maze_list)}")
            print(f"Number of open paths: {sum(row.count('O') for row in maze_list)}")
            exit_positions = [(i, j) for i, row in enumerate(maze_list) for j, cell in enumerate(row) if cell == 'E']
            print(f"Exit position: {exit_positions[0] if exit_positions else 'None'}")
            
            # Perform traditional dead end analysis
            generator.print_dead_end_analysis()
            
            # Perform Shannon complexity analysis
            generator.print_shannon_complexity_analysis()
            
            # Save maze data with enhanced metrics
            maze_data_enhanced = {
                'uuid': generator.get_maze_uuid(),
                'size': generator.maze_size,
                'maze_array': generator.maze.tolist(),
                'dead_end_analysis': generator.analyze_dead_ends(),
                'shannon_complexity': generator.calculate_shannon_complexity(),
                'difficulty_score': generator.calculate_difficulty_score()
            }
            
            # Save enhanced data
            filename = f"maze_{generator.get_maze_uuid()}.json"
            filepath = os.path.join(output_dir, filename)
            with open(filepath, 'w') as f:
                json.dump(maze_data_enhanced, f, indent=2)
            print(f"Saved enhanced maze data to: {filepath}")
            
            # Save maze image with Shannon complexity info
            output_file = os.path.join(output_dir, f"maze_{generator.get_maze_uuid()}.png")
            generator.save_maze_image(output_file, cell_size=20, show_grid=True, dpi=150, dead_end_factor=dead_end_factor)
            print(f"Saved maze image to: {output_file}")
            
            # Print a small section of the maze for verification
            print("\nMaze preview (top-left corner):")
            for row in maze_list[:5]:
                print(''.join(row[:5]))
    
    print(f"\n{'='*70}")
    print("ENHANCED SHANNON COMPLEXITY MAZE GENERATION COMPLETE")
    print("New features:")
    print("✅ Shannon entropy-based surprisingness calculation")
    print("✅ Trap detection and deceptiveness analysis") 
    print("✅ Weighted trap complexity scoring")
    print("✅ Enhanced difficulty calculation with entropy metrics")
    print("✅ Detailed trap analysis and reporting")
    print(f"{'='*70}")
