Atomizer/studies/bracket_stiffness_optimization/bracket_stiffness_extractor.py

"""
Bracket Stiffness Optimization - Results Extractor
==================================================

This extractor uses the generic tools from optimization_engine/extractors/ to:
1. Execute NX journal to export displacement field
2. Extract z-displacement from field data
3. Extract applied force from OP2 file
4. Calculate stiffness (k = F/δ)
5. Extract mass from OP2 file

This is a thin wrapper around generic extractors - all the heavy lifting
is done by reusable components.
"""

import sys
from pathlib import Path
from typing import Dict, Any, Tuple

# Add project root to path
project_root = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(project_root))

from optimization_engine.extractors.stiffness_calculator import StiffnessCalculator
from optimization_engine.extractors.bdf_mass_extractor import BDFMassExtractor
from optimization_engine.nx_solver import NXSolver

# Import central configuration
import config as atomizer_config


class BracketStiffnessExtractor:
    """
    Bracket-specific extractor that orchestrates generic tools.

    Extracts:
    - Stiffness (N/mm) from force and z-displacement
    - Mass (kg) from OP2 grid point weight
    """

    def __init__(
        self,
        model_dir: Path,
        sim_file: str = "Bracket_sim1.sim",
        export_journal: str = "export_displacement_field.py",
    ):
        """
        Args:
            model_dir: Directory containing model files
            sim_file: Name of .sim file
            export_journal: Name of journal that exports displacement field
        """
        self.model_dir = Path(model_dir)
        self.sim_file = self.model_dir / sim_file
        self.export_journal = self.model_dir / export_journal
        self.sim_base = Path(sim_file).stem  # e.g., "Bracket_sim1" -> "Bracket_sim1"

        # Expected output files from NX
        self.field_file = self.model_dir / "export_field_dz.fld"
        # NX creates OP2 with lowercase base name and solution suffix
        self.op2_file = self.model_dir / f"{self.sim_base.lower()}-solution_1.op2"
        # BDF/DAT file for mass extraction
        self.dat_file = self.model_dir / f"{self.sim_base.lower()}-solution_1.dat"

    def extract_results(self) -> Dict[str, Any]:
        """
        Extract stiffness and mass from FEA results.

        Returns:
            dict: {
                'stiffness': stiffness value (N/mm),
                'mass': mass in kg,
                'mass_g': mass in grams,
                'displacement': max z-displacement (mm),
                'force': applied force (N),
                'compliance': inverse stiffness (mm/N),
                'objectives': {
                    'stiffness': value for maximization,
                    'mass': value for constraint checking
                }
            }
        """
        # Step 1: Execute NX journal to export displacement field
        print(f"Executing journal to export displacement field...")
        self._export_displacement_field()

        # Verify field file was created
        if not self.field_file.exists():
            raise FileNotFoundError(f"Field file not created: {self.field_file}")

        # Verify OP2 file exists
        if not self.op2_file.exists():
            raise FileNotFoundError(f"OP2 file not found: {self.op2_file}")

        # Step 2: Calculate stiffness using generic calculator
        print(f"Calculating stiffness...")
        stiffness_calc = StiffnessCalculator(
            field_file=str(self.field_file),
            op2_file=str(self.op2_file),
            force_component="fz",  # Z-direction force
            displacement_component="z",  # Z-displacement
            displacement_aggregation="max_abs",  # Maximum absolute displacement
            applied_force=1000.0  # Applied load is 1000N (constant for this model)
        )

        stiffness_results = stiffness_calc.calculate()

        # Step 3: Extract mass from BDF/DAT file
        print(f"Extracting mass from BDF...")
        if not self.dat_file.exists():
            raise FileNotFoundError(f"DAT file not found: {self.dat_file}")

        bdf_extractor = BDFMassExtractor(bdf_file=str(self.dat_file))
        mass_results = bdf_extractor.extract_mass()

        # Step 4: Combine results
        results = {
            'stiffness': stiffness_results['stiffness'],
            'mass': mass_results['mass_kg'],
            'mass_g': mass_results['mass_g'],
            'displacement': stiffness_results['displacement'],
            'force': stiffness_results['force'],
            'compliance': stiffness_results['compliance'],
            'objectives': {
                'stiffness': stiffness_results['stiffness'],  # Maximize
                'mass': mass_results['mass_kg']  # Constrain ≤ 0.2 kg
            },
            'details': {
                'stiffness_stats': stiffness_results['displacement_stats'],
                'mass_cg': mass_results.get('cg'),
                'units': stiffness_results['units']
            }
        }

        print(f"\n[OK] Stiffness: {results['stiffness']:.2f} N/mm")
        print(f"[OK] Mass: {results['mass']:.6f} kg ({results['mass_g']:.2f} g)")
        print(f"[OK] Displacement: {results['displacement']:.6f} mm")
        print(f"[OK] Force: {results['force']:.2f} N")

        return results

    def _export_displacement_field(self):
        """
        Execute NX journal to export displacement field.

        The journal should:
        1. Open the simulation
        2. Export ResultProbe to field file (.fld)
        3. Save and close
        """
        if not self.export_journal.exists():
            raise FileNotFoundError(f"Export journal not found: {self.export_journal}")

        # Use NXSolver to execute journal
        # Note: This assumes NXSolver can run journals in non-solve mode
        # If not, we'll need to create a separate journal runner
        try:
            from optimization_engine.nx_solver import run_journal
            run_journal(str(self.export_journal))
        except ImportError:
            # Fallback: Execute journal directly via NX command line
            import subprocess
            nx_exe = atomizer_config.NX_RUN_JOURNAL
            if Path(nx_exe).exists():
                # Note: NX's run_journal.exe may return non-zero even on success due to sys.exit() handling
                # We check for field file existence instead of return code
                result = subprocess.run([nx_exe, str(self.export_journal)], capture_output=True, text=True)
                # If field file doesn't exist after running journal, something went wrong
                if not self.field_file.exists():
                    raise RuntimeError(
                        f"Journal execution completed but field file not created: {self.field_file}\n"
                        f"Journal output:\n{result.stdout}\n{result.stderr}"
                    )
            else:
                raise RuntimeError(
                    f"Cannot execute journal. NX executable not found: {nx_exe}\n"
                    f"Please execute journal manually: {self.export_journal}"
                )


