"""
pyNastran Research Agent - Phase 3

Automated research and code generation for OP2 result extraction using pyNastran.

This agent:
1. Searches pyNastran documentation
2. Finds relevant APIs for extraction tasks
3. Generates executable Python code for extractors
4. Stores patterns in knowledge base

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

from typing import Dict, Any, List, Optional
from dataclasses import dataclass
from pathlib import Path
import json


@dataclass
class ExtractionPattern:
    """Represents a learned pattern for OP2 extraction."""
    name: str
    description: str
    element_type: Optional[str]  # e.g., 'CBAR', 'CQUAD4', None for general
    result_type: str  # 'force', 'stress', 'displacement', 'strain'
    code_template: str
    api_path: str  # e.g., 'model.cbar_force[subcase]'
    data_structure: str  # Description of data array structure
    examples: List[str]  # Example usage


class PyNastranResearchAgent:
    """
    Research agent for pyNastran documentation and code generation.

    Uses a combination of:
    - Pre-learned patterns from documentation
    - WebFetch for dynamic lookup (future)
    - Knowledge base caching
    """

    def __init__(self, knowledge_base_path: Optional[Path] = None):
        """
        Initialize the research agent.

        Args:
            knowledge_base_path: Path to store learned patterns
        """
        if knowledge_base_path is None:
            knowledge_base_path = Path(__file__).parent.parent / "knowledge_base" / "pynastran_patterns"

        self.knowledge_base_path = Path(knowledge_base_path)
        self.knowledge_base_path.mkdir(parents=True, exist_ok=True)

        # Initialize with core patterns from documentation research
        self.patterns = self._initialize_core_patterns()

    def _initialize_core_patterns(self) -> Dict[str, ExtractionPattern]:
        """Initialize core extraction patterns from pyNastran docs."""
        patterns = {}

        # Displacement extraction
        patterns['displacement'] = ExtractionPattern(
            name='displacement',
            description='Extract displacement results',
            element_type=None,
            result_type='displacement',
            code_template='''def extract_displacement(op2_file: Path, subcase: int = 1):
    """Extract displacement results from OP2 file."""
    from pyNastran.op2.op2 import OP2
    import numpy as np

    model = OP2()
    model.read_op2(str(op2_file))

    disp = model.displacements[subcase]
    itime = 0  # static case

    # Extract translation components
    txyz = disp.data[itime, :, :3]  # [tx, ty, tz]

    # Calculate total displacement
    total_disp = np.linalg.norm(txyz, axis=1)
    max_disp = np.max(total_disp)

    # Get node info
    node_ids = [nid for (nid, grid_type) in disp.node_gridtype]
    max_disp_node = node_ids[np.argmax(total_disp)]

    return {
        'max_displacement': float(max_disp),
        'max_disp_node': int(max_disp_node),
        'max_disp_x': float(np.max(np.abs(txyz[:, 0]))),
        'max_disp_y': float(np.max(np.abs(txyz[:, 1]))),
        'max_disp_z': float(np.max(np.abs(txyz[:, 2])))
    }''',
            api_path='model.displacements[subcase]',
            data_structure='data[itime, :, :6] where :6=[tx, ty, tz, rx, ry, rz]',
            examples=['max_disp = extract_displacement(Path("results.op2"))']
        )

        # Stress extraction (solid elements)
        patterns['solid_stress'] = ExtractionPattern(
            name='solid_stress',
            description='Extract stress from solid elements (CTETRA, CHEXA)',
            element_type='CTETRA',
            result_type='stress',
            code_template='''def extract_solid_stress(op2_file: Path, subcase: int = 1, element_type: str = 'ctetra'):
    """Extract stress from solid elements."""
    from pyNastran.op2.op2 import OP2
    import numpy as np

    model = OP2()
    model.read_op2(str(op2_file))

    # Get stress object for element type
    # In pyNastran, stress is stored in model.op2_results.stress
    stress_attr = f"{element_type}_stress"

    if not hasattr(model, 'op2_results') or not hasattr(model.op2_results, 'stress'):
        raise ValueError(f"No stress results in OP2")

    stress_obj = model.op2_results.stress
    if not hasattr(stress_obj, stress_attr):
        raise ValueError(f"No {element_type} stress results in OP2")

    stress = getattr(stress_obj, stress_attr)[subcase]
    itime = 0

    # Extract von Mises if available
    if stress.is_von_mises:  # Property, not method
        von_mises = stress.data[itime, :, 9]  # Column 9 is von Mises
        max_stress = float(np.max(von_mises))

        # Get element info
        element_ids = [eid for (eid, node) in stress.element_node]
        max_stress_elem = element_ids[np.argmax(von_mises)]

        return {
            'max_von_mises': max_stress,
            'max_stress_element': int(max_stress_elem)
        }
    else:
        raise ValueError("von Mises stress not available")''',
            api_path='model.ctetra_stress[subcase] or model.chexa_stress[subcase]',
            data_structure='data[itime, :, 10] where column 9=von_mises',
            examples=['stress = extract_solid_stress(Path("results.op2"), element_type="ctetra")']
        )

        # CBAR force extraction
        patterns['cbar_force'] = ExtractionPattern(
            name='cbar_force',
            description='Extract forces from CBAR elements',
            element_type='CBAR',
            result_type='force',
            code_template='''def extract_cbar_force(op2_file: Path, subcase: int = 1, direction: str = 'Z'):
    """
    Extract forces from CBAR elements.

    Args:
        op2_file: Path to OP2 file
        subcase: Subcase ID
        direction: Force direction ('X', 'Y', 'Z', 'axial', 'torque')

    Returns:
        Dict with force statistics
    """
    from pyNastran.op2.op2 import OP2
    import numpy as np

    model = OP2()
    model.read_op2(str(op2_file))

    if not hasattr(model, 'cbar_force'):
        raise ValueError("No CBAR force results in OP2")

    force = model.cbar_force[subcase]
    itime = 0

    # CBAR force data structure:
    # [bending_moment_a1, bending_moment_a2,
    #  bending_moment_b1, bending_moment_b2,
    #  shear1, shear2, axial, torque]

    direction_map = {
        'shear1': 4,
        'shear2': 5,
        'axial': 6,
        'Z': 6,  # Commonly axial is Z direction
        'torque': 7
    }

    col_idx = direction_map.get(direction, direction_map.get(direction.lower(), 6))
    forces = force.data[itime, :, col_idx]

    return {
        f'max_{direction}_force': float(np.max(np.abs(forces))),
        f'avg_{direction}_force': float(np.mean(np.abs(forces))),
        f'min_{direction}_force': float(np.min(np.abs(forces))),
        'forces_array': forces.tolist()
    }''',
            api_path='model.cbar_force[subcase]',
            data_structure='data[ntimes, nelements, 8] where 8=[bm_a1, bm_a2, bm_b1, bm_b2, shear1, shear2, axial, torque]',
            examples=['forces = extract_cbar_force(Path("results.op2"), direction="Z")']
        )

        return patterns

    def research_extraction(self, request: Dict[str, Any]) -> ExtractionPattern:
        """
        Research and find/generate extraction pattern for a request.

        Args:
            request: Dict with:
                - action: e.g., 'extract_1d_element_forces'
                - domain: e.g., 'result_extraction'
                - params: {'element_types': ['CBAR'], 'result_type': 'element_force', 'direction': 'Z'}

        Returns:
            ExtractionPattern with code template
        """
        action = request.get('action', '')
        params = request.get('params', {})

        # Determine result type
        if 'displacement' in action.lower():
            return self.patterns['displacement']

        elif 'stress' in action.lower():
            element_types = params.get('element_types', [])
            if any(et in ['CTETRA', 'CHEXA', 'CPENTA'] for et in element_types):
                return self.patterns['solid_stress']
            # Could add plate stress pattern here
            return self.patterns['solid_stress']  # Default to solid for now

        elif 'force' in action.lower() or 'element_force' in params.get('result_type', ''):
            element_types = params.get('element_types', [])
            if 'CBAR' in element_types or '1d' in action.lower():
                return self.patterns['cbar_force']

        # Fallback: return generic pattern
        return self._generate_generic_pattern(request)

    def _generate_generic_pattern(self, request: Dict[str, Any]) -> ExtractionPattern:
        """Generate a generic extraction pattern as fallback."""
        return ExtractionPattern(
            name='generic_extraction',
            description=f"Generic extraction for {request.get('action', 'unknown')}",
            element_type=None,
            result_type='unknown',
            code_template='''def extract_generic(op2_file: Path):
    """Generic OP2 extraction - needs customization."""
    from pyNastran.op2.op2 import OP2

    model = OP2()
    model.read_op2(str(op2_file))

    # TODO: Customize extraction based on requirements
    # Available: model.displacements, model.ctetra_stress, etc.
    # Use model.get_op2_stats() to see available results

    return {'result': None}''',
            api_path='model.<result_type>[subcase]',
            data_structure='Varies by result type',
            examples=['# Needs customization']
        )

    def generate_extractor_code(self, request: Dict[str, Any]) -> str:
        """
        Generate complete extractor code for a request.

        Args:
            request: Extraction request from Phase 2.7 LLM

        Returns:
            Complete Python code as string
        """
        pattern = self.research_extraction(request)

        # Generate module header
        description = request.get('description', pattern.description)

        code = f'''"""
{description}
Auto-generated by Atomizer Phase 3 - pyNastran Research Agent

