import numpy as np
import math
from typing import Tuple, Optional


class LevyFlightMixin:
    """
    Mixin class to add Levy flight capabilities to any optimization algorithm.
    
    Simply inherit from this mixin in your algorithm class to gain Levy flight support.
    
    Example:
        class MyAlgorithm(LevyFlightMixin, BaseOptimizer):
            def optimize(self, ...):
                # Use Levy flight
                new_position = self.levy_flight_step(position, best, bounds, iteration)
    """
    
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._levy_beta = kwargs.get('levy_beta', 1.5)
        self._levy_enabled = kwargs.get('levy_enabled', True)
        self._levy_frequency = kwargs.get('levy_frequency', 5)  # Apply every N iterations
    
    def levy_flight_step(self, 
                        position: np.ndarray, 
                        best_position: np.ndarray, 
                        bounds: Tuple[float, float],
                        iteration: int,
                        exploration_mode: bool = True) -> np.ndarray:
        """
        Apply Levy flight to a position.
        
        Parameters
        ----------
        position : np.ndarray
            Current position
        best_position : np.ndarray
            Best known position
        bounds : tuple
            (lower_bound, upper_bound)
        iteration : int
            Current iteration number
        exploration_mode : bool
            If True, emphasize exploration; if False, emphasize exploitation
        
        Returns
        -------
        new_position : np.ndarray
            Position after Levy flight
        """
        if not self._levy_enabled or iteration % self._levy_frequency != 0:
            return position
        
        dimensions = len(position)
        levy_step = self._calculate_levy_step(dimensions)
        
        # Adaptive step size based on iteration
        progress = iteration / getattr(self, 'max_iterations', 100)
        adaptive_factor = np.exp(-2 * progress)  # Decreases over time
        
        if exploration_mode:
            # Random direction for exploration
            direction = np.random.randn(dimensions)
            direction = direction / (np.linalg.norm(direction) + 1e-10)
            new_position = position + adaptive_factor * levy_step * direction
        else:
            # Toward best for exploitation
            direction = best_position - position
            direction = direction / (np.linalg.norm(direction) + 1e-10)
            new_position = position + adaptive_factor * levy_step * direction
        
        # Clip to bounds
        lb, ub = bounds
        new_position = np.clip(new_position, lb, ub)
        
        return new_position
    
    def _calculate_levy_step(self, dimensions: int) -> np.ndarray:
        """Calculate Levy flight step using Mantegna's algorithm"""
        beta = self._levy_beta
        
        # Calculate sigma
        numerator = math.gamma(1 + beta) * np.sin(np.pi * beta / 2)
        denominator = math.gamma((1 + beta) / 2) * beta * (2 ** ((beta - 1) / 2))
        sigma = (numerator / denominator) ** (1 / beta)
        
        # Generate step
        u = np.random.randn(dimensions) * sigma
        v = np.random.randn(dimensions)
        step = u / (np.abs(v) ** (1 / beta))
        
        return 0.01 * step
    
    def should_apply_levy(self, iteration: int) -> bool:
        """Check if Levy flight should be applied in this iteration"""
        return self._levy_enabled and (iteration % self._levy_frequency == 0)


def add_levy_flight_to_position(position: np.ndarray,
                                best_position: np.ndarray,
                                bounds: Tuple[float, float],
                                iteration: int,
                                max_iterations: int,
                                beta: float = 1.5) -> np.ndarray:
    """
    Standalone function to add Levy flight to any position.
    
    This can be called from any algorithm without inheriting from LevyFlightMixin.
    
    Parameters
    ----------
    position : np.ndarray
        Current position
    best_position : np.ndarray
        Best known position
    bounds : tuple
        (lower_bound, upper_bound)
    iteration : int
        Current iteration
    max_iterations : int
        Maximum iterations
    beta : float
        Levy exponent (default: 1.5)
    
    Returns
    -------
    new_position : np.ndarray
        Position after Levy flight
    
    Example
    -------
    >>> from mha_toolbox.algorithms.levy_flight_universal import add_levy_flight_to_position
    >>> new_pos = add_levy_flight_to_position(current_pos, best_pos, bounds, iter, max_iter)
    """
    dimensions = len(position)
    
    # Calculate Levy step
    numerator = math.gamma(1 + beta) * np.sin(np.pi * beta / 2)
    denominator = math.gamma((1 + beta) / 2) * beta * (2 ** ((beta - 1) / 2))
    sigma = (numerator / denominator) ** (1 / beta)
    
    u = np.random.randn(dimensions) * sigma
    v = np.random.randn(dimensions)
    levy_step = u / (np.abs(v) ** (1 / beta))
    levy_step = 0.01 * levy_step
    
    # Adaptive scaling
    progress = iteration / max_iterations
    adaptive_factor = np.exp(-2 * progress)
    
    # Apply toward best position
    direction = best_position - position
    direction = direction / (np.linalg.norm(direction) + 1e-10)
    new_position = position + adaptive_factor * levy_step * direction
    
    # Clip to bounds
    lb, ub = bounds
    new_position = np.clip(new_position, lb, ub)
    
    return new_position


