Atomizer/optimization_engine/targeted_research_planner.py

"""
Targeted Research Planner

Creates focused research plans that target ONLY the actual knowledge gaps,
leveraging similar existing capabilities when available.

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

from typing import List, Dict, Any
from pathlib import Path

from optimization_engine.capability_matcher import CapabilityMatch, StepMatch


class TargetedResearchPlanner:
    """Creates research plan focused on actual gaps."""

    def __init__(self):
        pass

    def plan(self, capability_match: CapabilityMatch) -> List[Dict[str, Any]]:
        """
        Create targeted research plan for missing capabilities.

        For gap='strain_from_op2', similar_to='stress_from_op2':

        Research Plan:
        1. Read existing op2_extractor_example.py to understand pattern
        2. Search pyNastran docs for strain extraction API
        3. If not found, ask user for strain extraction example
        4. Generate extract_strain() function following same pattern as extract_stress()
        """
        if not capability_match.unknown_steps:
            return []

        research_steps = []

        for unknown_step in capability_match.unknown_steps:
            steps_for_this_gap = self._plan_for_gap(unknown_step)
            research_steps.extend(steps_for_this_gap)

        return research_steps

    def _plan_for_gap(self, step_match: StepMatch) -> List[Dict[str, Any]]:
        """Create research plan for a single gap."""
        step = step_match.step
        similar = step_match.similar_capabilities

        plan_steps = []

        # If we have similar capabilities, start by studying them
        if similar:
            plan_steps.append({
                'action': 'read_existing_code',
                'description': f'Study existing {similar[0]} implementation to understand pattern',
                'details': {
                    'capability': similar[0],
                    'category': step.domain,
                    'purpose': f'Learn pattern for {step.action}'
                },
                'expected_confidence': 0.7,
                'priority': 1
            })

        # Search knowledge base for previous similar work
        plan_steps.append({
            'action': 'search_knowledge_base',
            'description': f'Search for previous {step.domain} work',
            'details': {
                'query': f"{step.domain} {step.action}",
                'required_params': step.params
            },
            'expected_confidence': 0.8 if similar else 0.5,
            'priority': 2
        })

        # For result extraction, search pyNastran docs
        if step.domain == 'result_extraction':
            result_type = step.params.get('result_type', '')
            plan_steps.append({
                'action': 'search_pynastran_docs',
                'description': f'Search pyNastran documentation for {result_type} extraction',
                'details': {
                    'query': f'pyNastran OP2 {result_type} extraction',
                    'library': 'pyNastran',
                    'expected_api': f'op2.{result_type}s or similar'
                },
                'expected_confidence': 0.85,
                'priority': 3
            })

        # For simulation, search NX docs
        elif step.domain == 'simulation':
            solver = step.params.get('solver', '')
            plan_steps.append({
                'action': 'query_nx_docs',
                'description': f'Search NX documentation for {solver}',
                'details': {
                    'query': f'NX Nastran {solver} solver',
                    'solver_type': solver
                },
                'expected_confidence': 0.85,
                'priority': 3
            })

        # As fallback, ask user for example
        plan_steps.append({
            'action': 'ask_user_for_example',
            'description': f'Request example from user for {step.action}',
            'details': {
                'prompt': f"Could you provide an example of {step.action.replace('_', ' ')}?",
                'suggested_file_types': self._get_suggested_file_types(step.domain),
                'params_needed': step.params
            },
            'expected_confidence': 0.95,  # User examples have high confidence
            'priority': 4
        })

        return plan_steps

    def _get_suggested_file_types(self, domain: str) -> List[str]:
        """Get suggested file types for user examples based on domain."""
        suggestions = {
            'materials': ['.xml', '.mtl'],
            'geometry': ['.py', '.prt'],
            'loads_bc': ['.py', '.xml'],
            'mesh': ['.py', '.dat'],
            'result_extraction': ['.py', '.txt'],
            'optimization': ['.py', '.json']
        }
        return suggestions.get(domain, ['.py', '.txt'])

    def get_plan_summary(self, plan: List[Dict[str, Any]]) -> str:
        """Get human-readable summary of research plan."""
        if not plan:
            return "No research needed - all capabilities are known!"

        lines = [
            "Targeted Research Plan",
            "=" * 80,
            "",
            f"Research steps needed: {len(plan)}",
            ""
        ]

        current_gap = None
        for i, step in enumerate(plan, 1):
            # Group by action for clarity
            if step['action'] != current_gap:
                current_gap = step['action']
                lines.append(f"\nStep {i}: {step['description']}")
                lines.append("-" * 80)
            else:
                lines.append(f"\nStep {i}: {step['description']}")

            lines.append(f"  Action: {step['action']}")

            if 'details' in step:
                if 'capability' in step['details']:
                    lines.append(f"  Study: {step['details']['capability']}")
                if 'query' in step['details']:
                    lines.append(f"  Query: \"{step['details']['query']}\"")
                if 'prompt' in step['details']:
                    lines.append(f"  Prompt: \"{step['details']['prompt']}\"")

            lines.append(f"  Expected confidence: {step['expected_confidence']:.0%}")

        lines.append("")
        lines.append("=" * 80)

        # Add strategic summary
        lines.append("\nResearch Strategy:")
        lines.append("-" * 80)

        has_existing_code = any(s['action'] == 'read_existing_code' for s in plan)
        if has_existing_code:
            lines.append("  - Will adapt from existing similar code patterns")
            lines.append("  - Lower risk: Can follow proven implementation")
        else:
            lines.append("  - New domain: Will need to research from scratch")
            lines.append("  - Higher risk: No existing patterns to follow")

        return "\n".join(lines)


def main():
    """Test the targeted research planner."""
    from optimization_engine.codebase_analyzer import CodebaseCapabilityAnalyzer
    from optimization_engine.workflow_decomposer import WorkflowDecomposer
    from optimization_engine.capability_matcher import CapabilityMatcher

    print("Targeted Research Planner Test")
    print("=" * 80)
    print()

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

    # 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()

    # Full pipeline
    print("Phase 2.5 Pipeline:")
    print("-" * 80)
    print("1. Decompose workflow...")
    steps = decomposer.decompose(test_request)
    print(f"   Found {len(steps)} workflow steps")

    print("\n2. Match to codebase capabilities...")
    match = matcher.match(steps)
    print(f"   Known: {len(match.known_steps)}/{len(steps)}")
    print(f"   Unknown: {len(match.unknown_steps)}/{len(steps)}")
    print(f"   Overall confidence: {match.overall_confidence:.0%}")

    print("\n3. Create targeted research plan...")
    plan = planner.plan(match)
    print(f"   Generated {len(plan)} research steps")

    print("\n" + "=" * 80)
    print()

    # Display the plan
    print(planner.get_plan_summary(plan))

    # Show what's being researched
    print("\n\nWhat will be researched:")
    print("-" * 80)
    for unknown_step in match.unknown_steps:
        step = unknown_step.step
        print(f"  Missing: {step.action} ({step.domain})")
        print(f"  Required params: {step.params}")
        if unknown_step.similar_capabilities:
            print(f"  Can adapt from: {', '.join(unknown_step.similar_capabilities)}")
        print()

    print("\nWhat will NOT be researched (already known):")
    print("-" * 80)
    for known_step in match.known_steps:
        step = known_step.step
        print(f"  - {step.action} ({step.domain})")
    print()


if __name__ == '__main__':
    main()