"""
Universal Levy Flight Engine for All MHA Algorithms
===================================================

This module provides Levy flight capabilities for both exploration and exploitation
phases in all metaheuristic algorithms. It can be integrated with any algorithm to
enhance its search capabilities.

Levy flights are random walks whose step lengths have a Levy distribution, which
enables long jumps (exploration) while maintaining local search (exploitation).

Author: MHA Flow Development Team
License: MIT
"""

import numpy as np
from scipy.stats import levy_stable
from typing import Optional, Tuple, Union
import math


class LevyFlightEngine:
    """
    Universal Levy Flight implementation that can be integrated with any MHA algorithm.
    
    Features:
    - Adaptive Levy flight for both exploration and exploitation
    - Parameter control based on iteration progress
    - Multiple Levy flight strategies
    - Boundary handling
    - Performance monitoring
    """
    
    def __init__(self, 
                 beta: float = 1.5,
                 sigma_factor: float = 0.01,
                 strategy: str = 'mantegna',
                 adaptive: bool = True):
        """
        Initialize Levy Flight Engine.
        
        Parameters
        ----------
        beta : float, default=1.5
            Levy exponent (0 < beta <= 2). Controls heavy tail behavior.
            - beta = 1.5: Good balance for most problems
            - beta → 2: Approaches Gaussian
            - beta → 0: More heavy-tailed
        sigma_factor : float, default=0.01
            Scaling factor for step size
        strategy : str, default='mantegna'
            Levy flight calculation strategy:
            - 'mantegna': Mantegna's algorithm (fast, recommended)
            - 'scipy': Using scipy.stats.levy_stable (precise)
            - 'simplified': Simplified power-law (fastest)
        adaptive : bool, default=True
            Whether to adapt parameters during optimization
        """
        self.beta = beta
        self.sigma_factor = sigma_factor
        self.strategy = strategy
        self.adaptive = adaptive
        self.iteration = 0
        self.max_iterations = 100
        
        # Pre-calculate sigma for Mantegna algorithm
        if strategy == 'mantegna':
            self._precalculate_mantegna_sigma()
    
    def _precalculate_mantegna_sigma(self):
        """Pre-calculate sigma value for Mantegna algorithm"""
        numerator = math.gamma(1 + self.beta) * np.sin(np.pi * self.beta / 2)
        denominator = math.gamma((1 + self.beta) / 2) * self.beta * (2 ** ((self.beta - 1) / 2))
        self.sigma_mantegna = (numerator / denominator) ** (1 / self.beta)
    
    def set_iteration_context(self, current_iter: int, max_iter: int):
        """
        Set iteration context for adaptive behavior.
        
        Parameters
        ----------
        current_iter : int
            Current iteration number
        max_iter : int
            Maximum number of iterations
        """
        self.iteration = current_iter
        self.max_iterations = max_iter
    
    def get_step_size(self, dimensions: int) -> np.ndarray:
        """
        Calculate Levy flight step size.
        
        Parameters
        ----------
        dimensions : int
            Problem dimensionality
            
        Returns
        -------
        step : np.ndarray
            Levy flight step vector
        """
        if self.strategy == 'mantegna':
            return self._mantegna_step(dimensions)
        elif self.strategy == 'scipy':
            return self._scipy_step(dimensions)
        elif self.strategy == 'simplified':
            return self._simplified_step(dimensions)
        else:
            raise ValueError(f"Unknown strategy: {self.strategy}")
    
    def _mantegna_step(self, dimensions: int) -> np.ndarray:
        """Mantegna's algorithm for Levy flight"""
        u = np.random.randn(dimensions) * self.sigma_mantegna
        v = np.random.randn(dimensions)
        step = u / (np.abs(v) ** (1 / self.beta))
        
        # Apply adaptive scaling
        if self.adaptive:
            scale = self._get_adaptive_scale()
            step *= scale
        
        return step * self.sigma_factor
    
    def _scipy_step(self, dimensions: int) -> np.ndarray:
        """Use scipy's levy_stable for precise Levy flight"""
        step = levy_stable.rvs(alpha=self.beta, beta=0, size=dimensions)
        
        if self.adaptive:
            scale = self._get_adaptive_scale()
            step *= scale
        
        return step * self.sigma_factor
    
    def _simplified_step(self, dimensions: int) -> np.ndarray:
        """Simplified power-law Levy flight (fastest)"""
        u = np.random.rand(dimensions)
        step = (u ** (-1 / self.beta)) - 1
        
        if self.adaptive:
            scale = self._get_adaptive_scale()
            step *= scale
        
        return step * self.sigma_factor
    
    def _get_adaptive_scale(self) -> float:
        """
        Calculate adaptive scaling factor based on iteration progress.
        
        Early iterations: Larger steps (exploration)
        Later iterations: Smaller steps (exploitation)
        """
        progress = self.iteration / self.max_iterations
        # Exponential decay from exploration to exploitation
        scale = np.exp(-2 * progress)
        return max(0.1, scale)  # Minimum scale of 0.1
    
    def apply_levy_flight(self,
                         position: np.ndarray,
                         best_position: np.ndarray,
                         bounds: Tuple[float, float],
                         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 (for exploitation)
        bounds : tuple
            Search space bounds (min, max)
        exploration_mode : bool
            If True, emphasize exploration; if False, emphasize exploitation
            
        Returns
        -------
        new_position : np.ndarray
            New position after Levy flight
        """
        dimensions = len(position)
        levy_step = self.get_step_size(dimensions)
        
        if exploration_mode:
            # Exploration: Random direction with Levy flight
            direction = np.random.randn(dimensions)
            direction = direction / (np.linalg.norm(direction) + 1e-10)
            new_position = position + levy_step * direction
        else:
            # Exploitation: Toward best position with Levy flight
            direction = best_position - position
            direction = direction / (np.linalg.norm(direction) + 1e-10)
            new_position = position + levy_step * direction
        
        # Clip to bounds
        new_position = np.clip(new_position, bounds[0], bounds[1])
        
        return new_position
    
    def enhance_population_diversity(self,
                                    population: np.ndarray,
                                    best_solution: np.ndarray,
                                    bounds: Tuple[float, float],
                                    diversity_rate: float = 0.2) -> 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
            Search space bounds
        diversity_rate : float
            Fraction of population to diversify
            
        Returns
        -------
        new_population : np.ndarray
            Enhanced population
        """
        n_individuals = len(population)
        n_to_diversify = int(n_individuals * diversity_rate)
        
        # Select random individuals to diversify
        indices = np.random.choice(n_individuals, n_to_diversify, replace=False)
        
        new_population = population.copy()
        for idx in indices:
            new_population[idx] = self.apply_levy_flight(
                population[idx],
                best_solution,
                bounds,
                exploration_mode=True
            )
        
        return new_population


class AdaptiveLevyFlight(LevyFlightEngine):
    """
    Advanced adaptive Levy flight with multiple strategies.
    
    Features:
    - Auto-switching between exploration and exploitation
    - Performance-based adaptation
    - Multi-strategy approach
    """
    
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.no_improvement_count = 0
        self.best_fitness_history = []
        self.stagnation_threshold = 10
    
    def update_performance(self, current_best_fitness: float):
        """
        Update performance tracking for adaptive behavior.
        
        Parameters
        ----------
        current_best_fitness : float
            Current best fitness value
        """
        if len(self.best_fitness_history) == 0:
            self.best_fitness_history.append(current_best_fitness)
            self.no_improvement_count = 0
        else:
            if abs(current_best_fitness - self.best_fitness_history[-1]) < 1e-6:
                self.no_improvement_count += 1
            else:
                self.no_improvement_count = 0
            self.best_fitness_history.append(current_best_fitness)
    
    def is_stagnant(self) -> bool:
        """Check if optimization is stagnant"""
        return self.no_improvement_count >= self.stagnation_threshold
    
    def get_adaptive_step(self, dimensions: int, force_exploration: bool = False) -> np.ndarray:
        """
        Get adaptive step size based on performance.
        
        Parameters
        ----------
        dimensions : int
            Problem dimensionality
        force_exploration : bool
            Force exploration mode if stagnant
            
        Returns
        -------
        step : np.ndarray
            Adaptive Levy flight step
        """
        # If stagnant, increase exploration
        if self.is_stagnant() or force_exploration:
            original_sigma = self.sigma_factor
            self.sigma_factor *= 2.0  # Increase step size for exploration
            step = self.get_step_size(dimensions)
            self.sigma_factor = original_sigma
            return step
        else:
            return self.get_step_size(dimensions)


def integrate_levy_flight(algorithm_instance,
                         position: np.ndarray,
                         best_position: np.ndarray,
                         bounds: Tuple[float, float],
                         iteration: int,
                         max_iterations: int,
                         exploration_rate: float = 0.5) -> np.ndarray:
    """
    Universal function to integrate Levy flight into any algorithm.
    
    This function can be called from any MHA algorithm to add Levy flight capability.
    
    Parameters
    ----------
    algorithm_instance : object
        Instance of the algorithm (can be any MHA)
    position : np.ndarray
        Current position
    best_position : np.ndarray
        Best known position
    bounds : tuple
        Search space bounds
    iteration : int
        Current iteration
    max_iterations : int
        Maximum iterations
    exploration_rate : float
        Balance between exploration (1.0) and exploitation (0.0)
        
    Returns
    -------
    new_position : np.ndarray
        Position after Levy flight
    
    Example
    -------
    >>> # In any algorithm's optimize method:
    >>> from mha_toolbox.levy_flight_engine import integrate_levy_flight
    >>> new_position = integrate_levy_flight(
    ...     self, current_pos, best_pos, bounds, iter, max_iter
    ... )
    """
    # Create or reuse Levy engine
    if not hasattr(algorithm_instance, '_levy_engine'):
        algorithm_instance._levy_engine = AdaptiveLevyFlight(
            beta=1.5,
            strategy='mantegna',
            adaptive=True
        )
    
    levy_engine = algorithm_instance._levy_engine
    levy_engine.set_iteration_context(iteration, max_iterations)
    
    # Decide between exploration and exploitation
    is_exploration = np.random.rand() < exploration_rate
    
    return levy_engine.apply_levy_flight(
        position, best_position, bounds, is_exploration
    )


# Convenience functions for common use cases

def levy_exploration(position: np.ndarray, 
                     bounds: Tuple[float, float],
                     step_size: float = 0.01) -> np.ndarray:
    """Quick exploration using Levy flight"""
    engine = LevyFlightEngine(sigma_factor=step_size, strategy='mantegna')
    best_pos = np.mean(bounds) * np.ones_like(position)  # Use center as reference
    return engine.apply_levy_flight(position, best_pos, bounds, exploration_mode=True)


def levy_exploitation(position: np.ndarray,
                      target: np.ndarray,
                      bounds: Tuple[float, float],
                      step_size: float = 0.005) -> np.ndarray:
    """Quick exploitation toward target using Levy flight"""
    engine = LevyFlightEngine(sigma_factor=step_size, strategy='mantegna')
    return engine.apply_levy_flight(position, target, bounds, exploration_mode=False)


if __name__ == "__main__":
    # Demonstration
    print("=" * 60)
    print("Universal Levy Flight Engine - Demonstration")
    print("=" * 60)
    
    # Create engine
    engine = AdaptiveLevyFlight(beta=1.5, strategy='mantegna', adaptive=True)
    
    # Simulate optimization
    dimensions = 10
    bounds = (-100, 100)
    position = np.random.uniform(bounds[0], bounds[1], dimensions)
    best_position = np.zeros(dimensions)
    
    print(f"\nInitial position: {position[:3]}...")
    print(f"Best position: {best_position[:3]}...")
    
    # Simulate iterations
    for i in range(20):
        engine.set_iteration_context(i, 20)
        
        # Apply Levy flight
        if i < 10:
            # Early: Exploration
            new_position = engine.apply_levy_flight(
                position, best_position, bounds, exploration_mode=True
            )
            mode = "Exploration"
        else:
            # Late: Exploitation
            new_position = engine.apply_levy_flight(
                position, best_position, bounds, exploration_mode=False
            )
            mode = "Exploitation"
        
        distance_moved = np.linalg.norm(new_position - position)
        distance_to_best = np.linalg.norm(new_position - best_position)
        
        if i % 5 == 0:
            print(f"\nIteration {i} ({mode}):")
            print(f"  Moved: {distance_moved:.4f}")
            print(f"  Distance to best: {distance_to_best:.4f}")
        
        position = new_position
    
    print("\n" + "=" * 60)
    print("Demonstration Complete!")
    print("=" * 60)