def extract_bracket_stiffness(
    model_dir: str,
    sim_file: str = "Bracket_sim1.sim"
) -> Tuple[float, float]:
    """
    Convenience function to extract stiffness and mass.

    Args:
        model_dir: Directory containing bracket model files
        sim_file: Name of simulation file

    Returns:
        (stiffness, mass): Stiffness in N/mm, mass in kg
    """
    extractor = BracketStiffnessExtractor(
        model_dir=Path(model_dir),
        sim_file=sim_file
    )
    results = extractor.extract_results()
    return results['stiffness'], results['mass']


if __name__ == "__main__":
    # Example usage / testing
    import sys

    if len(sys.argv) > 1:
        model_dir = sys.argv[1]
    else:
        # Default to study model directory
        model_dir = Path(__file__).parent / "1_setup" / "model"

    print(f"Testing bracket stiffness extractor...")
    print(f"Model directory: {model_dir}\n")

    extractor = BracketStiffnessExtractor(model_dir=model_dir)

    try:
        results = extractor.extract_results()

        print("\n" + "="*60)
        print("EXTRACTION RESULTS")
        print("="*60)
        print(f"Stiffness:    {results['stiffness']:.2f} N/mm")
        print(f"Mass:         {results['mass']:.6f} kg ({results['mass_g']:.2f} g)")
        print(f"Displacement: {results['displacement']:.6f} mm")
        print(f"Force:        {results['force']:.2f} N")
        print(f"Compliance:   {results['compliance']:.6e} mm/N")
        print("="*60)

        # Check constraint
        max_mass_kg = 0.2
        if results['mass'] <= max_mass_kg:
            print(f"[OK] Mass constraint satisfied: {results['mass']:.6f} kg <= {max_mass_kg} kg")
        else:
            print(f"[X] Mass constraint violated: {results['mass']:.6f} kg > {max_mass_kg} kg")

    except Exception as e:
        print(f"\n[X] Extraction failed: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)