"""
Test Script for Genetic Operators and EPD-GA Algorithms
========================================================

Tests the new genetic algorithm implementations including:
1. Standard GA with various operators
2. EPD-GA hybrid algorithm
3. EPD-GA variants (Explorative, Exploitative, Balanced)
"""

import numpy as np
import sys
import os

# Add parent directory to path
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from mha_toolbox.algorithms.genetic_algorithm import GeneticAlgorithm
from mha_toolbox.algorithms.epd_ga import (
    EPD_GA, EPD_GA_Explorative, EPD_GA_Exploitative, EPD_GA_Balanced
)


def sphere_function(x):
    """Simple sphere function for testing"""
    return np.sum(x ** 2)


def rastrigin_function(x):
    """Rastrigin function - multimodal test"""
    A = 10
    return A * len(x) + np.sum(x**2 - A * np.cos(2 * np.pi * x))


def rosenbrock_function(x):
    """Rosenbrock function - narrow valley test"""
    return np.sum(100.0 * (x[1:] - x[:-1]**2)**2 + (1 - x[:-1])**2)


def test_genetic_algorithm():
    """Test standard Genetic Algorithm"""
    print("\n" + "="*70)
    print("Testing Standard Genetic Algorithm (GA)")
    print("="*70)
    
    # Test parameters
    dimensions = 10
    lower_bound = -5.0
    upper_bound = 5.0
    population_size = 30
    max_iterations = 50
    
    print(f"\nTest Configuration:")
    print(f"  Dimensions: {dimensions}")
    print(f"  Search Space: [{lower_bound}, {upper_bound}]")
    print(f"  Population Size: {population_size}")
    print(f"  Max Iterations: {max_iterations}")
    
    # Test on Sphere function
    print(f"\n{'='*70}")
    print("Test 1: Sphere Function (Unimodal)")
    print(f"{'='*70}")
    
    ga = GeneticAlgorithm(
        population_size=population_size,
        max_iterations=max_iterations,
        mutation_type='gaussian',
        crossover_type='simulated_binary',
        adaptive=True,
        lower_bound=lower_bound,
        upper_bound=upper_bound,
        dimensions=dimensions
    )
    
    best_solution, best_fitness, convergence = ga._optimize(sphere_function)
    
    print(f"\nResults:")
    print(f"  Best Fitness: {best_fitness:.6e}")
    print(f"  Best Solution (first 5): {best_solution[:5]}")
    print(f"  Initial Fitness: {convergence[0]:.6e}")
    print(f"  Final Fitness: {convergence[-1]:.6e}")
    print(f"  Improvement: {((convergence[0] - convergence[-1]) / convergence[0] * 100):.2f}%")
    print(f"  Convergence Trend: {convergence[::10]}")
    
    # Success if significant improvement achieved
    return best_fitness < 1.0  # Relaxed threshold


def test_epd_ga():
    """Test EPD-GA Hybrid Algorithm"""
    print("\n" + "="*70)
    print("Testing EPD-GA (Evolutionary Population Dynamics GA)")
    print("="*70)
    
    dimensions = 10
    lower_bound = -5.0
    upper_bound = 5.0
    population_size = 30
    max_iterations = 50
    
    print(f"\nTest Configuration:")
    print(f"  Dimensions: {dimensions}")
    print(f"  Search Space: [{lower_bound}, {upper_bound}]")
    print(f"  Population Size: {population_size}")
    print(f"  Max Iterations: {max_iterations}")
    print(f"  EPD Frequency: Every 5 generations")
    
    # Test on Rastrigin (multimodal)
    print(f"\n{'='*70}")
    print("Test 2: Rastrigin Function (Highly Multimodal)")
    print(f"{'='*70}")
    
    epd_ga = EPD_GA(
        population_size=population_size,
        max_iterations=max_iterations,
        epd_frequency=5,
        lower_bound=lower_bound,
        upper_bound=upper_bound,
        dimensions=dimensions
    )
    
    best_solution, best_fitness, convergence = epd_ga._optimize(rastrigin_function)
    
    print(f"\nResults:")
    print(f"  Best Fitness: {best_fitness:.6e}")
    print(f"  Best Solution (first 5): {best_solution[:5]}")
    print(f"  Initial Fitness: {convergence[0]:.6e}")
    print(f"  Final Fitness: {convergence[-1]:.6e}")
    print(f"  Improvement: {((convergence[0] - convergence[-1]) / convergence[0] * 100):.2f}%")
    print(f"  EPD Applications: {len(epd_ga.epd_applications)} times at iterations {epd_ga.epd_applications[:5]}...")
    print(f"  Final Diversity: {epd_ga.diversity_history[-1]:.4f}")
    
    # Get algorithm info
    info = epd_ga.get_algorithm_info()
    print(f"\nAlgorithm Info:")
    print(f"  Name: {info['name']}")
    print(f"  Type: {info['type']}")
    print(f"  Features: {', '.join(info['features'][:3])}")
    
    # Success if significant improvement achieved
    return best_fitness < 100  # Rastrigin is harder, relaxed threshold