Pattern: {pattern.name}
Element Type: {pattern.element_type or 'General'}
Result Type: {pattern.result_type}
API: {pattern.api_path}
"""

from pathlib import Path
from typing import Dict, Any
import numpy as np
from pyNastran.op2.op2 import OP2


{pattern.code_template}


if __name__ == '__main__':
    # Example usage
    import sys
    if len(sys.argv) > 1:
        op2_file = Path(sys.argv[1])
        result = {pattern.code_template.split('(')[0].split()[-1]}(op2_file)
        print(f"Extraction result: {{result}}")
    else:
        print("Usage: python {{sys.argv[0]}} <op2_file>")
'''

        return code

    def save_pattern(self, pattern: ExtractionPattern):
        """Save a pattern to the knowledge base."""
        pattern_file = self.knowledge_base_path / f"{pattern.name}.json"

        pattern_dict = {
            'name': pattern.name,
            'description': pattern.description,
            'element_type': pattern.element_type,
            'result_type': pattern.result_type,
            'code_template': pattern.code_template,
            'api_path': pattern.api_path,
            'data_structure': pattern.data_structure,
            'examples': pattern.examples
        }

        with open(pattern_file, 'w') as f:
            json.dump(pattern_dict, f, indent=2)

    def load_pattern(self, name: str) -> Optional[ExtractionPattern]:
        """Load a pattern from the knowledge base."""
        pattern_file = self.knowledge_base_path / f"{name}.json"

        if not pattern_file.exists():
            return None

        with open(pattern_file, 'r') as f:
            data = json.load(f)

        return ExtractionPattern(**data)


def main():
    """Test the pyNastran research agent."""
    print("=" * 80)
    print("Phase 3: pyNastran Research Agent Test")
    print("=" * 80)
    print()

    agent = PyNastranResearchAgent()

    # Test request: CBAR force extraction (from Phase 2.7 example)
    test_request = {
        "action": "extract_1d_element_forces",
        "domain": "result_extraction",
        "description": "Extract element forces from CBAR in Z direction from OP2",
        "params": {
            "element_types": ["CBAR"],
            "result_type": "element_force",
            "direction": "Z"
        }
    }

    print("Test Request:")
    print(f"  Action: {test_request['action']}")
    print(f"  Description: {test_request['description']}")
    print()

    print("1. Researching extraction pattern...")
    pattern = agent.research_extraction(test_request)
    print(f"   Found pattern: {pattern.name}")
    print(f"   API path: {pattern.api_path}")
    print()

    print("2. Generating extractor code...")
    code = agent.generate_extractor_code(test_request)
    print()

    print("=" * 80)
    print("Generated Extractor Code:")
    print("=" * 80)
    print(code)

    # Save to file
    output_file = Path("generated_extractors") / "cbar_force_extractor.py"
    output_file.parent.mkdir(exist_ok=True)
    with open(output_file, 'w') as f:
        f.write(code)

    print()
    print(f"[OK] Saved to: {output_file}")


if __name__ == '__main__':
    main()