Atomizer/mcp_server/tools/optimization_config.py

"""
MCP Tool: Build Optimization Configuration

Wraps the OptimizationConfigBuilder to create an MCP-compatible tool
that helps LLMs guide users through building optimization configurations.

This tool:
1. Discovers the FEA model (design variables)
2. Lists available objectives and constraints
3. Builds a complete optimization_config.json based on user selections
"""

from pathlib import Path
from typing import Dict, Any, List, Optional
import json
import sys

# Add project root to path for imports
project_root = Path(__file__).parent.parent.parent
sys.path.insert(0, str(project_root))

from optimization_engine.optimization_config_builder import OptimizationConfigBuilder
from mcp_server.tools.model_discovery import discover_fea_model


def build_optimization_config(
    sim_file_path: str,
    design_variables: List[Dict[str, Any]],
    objectives: List[Dict[str, Any]],
    constraints: Optional[List[Dict[str, Any]]] = None,
    optimization_settings: Optional[Dict[str, Any]] = None,
    output_path: Optional[str] = None
) -> Dict[str, Any]:
    """
    MCP Tool: Build Optimization Configuration

    Creates a complete optimization configuration file from user selections.

    Args:
        sim_file_path: Absolute path to .sim file
        design_variables: List of design variable definitions
            [
                {
                    'name': 'tip_thickness',
                    'lower_bound': 15.0,
                    'upper_bound': 25.0
                },
                ...
            ]
        objectives: List of objective definitions
            [
                {
                    'objective_key': 'minimize_mass',
                    'weight': 5.0,  # optional
                    'target': None   # optional, for goal programming
                },
                ...
            ]
        constraints: Optional list of constraint definitions
            [
                {
                    'constraint_key': 'max_stress_limit',
                    'limit_value': 200.0
                },
                ...
            ]
        optimization_settings: Optional dict with algorithm settings
            {
                'n_trials': 100,
                'sampler': 'TPE'
            }
        output_path: Optional path to save config JSON.
                    Defaults to 'optimization_config.json' in sim file directory

    Returns:
        Dictionary with status and configuration details

    Example:
        >>> result = build_optimization_config(
        ...     sim_file_path="C:/Projects/Bracket/analysis.sim",
        ...     design_variables=[
        ...         {'name': 'tip_thickness', 'lower_bound': 15.0, 'upper_bound': 25.0}
        ...     ],
        ...     objectives=[
        ...         {'objective_key': 'minimize_mass', 'weight': 5.0}
        ...     ],
        ...     constraints=[
        ...         {'constraint_key': 'max_stress_limit', 'limit_value': 200.0}
        ...     ]
        ... )
    """
    try:
        # Step 1: Discover model
        model_result = discover_fea_model(sim_file_path)

        if model_result['status'] != 'success':
            return {
                'status': 'error',
                'error_type': 'model_discovery_failed',
                'message': model_result.get('message', 'Failed to discover FEA model'),
                'suggestion': model_result.get('suggestion', 'Check that the .sim file is valid')
            }

        # Step 2: Create builder
        builder = OptimizationConfigBuilder(model_result)

        # Step 3: Validate and add design variables
        available_vars = {dv['name']: dv for dv in builder.list_available_design_variables()}

        for dv in design_variables:
            name = dv['name']
            if name not in available_vars:
                return {
                    'status': 'error',
                    'error_type': 'invalid_design_variable',
                    'message': f"Design variable '{name}' not found in model",
                    'available_variables': list(available_vars.keys()),
                    'suggestion': f"Choose from: {', '.join(available_vars.keys())}"
                }

            builder.add_design_variable(
                name=name,
                lower_bound=dv['lower_bound'],
                upper_bound=dv['upper_bound']
            )

        # Step 4: Add objectives
        available_objectives = builder.list_available_objectives()

        for obj in objectives:
            obj_key = obj['objective_key']
            if obj_key not in available_objectives:
                return {
                    'status': 'error',
                    'error_type': 'invalid_objective',
                    'message': f"Objective '{obj_key}' not recognized",
                    'available_objectives': list(available_objectives.keys()),
                    'suggestion': f"Choose from: {', '.join(available_objectives.keys())}"
                }

            builder.add_objective(
                objective_key=obj_key,
                weight=obj.get('weight'),
                target=obj.get('target')
            )

        # Step 5: Add constraints (optional)
        if constraints:
            available_constraints = builder.list_available_constraints()

            for const in constraints:
                const_key = const['constraint_key']
                if const_key not in available_constraints:
                    return {
                        'status': 'error',
                        'error_type': 'invalid_constraint',
                        'message': f"Constraint '{const_key}' not recognized",
                        'available_constraints': list(available_constraints.keys()),
                        'suggestion': f"Choose from: {', '.join(available_constraints.keys())}"
                    }

                builder.add_constraint(
                    constraint_key=const_key,
                    limit_value=const['limit_value']
                )

        # Step 6: Set optimization settings (optional)
        if optimization_settings:
            builder.set_optimization_settings(
                n_trials=optimization_settings.get('n_trials'),
                sampler=optimization_settings.get('sampler')
            )

        # Step 7: Build and validate configuration
        config = builder.build()

        # Step 8: Save to file
        if output_path is None:
            sim_path = Path(sim_file_path)
            output_path = sim_path.parent / 'optimization_config.json'
        else:
            output_path = Path(output_path)

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

        # Step 9: Return success with summary
        return {
            'status': 'success',
            'message': 'Optimization configuration created successfully',
            'config_file': str(output_path),
            'summary': {
                'design_variables': len(config['design_variables']),
                'objectives': len(config['objectives']),
                'constraints': len(config['constraints']),
                'n_trials': config['optimization_settings']['n_trials'],
                'sampler': config['optimization_settings']['sampler']
            },
            'config': config
        }

    except ValueError as e:
        return {
            'status': 'error',
            'error_type': 'validation_error',
            'message': str(e),
            'suggestion': 'Check that all required fields are provided correctly'
        }

    except Exception as e:
        return {
            'status': 'error',
            'error_type': 'unexpected_error',
            'message': str(e),
            'suggestion': 'This may be a bug. Please report this issue.'
        }