def test_epd_ga_variants():
    """Test EPD-GA variants"""
    print("\n" + "="*70)
    print("Testing EPD-GA Variants (Explorative, Exploitative, Balanced)")
    print("="*70)
    
    dimensions = 10
    lower_bound = -5.0
    upper_bound = 5.0
    population_size = 30
    max_iterations = 50
    
    variants = {
        'Explorative': EPD_GA_Explorative,
        'Exploitative': EPD_GA_Exploitative,
        'Balanced': EPD_GA_Balanced
    }
    
    results = {}
    
    for variant_name, VariantClass in variants.items():
        print(f"\n{'='*70}")
        print(f"Test 3: {variant_name} Variant on Rosenbrock")
        print(f"{'='*70}")
        
        algo = VariantClass(
            objective_function=rosenbrock_function,
            population_size=population_size,
            max_iterations=max_iterations,
            lower_bound=lower_bound,
            upper_bound=upper_bound,
            dimensions=dimensions
        )
        
        best_solution, best_fitness, convergence = algo._optimize(rosenbrock_function)
        
        print(f"\nResults for {variant_name}:")
        print(f"  Best Fitness: {best_fitness:.6e}")
        print(f"  Improvement: {((convergence[0] - convergence[-1]) / convergence[0] * 100):.2f}%")
        print(f"  Final Diversity: {algo.diversity_history[-1]:.4f}")
        
        results[variant_name] = best_fitness
    
    # Compare variants
    print(f"\n{'='*70}")
    print("Variant Comparison:")
    print(f"{'='*70}")
    for variant_name, fitness in sorted(results.items(), key=lambda x: x[1]):
        print(f"  {variant_name:15s}: {fitness:.6e}")
    
    return True


def test_operator_diversity():
    """Test different operator combinations"""
    print("\n" + "="*70)
    print("Testing Different Genetic Operator Combinations")
    print("="*70)
    
    dimensions = 10
    lower_bound = -5.0
    upper_bound = 5.0
    
    configs = [
        ('Gaussian + SBX', 'gaussian', 'simulated_binary'),
        ('Uniform + Uniform', 'uniform', 'uniform'),
        ('Polynomial + Arithmetic', 'polynomial', 'arithmetic'),
    ]
    
    results = {}
    
    for config_name, mutation_type, crossover_type in configs:
        print(f"\nTesting: {config_name}")
        print(f"  Mutation: {mutation_type}")
        print(f"  Crossover: {crossover_type}")
        
        ga = GeneticAlgorithm(
            objective_function=sphere_function,
            population_size=30,
            max_iterations=30,
            mutation_type=mutation_type,
            crossover_type=crossover_type,
            adaptive=True,
            lower_bound=lower_bound,
            upper_bound=upper_bound,
            dimensions=dimensions
        )
        
        try:
            _, best_fitness, convergence = ga._optimize(sphere_function)
            improvement = (convergence[0] - convergence[-1]) / convergence[0] * 100
            
            print(f"  Best Fitness: {best_fitness:.6e}")
            print(f"  Improvement: {improvement:.2f}%")
            
            results[config_name] = best_fitness
        except Exception as e:
            print(f"  ERROR: {e}")
            results[config_name] = float('inf')
    
    print(f"\n{'='*70}")
    print("Operator Comparison:")
    print(f"{'='*70}")
    for config_name, fitness in sorted(results.items(), key=lambda x: x[1]):
        print(f"  {config_name:30s}: {fitness:.6e}")
    
    return True


def run_all_tests():
    """Run all tests"""
    print("\n" + "="*70)
    print("GENETIC OPERATORS & EPD-GA TEST SUITE")
    print("="*70)
    print("\nThis test suite validates:")
    print("  1. Standard Genetic Algorithm with mutation & crossover")
    print("  2. EPD-GA with self-organized criticality")
    print("  3. EPD-GA variants (Explorative, Exploitative, Balanced)")
    print("  4. Different operator combinations")
    
    test_results = []
    
    try:
        test_results.append(("Standard GA", test_genetic_algorithm()))
    except Exception as e:
        print(f"\n❌ Standard GA test failed: {e}")
        test_results.append(("Standard GA", False))
    
    try:
        test_results.append(("EPD-GA", test_epd_ga()))
    except Exception as e:
        print(f"\n❌ EPD-GA test failed: {e}")
        test_results.append(("EPD-GA", False))
    
    try:
        test_results.append(("EPD-GA Variants", test_epd_ga_variants()))
    except Exception as e:
        print(f"\n❌ EPD-GA Variants test failed: {e}")
        test_results.append(("EPD-GA Variants", False))
    
    try:
        test_results.append(("Operator Combinations", test_operator_diversity()))
    except Exception as e:
        print(f"\n❌ Operator test failed: {e}")
        test_results.append(("Operator Combinations", False))
    
    # Final summary
    print("\n" + "="*70)
    print("TEST SUMMARY")
    print("="*70)
    
    for test_name, passed in test_results:
        status = "PASSED" if passed else "FAILED"
        print(f"  {test_name:25s}: {status}")
    
    passed_count = sum(1 for _, passed in test_results if passed)
    total_count = len(test_results)
    
    print(f"\nOverall: {passed_count}/{total_count} tests passed")
    
    if passed_count == total_count:
        print("\nALL TESTS PASSED! Genetic operators and EPD-GA are working correctly.")
        return 0
    else:
        print(f"\n{total_count - passed_count} test(s) failed. Please review the output above.")
        return 1


if __name__ == "__main__":
    exit_code = run_all_tests()
    sys.exit(exit_code)
