"""
Chimp Optimization Algorithm (ChOA)
====================================

Chimp Optimization Algorithm inspired by the individual intelligence
and sexual motivation of chimps in their group hunting.

Reference:
Khishe, M., & Mosavi, M. R. (2020). Chimp optimization algorithm.
Expert systems with applications, 149, 113338.

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 ChimpOptimizationAlgorithm(BaseOptimizer):
    """
    Chimp Optimization Algorithm (ChOA)
    
    Based on chimp social behavior and hunting strategies with four types:
    - Attacker: Best solution
    - Barrier: Second best
    - Chaser: Third best
    - Driver: Fourth best
    
    Parameters
    ----------
    population_size : int, default=30
        Number of chimps in the troop
    max_iterations : int, default=100
        Maximum number of iterations
    f : float, default=2.5
        Chaotic value parameter
    m : float, default=2
        Multiplier parameter
    """
    
    def __init__(self, population_size=30, max_iterations=100, f=2.5, m=2, **kwargs):
        super().__init__(population_size, max_iterations, **kwargs)
        self.algorithm_name = "Chimp Optimization Algorithm"
        self.aliases = ["choa", "chimp", "chimp_optimization"]
        self.f = f
        self.m = m
    
    def _optimize(self, objective_function, X=None, y=None, **kwargs):
        # Levy flight available via: add_levy_flight_to_position()
        """
        Execute the Chimp Optimization Algorithm
        
        Args:
            objective_function: Function to optimize
            X: Feature matrix (optional)
            y: Target values (optional)
            **kwargs: Additional parameters
            
        Returns:
            Tuple of (best_solution, best_fitness, global_fitness, local_fitness, local_positions)
        """
        # 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 chimps
        chimps = np.random.uniform(lb, ub, (self.population_size_, dimension))
        fitness = np.array([objective_function(chimp) for chimp in chimps])
        
        # Identify alpha, beta, gamma, delta chimps
        sorted_idx = np.argsort(fitness)
        attacker = chimps[sorted_idx[0]].copy()
        barrier = chimps[sorted_idx[1]].copy() if len(sorted_idx) > 1 else attacker.copy()
        chaser = chimps[sorted_idx[2]].copy() if len(sorted_idx) > 2 else attacker.copy()
        driver = chimps[sorted_idx[3]].copy() if len(sorted_idx) > 3 else attacker.copy()
        
        best_fitness = fitness[sorted_idx[0]]
        global_fitness = []
        local_fitness = []
        local_positions = []
        
        for iteration in range(self.max_iterations_):
            # Update a parameter (decreases linearly from 2 to 0)
            a = 2 * (1 - iteration / self.max_iterations_)
            
            for i in range(self.population_size_):
                # Generate chaotic values
                c1 = np.random.rand(dimension) * self.f
                c2 = np.random.rand(dimension) * self.f
                c3 = np.random.rand(dimension) * self.f
                c4 = np.random.rand(dimension) * self.f
                
                # Calculate coefficients for attacker
                r1 = np.random.rand(dimension)
                r2 = np.random.rand(dimension)
                A1 = 2 * a * r1 - a
                C1 = 2 * c1
                
                # Position update with respect to attacker
                D_attacker = np.abs(C1 * attacker - self.m * chimps[i])
                X1 = attacker - A1 * D_attacker
                
                # Position update with respect to barrier
                r3 = np.random.rand(dimension)
                r4 = np.random.rand(dimension)
                A2 = 2 * a * r3 - a
                C2 = 2 * c2
                D_barrier = np.abs(C2 * barrier - self.m * chimps[i])
                X2 = barrier - A2 * D_barrier
                
                # Position update with respect to chaser
                r5 = np.random.rand(dimension)
                r6 = np.random.rand(dimension)
                A3 = 2 * a * r5 - a
                C3 = 2 * c3
                D_chaser = np.abs(C3 * chaser - self.m * chimps[i])
                X3 = chaser - A3 * D_chaser
                
                # Position update with respect to driver
                r7 = np.random.rand(dimension)
                r8 = np.random.rand(dimension)
                A4 = 2 * a * r7 - a
                C4 = 2 * c4
                D_driver = np.abs(C4 * driver - self.m * chimps[i])
                X4 = driver - A4 * D_driver
                
                # Update chimp position (average of four positions)
                chimps[i] = (X1 + X2 + X3 + X4) / 4
                
                # Apply Levy flight in late iterations for fine-tuning
                if iteration > self.max_iterations_ * 0.6:
                    levy_step = self._levy_flight(dimension)
                    chimps[i] = chimps[i] + levy_step * (attacker - chimps[i])
                
                # Boundary control
                chimps[i] = np.clip(chimps[i], lb, ub)
                
                # Evaluate
                fitness[i] = objective_function(chimps[i])
            
            # Update leaders (attacker, barrier, chaser, driver)
            sorted_idx = np.argsort(fitness)
            attacker = chimps[sorted_idx[0]].copy()
            barrier = chimps[sorted_idx[1]].copy() if len(sorted_idx) > 1 else attacker.copy()
            chaser = chimps[sorted_idx[2]].copy() if len(sorted_idx) > 2 else attacker.copy()
            driver = chimps[sorted_idx[3]].copy() if len(sorted_idx) > 3 else attacker.copy()
            best_fitness = fitness[sorted_idx[0]]
            
            global_fitness.append(best_fitness)
            # Note: Add local_fitness and local_positions tracking as needed
            
            if self.verbose_ and iteration % 10 == 0:
                print(f"ChOA Iteration {iteration}: Best Fitness = {best_fitness:.6e}")
        
        return attacker, 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
