"""Test neural surrogate integration"""
import sys
from pathlib import Path

# Add project paths
project_root = Path(__file__).parent
sys.path.insert(0, str(project_root))
sys.path.insert(0, str(project_root / 'atomizer-field'))

from optimization_engine.neural_surrogate import create_parametric_surrogate_for_study

# Create surrogate
print("Creating parametric surrogate...")
surrogate = create_parametric_surrogate_for_study(project_root=project_root)

if surrogate:
    print('Surrogate created successfully!')
    print(f'Design vars: {surrogate.design_var_names}')
    print(f'Number of nodes: {surrogate.num_nodes}')

    # Test prediction with example params
    test_params = {name: 2.0 for name in surrogate.design_var_names}
    print(f'\nTest params: {test_params}')

    results = surrogate.predict(test_params)
    print(f'\nTest prediction:')
    print(f'  Mass: {results["mass"]:.2f}')
    print(f'  Frequency: {results["frequency"]:.2f}')
    print(f'  Max Displacement: {results["max_displacement"]:.6f}')
    print(f'  Max Stress: {results["max_stress"]:.2f}')
    print(f'  Inference time: {results["inference_time_ms"]:.2f} ms')

    print('\nSurrogate is ready for use in optimization!')
else:
    print('Failed to create surrogate')