Atomizer/tests/test_step_classifier.py

"""
Test Step Classifier - Phase 2.6

Tests the intelligent classification of workflow steps into:
- Engineering features (need research/documentation)
- Inline calculations (auto-generate simple math)
- Post-processing hooks (middleware scripts)
"""

import sys
from pathlib import Path

# Set UTF-8 encoding for Windows console
if sys.platform == 'win32':
    import codecs
    if not isinstance(sys.stdout, codecs.StreamWriter):
        if hasattr(sys.stdout, 'buffer'):
            sys.stdout = codecs.getwriter('utf-8')(sys.stdout.buffer, errors='replace')
            sys.stderr = codecs.getwriter('utf-8')(sys.stderr.buffer, errors='replace')

project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from optimization_engine.future.workflow_decomposer import WorkflowDecomposer
from optimization_engine.future.step_classifier import StepClassifier


def main():
    print("=" * 80)
    print("PHASE 2.6 TEST: Intelligent Step Classification")
    print("=" * 80)
    print()

    # Test with CBUSH optimization request
    request = """I want to extract forces in direction Z of all the 1D elements and find the average of it,
then find the maximum value and compare it to the average, then assign it to a objective metric that needs to be minimized.

I want to iterate on the FEA properties of the Cbush element stiffness in Z to make the objective function minimized.

I want to use optuna with TPE to iterate and optimize this"""

    print("User Request:")
    print(request)
    print()
    print("=" * 80)
    print()

    # Initialize
    decomposer = WorkflowDecomposer()
    classifier = StepClassifier()

    # Step 1: Decompose workflow
    print("[1] Decomposing Workflow")
    print("-" * 80)
    steps = decomposer.decompose(request)
    print(f"Identified {len(steps)} workflow steps:")
    print()
    for i, step in enumerate(steps, 1):
        print(f"  {i}. {step.action.replace('_', ' ').title()}")
        print(f"     Domain: {step.domain}")
        print(f"     Params: {step.params}")
        print()

    # Step 2: Classify steps
    print()
    print("[2] Classifying Steps")
    print("-" * 80)
    classified = classifier.classify_workflow(steps, request)

    # Display classification summary
    print(classifier.get_summary(classified))
    print()

    # Step 3: Analysis
    print()
    print("[3] Intelligence Analysis")
    print("-" * 80)
    print()

    eng_count = len(classified['engineering_features'])
    inline_count = len(classified['inline_calculations'])
    hook_count = len(classified['post_processing_hooks'])

    print(f"Total Steps: {len(steps)}")
    print(f"  Engineering Features: {eng_count} (need research/documentation)")
    print(f"  Inline Calculations: {inline_count} (auto-generate Python)")
    print(f"  Post-Processing Hooks: {hook_count} (generate middleware)")
    print()

    print("What This Means:")
    if eng_count > 0:
        print(f"  - Research needed for {eng_count} FEA/CAE operations")
        print(f"  - Create documented features for reuse")
    if inline_count > 0:
        print(f"  - Auto-generate {inline_count} simple math operations")
        print(f"  - No documentation overhead needed")
    if hook_count > 0:
        print(f"  - Generate {hook_count} post-processing scripts")
        print(f"  - Execute between engineering steps")
    print()

    # Step 4: Show expected behavior
    print()
    print("[4] Expected Atomizer Behavior")
    print("-" * 80)
    print()

    print("When user makes this request, Atomizer should:")
    print()

    if eng_count > 0:
        print("  1. RESEARCH & DOCUMENT (Engineering Features):")
        for item in classified['engineering_features']:
            step = item['step']
            print(f"     - {step.action} ({step.domain})")
            print(f"       > Search pyNastran docs for element force extraction")
            print(f"       > Create feature file with documentation")
    print()

    if inline_count > 0:
        print("  2. AUTO-GENERATE (Inline Calculations):")
        for item in classified['inline_calculations']:
            step = item['step']
            print(f"     - {step.action}")
            print(f"       > Generate Python: avg = sum(forces) / len(forces)")
            print(f"       > No feature file created")
    print()

    if hook_count > 0:
        print("  3. CREATE HOOK (Post-Processing):")
        for item in classified['post_processing_hooks']:
            step = item['step']
            print(f"     - {step.action}")
            print(f"       > Generate hook script with proper I/O")
            print(f"       > Execute between solve and optimize steps")
    print()

    print("  4. EXECUTE WORKFLOW:")
    print("     - Extract 1D element forces (FEA feature)")
    print("     - Calculate avg/max/compare (inline Python)")
    print("     - Update CBUSH stiffness (FEA feature)")
    print("     - Optimize with Optuna TPE (existing feature)")
    print()

    print("=" * 80)
    print("TEST COMPLETE")
    print("=" * 80)
    print()


if __name__ == '__main__':
    main()