def enhance_population_with_levy(population: np.ndarray,
                                 best_solution: np.ndarray,
                                 bounds: Tuple[float, float],
                                 diversity_rate: float = 0.2,
                                 beta: float = 1.5) -> np.ndarray:
    """
    Enhance population diversity using Levy flight.
    
    Parameters
    ----------
    population : np.ndarray
        Current population (n_individuals × n_dimensions)
    best_solution : np.ndarray
        Best solution found
    bounds : tuple
        (lower_bound, upper_bound)
    diversity_rate : float
        Fraction of population to diversify (0 to 1)
    beta : float
        Levy exponent
    
    Returns
    -------
    enhanced_population : np.ndarray
        Population with enhanced diversity
    """
    n_individuals, n_dimensions = population.shape
    n_to_enhance = int(n_individuals * diversity_rate)
    
    # Select random individuals
    indices = np.random.choice(n_individuals, n_to_enhance, replace=False)
    
    enhanced_pop = population.copy()
    
    for idx in indices:
        # Calculate Levy step
        numerator = math.gamma(1 + beta) * np.sin(np.pi * beta / 2)
        denominator = math.gamma((1 + beta) / 2) * beta * (2 ** ((beta - 1) / 2))
        sigma = (numerator / denominator) ** (1 / beta)
        
        u = np.random.randn(n_dimensions) * sigma
        v = np.random.randn(n_dimensions)
        levy_step = u / (np.abs(v) ** (1 / beta))
        levy_step = 0.01 * levy_step
        
        # Random exploration
        direction = np.random.randn(n_dimensions)
        direction = direction / (np.linalg.norm(direction) + 1e-10)
        
        enhanced_pop[idx] = population[idx] + levy_step * direction
        
        # Clip to bounds
        lb, ub = bounds
        enhanced_pop[idx] = np.clip(enhanced_pop[idx], lb, ub)
    
    return enhanced_pop


# Decorator to automatically add Levy flight to optimize method
def with_levy_flight(frequency: int = 5, beta: float = 1.5):
    """
    Decorator to automatically add Levy flight to an algorithm's optimize method.
    
    Parameters
    ----------
    frequency : int
        Apply Levy flight every N iterations
    beta : float
        Levy exponent
    
    Example
    -------
    >>> @with_levy_flight(frequency=10)
    >>> def optimize(self, objective_func, bounds, dimensions):
    >>>     # Your optimization code
    >>>     pass
    """
    def decorator(optimize_func):
        def wrapper(self, *args, **kwargs):
            # Store Levy parameters
            self._levy_frequency = frequency
            self._levy_beta = beta
            self._levy_enabled = True
            
            # Call original optimize
            return optimize_func(self, *args, **kwargs)
        
        return wrapper
    return decorator


if __name__ == "__main__":
    print("=" * 70)
    print("Universal Levy Flight Integration System")
    print("=" * 70)
    
    # Test the standalone function
    print("\n[Test] Standalone Levy Flight Function")
    dimensions = 10
    bounds = (-100, 100)
    current_pos = np.random.uniform(bounds[0], bounds[1], dimensions)
    best_pos = np.zeros(dimensions)
    
    print(f"Current position: {current_pos[:3]}...")
    new_pos = add_levy_flight_to_position(current_pos, best_pos, bounds, 10, 100)
    print(f"After Levy flight: {new_pos[:3]}...")
    print(f"Distance moved: {np.linalg.norm(new_pos - current_pos):.4f}")
    
    # Test population enhancement
    print("\n[Test] Population Enhancement")
    population = np.random.uniform(bounds[0], bounds[1], (30, dimensions))
    enhanced = enhance_population_with_levy(population, best_pos, bounds, diversity_rate=0.3)
    print(f"Population size: {population.shape}")
    print(f"Enhanced {int(30 * 0.3)} individuals")
    
    print("\n" + "=" * 70)
    print("✓ Universal Levy Flight System Ready!")
    print("  - Works with any algorithm")
    print("  - No code modification needed")
    print("  - Simply import and use")
    print("=" * 70)
