Atomizer/optimization_engine/config/builder.py

"""
Optimization Configuration Builder

Helps users build multi-objective optimization configurations by:
1. Discovering available design variables from FEA model
2. Listing available objectives and constraints
3. Creating structured optimization_config.json

Supports:
- Multi-objective optimization (minimize weight + stress simultaneously)
- Constraints (max displacement, stress limits, mass limits)
- User selection of which objectives/constraints to apply
"""

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


class OptimizationConfigBuilder:
    """
    Interactive builder for optimization configurations.

    Workflow:
    1. Discover model capabilities (design variables, analysis type)
    2. Present available objectives/constraints to user
    3. Build configuration based on user selections
    """

    # Available objectives that can be extracted from OP2 files
    AVAILABLE_OBJECTIVES = {
        'minimize_mass': {
            'description': 'Minimize total mass (weight reduction)',
            'extractor': 'mass_extractor',
            'metric': 'total_mass',
            'units': 'kg',
            'direction': 'minimize',
            'typical_weight': 5.0  # Higher priority in multi-objective
        },
        'minimize_max_stress': {
            'description': 'Minimize maximum von Mises stress',
            'extractor': 'stress_extractor',
            'metric': 'max_von_mises',
            'units': 'MPa',
            'direction': 'minimize',
            'typical_weight': 10.0  # Very important - failure prevention
        },
        'minimize_max_displacement': {
            'description': 'Minimize maximum displacement (increase stiffness)',
            'extractor': 'displacement_extractor',
            'metric': 'max_displacement',
            'units': 'mm',
            'direction': 'minimize',
            'typical_weight': 3.0
        },
        'minimize_volume': {
            'description': 'Minimize total volume (material usage)',
            'extractor': 'volume_extractor',
            'metric': 'total_volume',
            'units': 'mm^3',
            'direction': 'minimize',
            'typical_weight': 4.0
        }
    }

    # Available constraints
    AVAILABLE_CONSTRAINTS = {
        'max_stress_limit': {
            'description': 'Maximum allowable von Mises stress',
            'extractor': 'stress_extractor',
            'metric': 'max_von_mises',
            'units': 'MPa',
            'typical_value': 200.0,  # Below yield strength with safety factor
            'constraint_type': 'upper_bound'
        },
        'max_displacement_limit': {
            'description': 'Maximum allowable displacement',
            'extractor': 'displacement_extractor',
            'metric': 'max_displacement',
            'units': 'mm',
            'typical_value': 1.0,  # Stiffness requirement
            'constraint_type': 'upper_bound'
        },
        'min_mass_limit': {
            'description': 'Minimum required mass (structural integrity)',
            'extractor': 'mass_extractor',
            'metric': 'total_mass',
            'units': 'kg',
            'typical_value': 0.3,
            'constraint_type': 'lower_bound'
        },
        'max_mass_limit': {
            'description': 'Maximum allowable mass (weight budget)',
            'extractor': 'mass_extractor',
            'metric': 'total_mass',
            'units': 'kg',
            'typical_value': 0.5,
            'constraint_type': 'upper_bound'
        }
    }

    def __init__(self, model_discovery_result: Dict[str, Any]):
        """
        Initialize with model discovery results.

        Args:
            model_discovery_result: Output from discover_fea_model()
        """
        self.model_info = model_discovery_result
        self.config = {
            'design_variables': [],
            'objectives': [],
            'constraints': [],
            'optimization_settings': {
                'n_trials': 100,
                'sampler': 'TPE',
                'n_startup_trials': 20
            }
        }

    def list_available_design_variables(self) -> List[Dict[str, Any]]:
        """
        List all available design variables from model.

        Returns:
            List of design variable options
        """
        if 'expressions' not in self.model_info:
            return []

        design_vars = []
        for expr in self.model_info['expressions']:
            if expr['value'] is not None:  # Only variables with known values
                design_vars.append({
                    'name': expr['name'],
                    'current_value': expr['value'],
                    'units': expr['units'],
                    'type': expr.get('type', 'Unknown'),
                    'suggested_bounds': self._suggest_bounds(expr)
                })

        return design_vars

    def _suggest_bounds(self, expr: Dict[str, Any]) -> tuple:
        """
        Suggest reasonable optimization bounds for a design variable.

        Args:
            expr: Expression dictionary

        Returns:
            (lower_bound, upper_bound)
        """
        value = expr['value']
        expr_type = expr.get('type', '').lower()

        if 'angle' in expr_type or 'degrees' in expr.get('units', '').lower():
            # Angles: ±15 degrees
            return (max(0, value - 15), min(180, value + 15))
        elif 'thickness' in expr['name'].lower() or 'dimension' in expr_type:
            # Dimensions: ±30%
            return (value * 0.7, value * 1.3)
        elif 'radius' in expr['name'].lower() or 'diameter' in expr['name'].lower():
            # Radii/diameters: ±25%
            return (value * 0.75, value * 1.25)
        else:
            # Default: ±20%
            return (value * 0.8, value * 1.2)

    def list_available_objectives(self) -> Dict[str, Dict[str, Any]]:
        """
        List all available optimization objectives.

        Returns:
            Dictionary of objective options
        """
        return self.AVAILABLE_OBJECTIVES.copy()

    def list_available_constraints(self) -> Dict[str, Dict[str, Any]]:
        """
        List all available constraints.

        Returns:
            Dictionary of constraint options
        """
        return self.AVAILABLE_CONSTRAINTS.copy()

    def add_design_variable(self, name: str, lower_bound: float, upper_bound: float):
        """
        Add a design variable to the configuration.

        Args:
            name: Expression name from model
            lower_bound: Minimum value
            upper_bound: Maximum value
        """
        # Verify variable exists in model
        expr = next((e for e in self.model_info['expressions'] if e['name'] == name), None)
        if not expr:
            raise ValueError(f"Design variable '{name}' not found in model")

        self.config['design_variables'].append({
            'name': name,
            'type': 'continuous',
            'bounds': [lower_bound, upper_bound],
            'units': expr.get('units', ''),
            'initial_value': expr['value']
        })

    def add_objective(self, objective_key: str, weight: float = None, target: float = None):
        """
        Add an objective to the configuration.

        Args:
            objective_key: Key from AVAILABLE_OBJECTIVES
            weight: Importance weight (for multi-objective)
            target: Target value (optional, for goal programming)
        """
        if objective_key not in self.AVAILABLE_OBJECTIVES:
            raise ValueError(f"Unknown objective: {objective_key}")

        obj_info = self.AVAILABLE_OBJECTIVES[objective_key]

        objective = {
            'name': objective_key,
            'description': obj_info['description'],
            'extractor': obj_info['extractor'],
            'metric': obj_info['metric'],
            'direction': obj_info['direction'],
            'weight': weight or obj_info['typical_weight']
        }

        if target is not None:
            objective['target'] = target

        self.config['objectives'].append(objective)

    def add_constraint(self, constraint_key: str, limit_value: float):
        """
        Add a constraint to the configuration.

        Args:
            constraint_key: Key from AVAILABLE_CONSTRAINTS
            limit_value: Constraint limit value
        """
        if constraint_key not in self.AVAILABLE_CONSTRAINTS:
            raise ValueError(f"Unknown constraint: {constraint_key}")

        const_info = self.AVAILABLE_CONSTRAINTS[constraint_key]

        constraint = {
            'name': constraint_key,
            'description': const_info['description'],
            'extractor': const_info['extractor'],
            'metric': const_info['metric'],
            'type': const_info['constraint_type'],
            'limit': limit_value,
            'units': const_info['units']
        }

        self.config['constraints'].append(constraint)

    def set_optimization_settings(self, n_trials: int = None, sampler: str = None):
        """
        Configure optimization algorithm settings.

        Args:
            n_trials: Number of optimization iterations
            sampler: 'TPE', 'CMAES', 'GP', etc.
        """
        if n_trials:
            self.config['optimization_settings']['n_trials'] = n_trials
        if sampler:
            self.config['optimization_settings']['sampler'] = sampler

    def build(self) -> Dict[str, Any]:
        """
        Build and validate the configuration.

        Returns:
            Complete optimization configuration
        """
        # Validation
        if not self.config['design_variables']:
            raise ValueError("At least one design variable is required")

        if not self.config['objectives']:
            raise ValueError("At least one objective is required")

        # Add metadata
        self.config['model_info'] = {
            'sim_file': self.model_info.get('sim_file', ''),
            'solutions': self.model_info.get('solutions', [])
        }

        return self.config

    def save(self, output_path: Path):
        """
        Save configuration to JSON file.

        Args:
            output_path: Path to save configuration
        """
        config = self.build()

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

        print(f"Configuration saved to: {output_path}")

    def print_summary(self):
        """Print a human-readable summary of the configuration."""
        print("\n" + "="*60)
        print("OPTIMIZATION CONFIGURATION SUMMARY")
        print("="*60)

        print(f"\nModel: {self.model_info.get('sim_file', 'Unknown')}")

        print(f"\nDesign Variables ({len(self.config['design_variables'])}):")
        for dv in self.config['design_variables']:
            print(f"  • {dv['name']}: [{dv['bounds'][0]:.2f}, {dv['bounds'][1]:.2f}] {dv['units']}")

        print(f"\nObjectives ({len(self.config['objectives'])}):")
        for obj in self.config['objectives']:
            print(f"  • {obj['description']} (weight: {obj['weight']:.1f})")

        print(f"\nConstraints ({len(self.config['constraints'])}):")
        for const in self.config['constraints']:
            operator = '<=' if const['type'] == 'upper_bound' else '>='
            print(f"  • {const['description']}: {const['metric']} {operator} {const['limit']} {const['units']}")

        print(f"\nOptimization Settings:")
        print(f"  • Trials: {self.config['optimization_settings']['n_trials']}")
        print(f"  • Sampler: {self.config['optimization_settings']['sampler']}")

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


