"""
JS-PSO Hybrid Algorithm
=======================

Hybrid algorithm combining Jellyfish Search (JS) with Particle Swarm
Optimization (PSO) for enhanced exploration and exploitation balance.

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

import math
import numpy as np
from ...base import BaseOptimizer
from ..levy_flight_universal import add_levy_flight_to_position


class JSPSOHybrid(BaseOptimizer):
    """
    Jellyfish Search + PSO Hybrid Algorithm
    
    Combines:
    - JS: Ocean current following and swarm behavior
    - PSO: Velocity-based movement and personal/global best tracking
    
    Parameters
    ----------
    population_size : int, default=30
        Size of the hybrid swarm
    max_iterations : int, default=100
        Maximum iterations
    w : float, default=0.7
        PSO inertia weight
    c1 : float, default=1.5
        PSO cognitive parameter
    c2 : float, default=1.5
        PSO social parameter
    c0 : float, default=0.5
        JS control parameter
    """
    
    def __init__(self, population_size=30, max_iterations=100, 
                 w=0.7, c1=1.5, c2=1.5, c0=0.5, **kwargs):
        super().__init__(population_size, max_iterations, **kwargs)
        self.algorithm_name = "JS-PSO Hybrid"
        self.aliases = ["js_pso", "jellyfish_pso_hybrid"]
        self.w = w
        self.c1 = c1
        self.c2 = c2
        self.c0 = c0
    
    def _optimize(self, objective_function, X=None, y=None, **kwargs):
        # Levy flight available via: add_levy_flight_to_position()
        # Get dimensions and bounds
        if X is not None:
            dimension = X.shape[1]
            lb, ub = 0.0, 1.0
        else:
            dimension = kwargs.get('dimensions', getattr(self, 'dimensions_', 10))
            lb = kwargs.get('lower_bound', getattr(self, 'lower_bound_', -100.0))
            ub = kwargs.get('upper_bound', getattr(self, 'upper_bound_', 100.0))
            if hasattr(lb, '__getitem__'):
                lb = lb[0]
            if hasattr(ub, '__getitem__'):
                ub = ub[0]
        
        
        # Initialize positions and velocities
        positions = np.random.uniform(lb, ub, (self.population_size_, dimension))
        velocities = np.random.uniform(-1, 1, (self.population_size_, dimension))
        fitness = np.array([objective_function(p) for p in positions])
        
        # Personal best
        pbest = positions.copy()
        pbest_fitness = fitness.copy()
        
        # Global best
        best_idx = np.argmin(fitness)
        gbest = positions[best_idx].copy()
        best_fitness = fitness[best_idx]
        global_fitness = []
        local_fitness = []
        local_positions = []
        
        for iteration in range(self.max_iterations_):
            # JS time control parameter
            c = self.c0 * (1 - iteration / self.max_iterations_)
            
            for i in range(self.population_size_):
                # Hybrid strategy selection
                if np.random.rand() < 0.5:
                    # JS component
                    if np.random.rand() > (1 - c):
                        # Type A: Follow ocean current (global best)
                        trend = gbest - positions[i]
                        mu = np.random.rand()
                        positions[i] = positions[i] + np.random.rand() * trend * mu
                    else:
                        # Type B: Move within swarm
                        j = np.random.randint(0, self.population_size_)
                        if fitness[i] < fitness[j]:
                            direction = positions[i] - positions[j]
                        else:
                            direction = positions[j] - positions[i]
                        positions[i] = positions[i] + np.random.rand() * direction
                else:
                    # PSO component
                    r1, r2 = np.random.rand(), np.random.rand()
                    velocities[i] = (self.w * velocities[i] + 
                                    self.c1 * r1 * (pbest[i] - positions[i]) +
                                    self.c2 * r2 * (gbest - positions[i]))
                    positions[i] = positions[i] + velocities[i]
                
                # Apply Levy flight in exploitation phase
                if iteration > self.max_iterations_ * 0.7:
                    levy_step = self._levy_flight(dimension)
                    positions[i] = positions[i] + levy_step * (gbest - positions[i])
                
                # Boundary control
                positions[i] = np.clip(positions[i], lb, ub)
                
                # Evaluate
                new_fitness = objective_function(positions[i])
                
                if new_fitness < fitness[i]:
                    fitness[i] = new_fitness
                    if new_fitness < pbest_fitness[i]:
                        pbest[i] = positions[i].copy()
                        pbest_fitness[i] = new_fitness
                    if new_fitness < best_fitness:
                        gbest = positions[i].copy()
                        best_fitness = new_fitness
            
            global_fitness.append(best_fitness)
            # Note: Add local_fitness and local_positions tracking as needed
            
            if self.verbose_ and iteration % 10 == 0:
                print(f"JS-PSO Iteration {iteration}: Best Fitness = {best_fitness:.6e}")
        
        return gbest, best_fitness, global_fitness, local_fitness, local_positions
    
    def _levy_flight(self, dimension):
        """Generate Levy flight step"""
        beta = 1.5
        sigma = (math.gamma(1 + beta) * np.sin(np.pi * beta / 2) / 
                (math.gamma((1 + beta) / 2) * beta * 2 ** ((beta - 1) / 2))) ** (1 / beta)
        u = np.random.randn(dimension) * sigma
        v = np.random.randn(dimension)
        step = u / (np.abs(v) ** (1 / beta))
        return 0.01 * step
