"""
Standard Objective Functions for MHA Flow v3.0
Provides benchmark functions for algorithm testing
"""
import numpy as np


def sphere(x):
    """
    Sphere function: f(x) = sum(x^2)
    Global minimum: f(0,...,0) = 0
    Search domain: [-100, 100]
    """
    return np.sum(x**2)


def rastrigin(x):
    """
    Rastrigin function - highly multimodal
    Global minimum: f(0,...,0) = 0
    Search domain: [-5.12, 5.12]
    """
    A = 10
    n = len(x)
    return A * n + np.sum(x**2 - A * np.cos(2 * np.pi * x))


def rosenbrock(x):
    """
    Rosenbrock function - valley-shaped
    Global minimum: f(1,...,1) = 0
    Search domain: [-2.048, 2.048]
    """
    return np.sum(100 * (x[1:] - x[:-1]**2)**2 + (1 - x[:-1])**2)


def ackley(x):
    """
    Ackley function - multimodal with deep valley
    Global minimum: f(0,...,0) = 0
    Search domain: [-32, 32]
    """
    a, b, c = 20, 0.2, 2 * np.pi
    n = len(x)
    sum1 = np.sum(x**2)
    sum2 = np.sum(np.cos(c * x))
    return -a * np.exp(-b * np.sqrt(sum1 / n)) - np.exp(sum2 / n) + a + np.exp(1)


def griewank(x):
    """
    Griewank function - multimodal
    Global minimum: f(0,...,0) = 0
    Search domain: [-600, 600]
    """
    sum_part = np.sum(x**2) / 4000
    prod_part = np.prod(np.cos(x / np.sqrt(np.arange(1, len(x) + 1))))
    return sum_part - prod_part + 1


def schwefel(x):
    """
    Schwefel function - multimodal with many local minima
    Global minimum: f(420.9687,...,420.9687) = 0
    Search domain: [-500, 500]
    """
    n = len(x)
    return 418.9829 * n - np.sum(x * np.sin(np.sqrt(np.abs(x))))


def levy(x):
    """
    Levy function - multimodal
    Global minimum: f(1,...,1) = 0
    Search domain: [-10, 10]
    """
    w = 1 + (x - 1) / 4
    term1 = np.sin(np.pi * w[0])**2
    term3 = (w[-1] - 1)**2 * (1 + np.sin(2 * np.pi * w[-1])**2)
    
    wi = w[:-1]
    sum_term = np.sum((wi - 1)**2 * (1 + 10 * np.sin(np.pi * wi + 1)**2))
    
    return term1 + sum_term + term3


def zakharov(x):
    """
    Zakharov function - unimodal
    Global minimum: f(0,...,0) = 0
    Search domain: [-5, 10]
    """
    n = len(x)
    sum1 = np.sum(x**2)
    sum2 = np.sum(0.5 * np.arange(1, n + 1) * x)
    return sum1 + sum2**2 + sum2**4


def michalewicz(x):
    """
    Michalewicz function - multimodal, many local minima
    Global minimum: varies with dimensions
    Search domain: [0, pi]
    """
    m = 10
    i = np.arange(1, len(x) + 1)
    return -np.sum(np.sin(x) * np.sin(i * x**2 / np.pi)**(2 * m))


def dixon_price(x):
    """
    Dixon-Price function - unimodal
    Global minimum: f(x*) = 0 where x_i = 2^(-(2^i-2)/(2^i))
    Search domain: [-10, 10]
    """
    n = len(x)
    i = np.arange(2, n + 1)
    term1 = (x[0] - 1)**2
    term2 = np.sum(i * (2 * x[1:]**2 - x[:-1])**2)
    return term1 + term2


def styblinski_tang(x):
    """
    Styblinski-Tang function - multimodal
    Global minimum: f(-2.903534,...,-2.903534) = -39.16599*n
    Search domain: [-5, 5]
    """
    return np.sum(x**4 - 16 * x**2 + 5 * x) / 2


def sum_squares(x):
    """
    Sum Squares function - unimodal
    Global minimum: f(0,...,0) = 0
    Search domain: [-10, 10]
    """
    i = np.arange(1, len(x) + 1)
    return np.sum(i * x**2)


