Atomizer/optimization_engine/capability_matcher.py

"""
Capability Matcher

Matches required workflow steps to existing codebase capabilities and identifies
actual knowledge gaps.

Author: Atomizer Development Team
Version: 0.1.0 (Phase 2.5)
Last Updated: 2025-01-16
"""

from typing import Dict, List, Any, Optional
from dataclasses import dataclass

from optimization_engine.workflow_decomposer import WorkflowStep
from optimization_engine.codebase_analyzer import CodebaseCapabilityAnalyzer


@dataclass
class StepMatch:
    """Represents the match status of a workflow step."""
    step: WorkflowStep
    is_known: bool
    implementation: Optional[str] = None
    similar_capabilities: List[str] = None
    confidence: float = 0.0


@dataclass
class CapabilityMatch:
    """Complete matching result for a workflow."""
    known_steps: List[StepMatch]
    unknown_steps: List[StepMatch]
    overall_confidence: float
    coverage: float  # Percentage of steps that are known


class CapabilityMatcher:
    """Matches required workflow steps to existing capabilities."""

    def __init__(self, analyzer: Optional[CodebaseCapabilityAnalyzer] = None):
        self.analyzer = analyzer or CodebaseCapabilityAnalyzer()
        self.capabilities = self.analyzer.analyze_codebase()

        # Mapping from workflow actions to capability checks
        self.action_to_capability = {
            'identify_parameters': ('geometry', 'expression_filtering'),
            'update_parameters': ('optimization', 'parameter_updating'),
            'read_expression': ('geometry', 'parameter_extraction'),  # Reading expressions from .prt
            'run_analysis': ('simulation', 'nx_solver'),
            'optimize': ('optimization', 'optuna_integration'),
            'create_material': ('materials', 'xml_generation'),
            'apply_loads': ('loads_bc', 'load_application'),
            'generate_mesh': ('mesh', 'mesh_generation')
        }

    def match(self, workflow_steps: List[WorkflowStep]) -> CapabilityMatch:
        """
        Match workflow steps to existing capabilities.

        Returns:
        {
            'known_steps': [
                {'step': WorkflowStep(...), 'implementation': 'parameter_updater.py'},
                ...
            ],
            'unknown_steps': [
                {'step': WorkflowStep(...), 'similar_to': 'extract_stress', 'gap': 'strain_from_op2'}
            ],
            'overall_confidence': 0.80,  # 4/5 steps known
            'coverage': 0.80
        }
        """
        known_steps = []
        unknown_steps = []

        for step in workflow_steps:
            match = self._match_step(step)

            if match.is_known:
                known_steps.append(match)
            else:
                unknown_steps.append(match)

        # Calculate coverage
        total_steps = len(workflow_steps)
        coverage = len(known_steps) / total_steps if total_steps > 0 else 0.0

        # Calculate overall confidence
        # Known steps contribute 100%, unknown steps contribute based on similarity
        total_confidence = sum(m.confidence for m in known_steps)
        total_confidence += sum(m.confidence for m in unknown_steps)
        overall_confidence = total_confidence / total_steps if total_steps > 0 else 0.0

        return CapabilityMatch(
            known_steps=known_steps,
            unknown_steps=unknown_steps,
            overall_confidence=overall_confidence,
            coverage=coverage
        )

    def _match_step(self, step: WorkflowStep) -> StepMatch:
        """Match a single workflow step to capabilities."""

        # Special handling for extract_result action
        if step.action == 'extract_result':
            return self._match_extraction_step(step)

        # Special handling for run_analysis action
        if step.action == 'run_analysis':
            return self._match_simulation_step(step)

        # General capability matching
        if step.action in self.action_to_capability:
            category, capability_name = self.action_to_capability[step.action]

            if category in self.capabilities:
                if capability_name in self.capabilities[category]:
                    if self.capabilities[category][capability_name]:
                        # Found!
                        details = self.analyzer.get_capability_details(category, capability_name)
                        impl = details['implementation_files'][0] if details and details.get('implementation_files') else 'unknown'

                        return StepMatch(
                            step=step,
                            is_known=True,
                            implementation=impl,
                            confidence=1.0
                        )

        # Not found - check for similar capabilities
        similar = self._find_similar_capabilities(step)

        return StepMatch(
            step=step,
            is_known=False,
            similar_capabilities=similar,
            confidence=0.3 if similar else 0.0  # Some confidence if similar capabilities exist
        )

    def _match_extraction_step(self, step: WorkflowStep) -> StepMatch:
        """Special matching logic for result extraction steps."""
        result_type = step.params.get('result_type', '')

        if not result_type:
            return StepMatch(step=step, is_known=False, confidence=0.0)

        # Check if this extraction capability exists
        if 'result_extraction' in self.capabilities:
            if result_type in self.capabilities['result_extraction']:
                if self.capabilities['result_extraction'][result_type]:
                    # Found!
                    details = self.analyzer.get_capability_details('result_extraction', result_type)
                    impl = details['implementation_files'][0] if details and details.get('implementation_files') else 'unknown'

                    return StepMatch(
                        step=step,
                        is_known=True,
                        implementation=impl,
                        confidence=1.0
                    )

        # Not found - find similar extraction capabilities
        similar = self.analyzer.find_similar_capabilities(result_type, 'result_extraction')

        # For result extraction, if similar capabilities exist, confidence is higher
        # because the pattern is likely the same (just different OP2 attribute)
        confidence = 0.6 if similar else 0.0

        return StepMatch(
            step=step,
            is_known=False,
            similar_capabilities=similar,
            confidence=confidence
        )

    def _match_simulation_step(self, step: WorkflowStep) -> StepMatch:
        """Special matching logic for simulation steps."""
        solver = step.params.get('solver', '')

        # Check if NX solver exists
        if 'simulation' in self.capabilities:
            if self.capabilities['simulation'].get('nx_solver'):
                # NX solver exists - check specific solver type
                solver_lower = solver.lower()

                if solver_lower in self.capabilities['simulation']:
                    if self.capabilities['simulation'][solver_lower]:
                        # Specific solver supported
                        details = self.analyzer.get_capability_details('simulation', 'nx_solver')
                        impl = details['implementation_files'][0] if details and details.get('implementation_files') else 'unknown'

                        return StepMatch(
                            step=step,
                            is_known=True,
                            implementation=impl,
                            confidence=1.0
                        )

                # NX solver exists but specific solver type not verified
                # Still high confidence because solver is generic
                details = self.analyzer.get_capability_details('simulation', 'nx_solver')
                impl = details['implementation_files'][0] if details and details.get('implementation_files') else 'unknown'

                return StepMatch(
                    step=step,
                    is_known=True,  # Consider it known since NX solver is generic
                    implementation=impl,
                    confidence=0.9  # Slight uncertainty about specific solver
                )

        return StepMatch(step=step, is_known=False, confidence=0.0)

    def _find_similar_capabilities(self, step: WorkflowStep) -> List[str]:
        """Find capabilities similar to what's needed for this step."""
        similar = []

        # Check in the step's domain
        if step.domain in self.capabilities:
            # Look for capabilities with overlapping words
            step_words = set(step.action.lower().split('_'))

            for cap_name, exists in self.capabilities[step.domain].items():
                if not exists:
                    continue

                cap_words = set(cap_name.lower().split('_'))

                # If there's overlap, it's similar
                if step_words & cap_words:
                    similar.append(cap_name)

        return similar

    def get_match_summary(self, match: CapabilityMatch) -> str:
        """Get human-readable summary of capability matching."""
        lines = [
            "Workflow Component Analysis",
            "=" * 80,
            ""
        ]

        if match.known_steps:
            lines.append(f"Known Capabilities ({len(match.known_steps)} of {len(match.known_steps) + len(match.unknown_steps)}):")
            lines.append("-" * 80)

            for i, step_match in enumerate(match.known_steps, 1):
                step = step_match.step
                lines.append(f"{i}. {step.action.replace('_', ' ').title()}")
                lines.append(f"   Domain: {step.domain}")
                if step_match.implementation:
                    lines.append(f"   Implementation: {step_match.implementation}")
                lines.append(f"   Status: KNOWN")
                lines.append("")

        if match.unknown_steps:
            lines.append(f"Missing Capabilities ({len(match.unknown_steps)}):")
            lines.append("-" * 80)

            for i, step_match in enumerate(match.unknown_steps, 1):
                step = step_match.step
                lines.append(f"{i}. {step.action.replace('_', ' ').title()}")
                lines.append(f"   Domain: {step.domain}")
                if step.params:
                    lines.append(f"   Required: {step.params}")
                lines.append(f"   Status: MISSING")

                if step_match.similar_capabilities:
                    lines.append(f"   Similar capabilities found: {', '.join(step_match.similar_capabilities)}")
                    lines.append(f"   Confidence: {step_match.confidence:.0%} (can adapt from similar)")
                else:
                    lines.append(f"   Confidence: {step_match.confidence:.0%} (needs research)")
                lines.append("")

        lines.append("=" * 80)
        lines.append(f"Overall Coverage: {match.coverage:.0%}")
        lines.append(f"Overall Confidence: {match.overall_confidence:.0%}")
        lines.append("")

        return "\n".join(lines)


