Atomizer/tests/test_lifecycle_hook_integration.py

"""
Test Integration of Phase 2.9 Hooks with Phase 1 Lifecycle Hook System

This demonstrates how generated hooks integrate with the existing HookManager.
"""

import sys
from pathlib import Path

# Add project root to path
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from optimization_engine.hook_generator import HookGenerator
from optimization_engine.plugins.hook_manager import HookManager


def test_lifecycle_hook_generation():
    """Test generating lifecycle-compatible hooks."""
    print("=" * 80)
    print("Testing Lifecycle Hook Integration (Phase 2.9 + Phase 1)")
    print("=" * 80)
    print()

    generator = HookGenerator()

    # Test hook spec from LLM (Phase 2.7 output)
    hook_spec = {
        "action": "weighted_objective",
        "description": "Combine normalized stress (70%) and displacement (30%)",
        "params": {
            "inputs": ["norm_stress", "norm_disp"],
            "weights": [0.7, 0.3],
            "objective": "minimize"
        }
    }

    print("1. Generating lifecycle-compatible hook...")
    hook_module_content = generator.generate_lifecycle_hook(hook_spec)

    print("\nGenerated Hook Module:")
    print("=" * 80)
    print(hook_module_content)
    print("=" * 80)
    print()

    # Save the hook to post_calculation directory
    output_dir = project_root / "optimization_engine" / "plugins" / "post_calculation"
    output_file = output_dir / "weighted_objective_test.py"

    print(f"2. Saving hook to: {output_file}")
    with open(output_file, 'w') as f:
        f.write(hook_module_content)
    print(f"   [OK] Hook saved")
    print()

    # Test loading with HookManager
    print("3. Testing hook registration with HookManager...")
    hook_manager = HookManager()

    # Load the plugin
    hook_manager.load_plugins_from_directory(project_root / "optimization_engine" / "plugins")

    # Show summary
    summary = hook_manager.get_summary()
    print(f"\n   Hook Manager Summary:")
    print(f"   - Total hooks: {summary['total_hooks']}")
    print(f"   - Enabled: {summary['enabled_hooks']}")
    print(f"\n   Hooks by point:")
    for point, info in summary['by_hook_point'].items():
        if info['total'] > 0:
            print(f"     {point}: {info['total']} hook(s) - {info['names']}")
    print()

    # Test executing the hook
    print("4. Testing hook execution...")
    test_context = {
        'trial_number': 42,
        'calculations': {
            'norm_stress': 0.75,
            'norm_disp': 0.64
        },
        'results': {}
    }

    results = hook_manager.execute_hooks('post_calculation', test_context)

    print(f"\n   Execution results: {results}")
    if results and results[0]:
        weighted_obj = results[0].get('weighted_objective')
        expected = 0.7 * 0.75 + 0.3 * 0.64
        print(f"   Weighted objective: {weighted_obj:.6f}")
        print(f"   Expected: {expected:.6f}")
        print(f"   Match: {'YES' if abs(weighted_obj - expected) < 0.0001 else 'NO'}")
    print()

    print("=" * 80)
    print("Lifecycle Hook Integration Test Complete!")
    print("=" * 80)


def test_multiple_hook_types():
    """Test generating different hook types for lifecycle system."""
    print("\n" + "=" * 80)
    print("Testing Multiple Hook Types")
    print("=" * 80)
    print()

    generator = HookGenerator()
    output_dir = Path(project_root) / "optimization_engine" / "plugins" / "post_calculation"

    # Test different hook types
    hook_specs = [
        {
            "action": "custom_formula",
            "description": "Calculate safety factor",
            "params": {
                "inputs": ["max_stress", "yield_strength"],
                "formula": "yield_strength / max_stress",
                "output_name": "safety_factor"
            }
        },
        {
            "action": "comparison",
            "description": "Compare min force to average",
            "params": {
                "inputs": ["min_force", "avg_force"],
                "operation": "ratio",
                "output_name": "min_to_avg_ratio"
            }
        }
    ]

    for i, spec in enumerate(hook_specs, 1):
        action = spec['action']
        print(f"{i}. Generating {action} hook...")

        hook_content = generator.generate_lifecycle_hook(spec)

        # Infer filename
        if 'formula' in action:
            filename = f"safety_factor_hook.py"
        elif 'comparison' in action:
            filename = f"min_to_avg_ratio_hook.py"
        else:
            filename = f"{action}_hook.py"

        output_file = output_dir / filename
        with open(output_file, 'w') as f:
            f.write(hook_content)

        print(f"   [OK] Saved to: {filename}")

    print("\n" + "=" * 80)
    print("All hook types generated successfully!")
    print("=" * 80)


if __name__ == '__main__':
    test_lifecycle_hook_generation()
    test_multiple_hook_types()

    print("\n" + "=" * 80)
    print("Summary")
    print("=" * 80)
    print()
    print("Generated lifecycle hooks can be used interchangeably at ANY hook point:")
    print("  - pre_mesh: Before meshing")
    print("  - post_mesh: After meshing")
    print("  - pre_solve: Before FEA solve")
    print("  - post_solve: After FEA solve")
    print("  - post_extraction: After result extraction")
    print("  - post_calculation: After inline calculations (NEW!)")
    print()
    print("Simply change the hook_point parameter when generating!")
    print()