def powell(x):
    """
    Powell function - unimodal
    Global minimum: f(0,...,0) = 0
    Search domain: [-4, 5]
    """
    n = len(x)
    result = 0
    for i in range(0, n - 3, 4):
        result += (x[i] + 10 * x[i + 1])**2
        result += 5 * (x[i + 2] - x[i + 3])**2
        result += (x[i + 1] - 2 * x[i + 2])**4
        result += 10 * (x[i] - x[i + 3])**4
    return result


def schaffer_n2(x):
    """
    Schaffer N.2 function - multimodal
    Global minimum: f(0,0) = 0
    Search domain: [-100, 100]
    """
    if len(x) < 2:
        return 0
    return 0.5 + (np.sin(x[0]**2 - x[1]**2)**2 - 0.5) / (1 + 0.001 * (x[0]**2 + x[1]**2))**2


def trid(x):
    """
    Trid function
    Global minimum: varies with dimension
    Search domain: [-n^2, n^2]
    """
    return np.sum((x - 1)**2) - np.sum(x[1:] * x[:-1])


# Dictionary of all available functions
OBJECTIVE_FUNCTIONS = {
    'sphere': sphere,
    'rastrigin': rastrigin,
    'rosenbrock': rosenbrock,
    'ackley': ackley,
    'griewank': griewank,
    'schwefel': schwefel,
    'levy': levy,
    'zakharov': zakharov,
    'michalewicz': michalewicz,
    'dixon_price': dixon_price,
    'styblinski_tang': styblinski_tang,
    'sum_squares': sum_squares,
    'powell': powell,
    'schaffer_n2': schaffer_n2,
    'trid': trid
}

# Recommended bounds for each function
FUNCTION_BOUNDS = {
    'sphere': (-100, 100),
    'rastrigin': (-5.12, 5.12),
    'rosenbrock': (-2.048, 2.048),
    'ackley': (-32, 32),
    'griewank': (-600, 600),
    'schwefel': (-500, 500),
    'levy': (-10, 10),
    'zakharov': (-5, 10),
    'michalewicz': (0, np.pi),
    'dixon_price': (-10, 10),
    'styblinski_tang': (-5, 5),
    'sum_squares': (-10, 10),
    'powell': (-4, 5),
    'schaffer_n2': (-100, 100),
    'trid': (-100, 100)
}

# Function properties
FUNCTION_PROPERTIES = {
    'sphere': {'type': 'unimodal', 'separable': True, 'differentiable': True},
    'rastrigin': {'type': 'multimodal', 'separable': True, 'differentiable': True},
    'rosenbrock': {'type': 'unimodal', 'separable': False, 'differentiable': True},
    'ackley': {'type': 'multimodal', 'separable': False, 'differentiable': True},
    'griewank': {'type': 'multimodal', 'separable': False, 'differentiable': True},
    'schwefel': {'type': 'multimodal', 'separable': True, 'differentiable': True},
    'levy': {'type': 'multimodal', 'separable': False, 'differentiable': True},
    'zakharov': {'type': 'unimodal', 'separable': False, 'differentiable': True},
    'michalewicz': {'type': 'multimodal', 'separable': True, 'differentiable': True},
    'dixon_price': {'type': 'unimodal', 'separable': False, 'differentiable': True},
    'styblinski_tang': {'type': 'multimodal', 'separable': True, 'differentiable': True},
    'sum_squares': {'type': 'unimodal', 'separable': True, 'differentiable': True},
    'powell': {'type': 'unimodal', 'separable': False, 'differentiable': True},
    'schaffer_n2': {'type': 'multimodal', 'separable': False, 'differentiable': True},
    'trid': {'type': 'unimodal', 'separable': False, 'differentiable': True}
}


if __name__ == '__main__':
    # Test all functions
    print("Testing objective functions...")
    print("=" * 60)
    
    for name, func in OBJECTIVE_FUNCTIONS.items():
        x = np.zeros(10)  # Test at origin
        try:
            result = func(x)
            print(f"✓ {name:20s}: f(0) = {result:12.6f}")
        except Exception as e:
            print(f"✗ {name:20s}: ERROR - {e}")
    
    print("=" * 60)
    print(f"Total functions: {len(OBJECTIVE_FUNCTIONS)}")