# Example usage
if __name__ == "__main__":
    from optimization_engine.model_discovery import discover_fea_model

    # Step 1: Discover model
    print("Step 1: Discovering FEA model...")
    model_result = discover_fea_model("tests/Bracket_sim1.sim")

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

    # Step 3: Show available options
    print("\n" + "="*60)
    print("AVAILABLE DESIGN VARIABLES:")
    print("="*60)
    for dv in builder.list_available_design_variables():
        print(f"\n• {dv['name']}")
        print(f"  Current value: {dv['current_value']} {dv['units']}")
        print(f"  Suggested bounds: {dv['suggested_bounds']}")

    print("\n" + "="*60)
    print("AVAILABLE OBJECTIVES:")
    print("="*60)
    for key, obj in builder.list_available_objectives().items():
        print(f"\n• {key}")
        print(f"  Description: {obj['description']}")
        print(f"  Default weight: {obj['typical_weight']}")

    print("\n" + "="*60)
    print("AVAILABLE CONSTRAINTS:")
    print("="*60)
    for key, const in builder.list_available_constraints().items():
        print(f"\n• {key}")
        print(f"  Description: {const['description']}")
        print(f"  Typical value: {const['typical_value']} {const['units']}")

    # Step 4: Build a multi-objective configuration
    print("\n" + "="*60)
    print("BUILDING CONFIGURATION:")
    print("="*60)

    # Add design variables
    builder.add_design_variable('tip_thickness', 15.0, 25.0)
    builder.add_design_variable('support_angle', 20.0, 40.0)
    builder.add_design_variable('support_blend_radius', 5.0, 15.0)

    # Add objectives: minimize weight AND minimize stress
    builder.add_objective('minimize_mass', weight=5.0)
    builder.add_objective('minimize_max_stress', weight=10.0)

    # Add constraints: max displacement < 1.0 mm, max stress < 200 MPa
    builder.add_constraint('max_displacement_limit', limit_value=1.0)
    builder.add_constraint('max_stress_limit', limit_value=200.0)

    # Set optimization settings
    builder.set_optimization_settings(n_trials=150, sampler='TPE')

    # Print summary
    builder.print_summary()

    # Save configuration
    builder.save(Path('optimization_config.json'))

    print("\nConfiguration ready for optimization!")