def list_optimization_options(sim_file_path: str) -> Dict[str, Any]:
    """
    Helper tool: List all available optimization options for a model.

    This is useful for LLMs to show users what they can choose from.

    Args:
        sim_file_path: Absolute path to .sim file

    Returns:
        Dictionary with all available options
    """
    try:
        # Discover model
        model_result = discover_fea_model(sim_file_path)

        if model_result['status'] != 'success':
            return model_result

        # Create builder to get options
        builder = OptimizationConfigBuilder(model_result)

        # Get all available options
        design_vars = builder.list_available_design_variables()
        objectives = builder.list_available_objectives()
        constraints = builder.list_available_constraints()

        return {
            'status': 'success',
            'sim_file': sim_file_path,
            'available_design_variables': design_vars,
            'available_objectives': objectives,
            'available_constraints': constraints,
            'model_info': {
                'solutions': model_result.get('solutions', []),
                'expression_count': len(model_result.get('expressions', []))
            }
        }

    except Exception as e:
        return {
            'status': 'error',
            'error_type': 'unexpected_error',
            'message': str(e)
        }


def format_optimization_options_for_llm(options: Dict[str, Any]) -> str:
    """
    Format optimization options for LLM consumption (Markdown).

    Args:
        options: Output from list_optimization_options()

    Returns:
        Markdown-formatted string
    """
    if options['status'] != 'success':
        return f"❌ **Error**: {options['message']}\n\n💡 {options.get('suggestion', '')}"

    md = []
    md.append(f"# Optimization Configuration Options\n")
    md.append(f"**Model**: `{options['sim_file']}`\n")

    # Design Variables
    md.append(f"## Available Design Variables ({len(options['available_design_variables'])})\n")
    if options['available_design_variables']:
        md.append("| Name | Current Value | Units | Suggested Bounds |")
        md.append("|------|---------------|-------|------------------|")
        for dv in options['available_design_variables']:
            bounds = dv['suggested_bounds']
            md.append(f"| `{dv['name']}` | {dv['current_value']} | {dv['units']} | [{bounds[0]:.2f}, {bounds[1]:.2f}] |")
    else:
        md.append("⚠️ No design variables found. Model may not be parametric.")
    md.append("")

    # Objectives
    md.append(f"## Available Objectives\n")
    for key, obj in options['available_objectives'].items():
        md.append(f"### `{key}`")
        md.append(f"- **Description**: {obj['description']}")
        md.append(f"- **Metric**: {obj['metric']} ({obj['units']})")
        md.append(f"- **Default Weight**: {obj['typical_weight']}")
        md.append(f"- **Extractor**: `{obj['extractor']}`")
        md.append("")

    # Constraints
    md.append(f"## Available Constraints\n")
    for key, const in options['available_constraints'].items():
        md.append(f"### `{key}`")
        md.append(f"- **Description**: {const['description']}")
        md.append(f"- **Metric**: {const['metric']} ({const['units']})")
        md.append(f"- **Typical Value**: {const['typical_value']}")
        md.append(f"- **Type**: {const['constraint_type']}")
        md.append(f"- **Extractor**: `{const['extractor']}`")
        md.append("")

    return "\n".join(md)


# For testing
if __name__ == "__main__":
    import sys

    if len(sys.argv) < 2:
        print("Usage: python optimization_config.py <path_to_sim_file>")
        sys.exit(1)

    sim_path = sys.argv[1]

    # Test 1: List options
    print("=" * 60)
    print("TEST 1: List Available Options")
    print("=" * 60)
    options = list_optimization_options(sim_path)
    print(format_optimization_options_for_llm(options))

    # Test 2: Build configuration
    print("\n" + "=" * 60)
    print("TEST 2: Build Optimization Configuration")
    print("=" * 60)

    result = build_optimization_config(
        sim_file_path=sim_path,
        design_variables=[
            {'name': 'tip_thickness', 'lower_bound': 15.0, 'upper_bound': 25.0},
            {'name': 'support_angle', 'lower_bound': 20.0, 'upper_bound': 40.0},
        ],
        objectives=[
            {'objective_key': 'minimize_mass', 'weight': 5.0},
            {'objective_key': 'minimize_max_stress', 'weight': 10.0}
        ],
        constraints=[
            {'constraint_key': 'max_displacement_limit', 'limit_value': 1.0},
            {'constraint_key': 'max_stress_limit', 'limit_value': 200.0}
        ],
        optimization_settings={
            'n_trials': 150,
            'sampler': 'TPE'
        }
    )

    if result['status'] == 'success':
        print(f"SUCCESS: Configuration saved to: {result['config_file']}")
        print(f"\nSummary:")
        for key, value in result['summary'].items():
            print(f"  - {key}: {value}")
    else:
        print(f"ERROR: {result['message']}")
        print(f"Suggestion: {result.get('suggestion', '')}")