Atomizer/tests/test_cbush_optimization.py

"""
Test Phase 2.5 with CBUSH Element Stiffness Optimization Request

Tests the intelligent gap detection with a 1D element force optimization request.
"""

import sys
from pathlib import Path

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

from optimization_engine.utils.codebase_analyzer import CodebaseCapabilityAnalyzer
from optimization_engine.future.workflow_decomposer import WorkflowDecomposer
from optimization_engine.config.capability_matcher import CapabilityMatcher
from optimization_engine.future.targeted_research_planner import TargetedResearchPlanner


def main():
    user_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 compere 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 uptuna with TPE to iterate and optimize this"""

    print('=' * 80)
    print('PHASE 2.5 TEST: 1D Element Forces Optimization with CBUSH Stiffness')
    print('=' * 80)
    print()
    print('User Request:')
    print(user_request)
    print()
    print('=' * 80)
    print()

    # Initialize
    analyzer = CodebaseCapabilityAnalyzer()
    decomposer = WorkflowDecomposer()
    matcher = CapabilityMatcher(analyzer)
    planner = TargetedResearchPlanner()

    # Step 1: Decompose
    print('[1] Decomposing Workflow')
    print('-' * 80)
    steps = decomposer.decompose(user_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}')
        if step.params:
            print(f'     Params: {step.params}')
        print()

    # Step 2: Match to capabilities
    print()
    print('[2] Matching to Existing Capabilities')
    print('-' * 80)
    match = matcher.match(steps)
    print(f'Coverage: {match.coverage:.0%} ({len(match.known_steps)}/{len(steps)} steps)')
    print(f'Confidence: {match.overall_confidence:.0%}')
    print()

    print('KNOWN Steps (Already Implemented):')
    for i, known in enumerate(match.known_steps, 1):
        print(f'  {i}. {known.step.action.replace("_", " ").title()} ({known.step.domain})')
        if known.implementation != 'unknown':
            impl_name = Path(known.implementation).name if ('\\' in known.implementation or '/' in known.implementation) else known.implementation
            print(f'     File: {impl_name}')
    print()

    print('MISSING Steps (Need Research):')
    if match.unknown_steps:
        for i, unknown in enumerate(match.unknown_steps, 1):
            print(f'  {i}. {unknown.step.action.replace("_", " ").title()} ({unknown.step.domain})')
            print(f'     Required: {unknown.step.params}')
            if unknown.similar_capabilities:
                similar_str = ', '.join(unknown.similar_capabilities)
                print(f'     Similar to: {similar_str}')
                print(f'     Confidence: {unknown.confidence:.0%} (can adapt)')
            else:
                print(f'     Confidence: {unknown.confidence:.0%} (needs research)')
            print()
    else:
        print('  None - all capabilities are known!')
    print()

    # Step 3: Create research plan
    print()
    print('[3] Creating Targeted Research Plan')
    print('-' * 80)
    plan = planner.plan(match)
    print(f'Research steps needed: {len(plan)}')
    print()

    if plan:
        for i, step in enumerate(plan, 1):
            print(f'Step {i}: {step["description"]}')
            print(f'  Action: {step["action"]}')
            details = step.get('details', {})
            if 'capability' in details:
                print(f'  Study: {details["capability"]}')
            if 'query' in details:
                print(f'  Query: "{details["query"]}"')
            print(f'  Expected confidence: {step["expected_confidence"]:.0%}')
            print()
    else:
        print('No research needed - all capabilities exist!')
        print()

    print()
    print('=' * 80)
    print('ANALYSIS SUMMARY')
    print('=' * 80)
    print()
    print('Request Complexity:')
    print('  - Extract forces from 1D elements (Z direction)')
    print('  - Calculate average and maximum forces')
    print('  - Define custom objective metric (max vs avg comparison)')
    print('  - Modify CBUSH element stiffness properties')
    print('  - Optuna TPE optimization')
    print()
    print(f'System Analysis:')
    print(f'  Known capabilities: {len(match.known_steps)}/{len(steps)} ({match.coverage:.0%})')
    print(f'  Missing capabilities: {len(match.unknown_steps)}/{len(steps)}')
    print(f'  Overall confidence: {match.overall_confidence:.0%}')
    print()

    if match.unknown_steps:
        print('What needs research:')
        for unknown in match.unknown_steps:
            print(f'  - {unknown.step.action} ({unknown.step.domain})')
    else:
        print('All capabilities already exist in Atomizer!')

    print()


if __name__ == '__main__':
    main()