Atomizer/optimization_engine/extractors/test_phase3_extractors.py

"""
Test script for Phase 3 extractors (Multi-Physics)
===================================================

Tests:
- Temperature extraction (extract_temperature)
- Thermal gradient extraction (extract_temperature_gradient)
- Modal mass extraction (extract_modal_mass)

Run with:
    conda activate atomizer
    python -m optimization_engine.extractors.test_phase3_extractors
"""

import sys
from pathlib import Path

# Add project root to path
project_root = Path(__file__).parents[2]
sys.path.insert(0, str(project_root))


def test_imports():
    """Test that all Phase 3 extractors can be imported."""
    print("\n" + "=" * 60)
    print("Testing Phase 3 Extractor Imports")
    print("=" * 60)

    try:
        from optimization_engine.extractors import (
            extract_temperature,
            extract_temperature_gradient,
            extract_heat_flux,
            get_max_temperature,
            extract_modal_mass,
            extract_frequencies,
            get_first_frequency,
            get_modal_mass_ratio,
        )
        print("✓ All Phase 3 extractors imported successfully!")
        return True
    except ImportError as e:
        print(f"✗ Import failed: {e}")
        return False


def test_temperature_extractor(op2_file: str = None):
    """Test temperature extraction."""
    print("\n" + "=" * 60)
    print("Testing extract_temperature")
    print("=" * 60)

    from optimization_engine.extractors import extract_temperature

    if op2_file and Path(op2_file).exists():
        result = extract_temperature(op2_file, subcase=1)
        print(f"Result: {result}")

        if result['success']:
            print(f"  Max temperature: {result['max_temperature']}")
            print(f"  Min temperature: {result['min_temperature']}")
            print(f"  Avg temperature: {result['avg_temperature']}")
            print(f"  Node count: {result['node_count']}")
        else:
            print(f"  Note: {result['error']}")
            print("  (This is expected for non-thermal OP2 files)")
    else:
        # Test with non-existent file to verify error handling
        result = extract_temperature("nonexistent.op2")
        if not result['success'] and 'not found' in result['error']:
            print("✓ Error handling works correctly for missing files")
        else:
            print("✗ Error handling issue")

    return True


def test_temperature_gradient(op2_file: str = None):
    """Test temperature gradient extraction."""
    print("\n" + "=" * 60)
    print("Testing extract_temperature_gradient")
    print("=" * 60)

    from optimization_engine.extractors import extract_temperature_gradient

    if op2_file and Path(op2_file).exists():
        result = extract_temperature_gradient(op2_file, subcase=1)
        print(f"Result: {result}")
    else:
        result = extract_temperature_gradient("nonexistent.op2")
        if not result['success']:
            print("✓ Error handling works correctly")

    return True


def test_modal_mass_extractor(f06_file: str = None):
    """Test modal mass extraction."""
    print("\n" + "=" * 60)
    print("Testing extract_modal_mass")
    print("=" * 60)

    from optimization_engine.extractors import extract_modal_mass, get_first_frequency

    if f06_file and Path(f06_file).exists():
        # Test all modes
        result = extract_modal_mass(f06_file, mode=None)
        print(f"All modes result:")
        if result['success']:
            print(f"  Mode count: {result['mode_count']}")
            print(f"  Frequencies: {result['frequencies'][:5]}..." if len(result.get('frequencies', [])) > 5 else f"  Frequencies: {result.get('frequencies', [])}")
        else:
            print(f"  Note: {result['error']}")

        # Test specific mode
        result = extract_modal_mass(f06_file, mode=1)
        print(f"\nMode 1 result:")
        if result['success']:
            print(f"  Frequency: {result['frequency']} Hz")
            print(f"  Modal mass X: {result.get('modal_mass_x')}")
            print(f"  Modal mass Y: {result.get('modal_mass_y')}")
            print(f"  Modal mass Z: {result.get('modal_mass_z')}")
        else:
            print(f"  Note: {result['error']}")

        # Test convenience function
        freq = get_first_frequency(f06_file)
        print(f"\nFirst frequency (convenience): {freq} Hz")
    else:
        result = extract_modal_mass("nonexistent.f06")
        if not result['success'] and 'not found' in result['error']:
            print("✓ Error handling works correctly for missing files")

    return True


def find_test_files():
    """Find available test files in studies."""
    studies_dir = project_root / "studies"
    op2_files = list(studies_dir.rglob("*.op2"))
    f06_files = list(studies_dir.rglob("*.f06"))

    print(f"\nFound {len(op2_files)} OP2 files and {len(f06_files)} F06 files")

    return op2_files, f06_files


def main():
    print("=" * 60)
    print("PHASE 3 EXTRACTOR TESTS")
    print("Multi-Physics: Thermal & Dynamic")
    print("=" * 60)

    # Test imports first
    if not test_imports():
        print("\n✗ Import test failed. Cannot continue.")
        return 1

    # Find test files
    op2_files, f06_files = find_test_files()

    # Use first available files for testing
    op2_file = str(op2_files[0]) if op2_files else None
    f06_file = str(f06_files[0]) if f06_files else None

    if op2_file:
        print(f"\nUsing OP2 file: {op2_file}")
    if f06_file:
        print(f"Using F06 file: {f06_file}")

    # Run tests
    test_temperature_extractor(op2_file)
    test_temperature_gradient(op2_file)
    test_modal_mass_extractor(f06_file)

    print("\n" + "=" * 60)
    print("PHASE 3 TESTS COMPLETE")
    print("=" * 60)

    print("""
Summary:
- Temperature extraction: Ready for thermal OP2 files (SOL 153/159)
- Thermal gradient: Ready (approximation based on temperature range)
- Heat flux: Ready for thermal OP2 files
- Modal mass: Ready for modal F06 files (SOL 103)

Note: Full testing requires thermal and modal analysis result files.
The extractors will return appropriate error messages for non-thermal/modal data.
""")

    return 0


if __name__ == "__main__":
    sys.exit(main())