def main():
    """Test the capability matcher."""
    from optimization_engine.workflow_decomposer import WorkflowDecomposer

    print("Capability Matcher Test")
    print("=" * 80)
    print()

    # Initialize components
    analyzer = CodebaseCapabilityAnalyzer()
    decomposer = WorkflowDecomposer()
    matcher = CapabilityMatcher(analyzer)

    # Test with strain optimization request
    test_request = "I want to evaluate strain on a part with sol101 and optimize this (minimize) using iterations and optuna to lower it varying all my geometry parameters that contains v_ in its expression"

    print("Request:")
    print(test_request)
    print()

    # Decompose workflow
    print("Step 1: Decomposing workflow...")
    steps = decomposer.decompose(test_request)
    print(f"  Identified {len(steps)} workflow steps")
    print()

    # Match to capabilities
    print("Step 2: Matching to existing capabilities...")
    match = matcher.match(steps)
    print()

    # Display results
    print(matcher.get_match_summary(match))

    # Show what needs to be researched
    if match.unknown_steps:
        print("\nResearch Needed:")
        print("-" * 80)
        for step_match in match.unknown_steps:
            step = step_match.step
            print(f"  Topic: How to {step.action.replace('_', ' ')}")
            print(f"  Domain: {step.domain}")

            if step_match.similar_capabilities:
                print(f"  Strategy: Adapt from {step_match.similar_capabilities[0]}")
                print(f"            (follow same pattern, different OP2 attribute)")
            else:
                print(f"  Strategy: Research from scratch")
                print(f"            (search docs, ask user for examples)")
            print()


if __name__ == '__main__':
    main()