Atomizer/projects/hydrotech-beam/studies/01_doe_landscape/nx_interface.py

"""NX automation interface for Hydrotech Beam optimization.

This module uses the EXISTING Atomizer optimization engine for NX integration:
- optimization_engine.nx.updater.NXParameterUpdater (expression updates via .exp import)
- optimization_engine.nx.solver.NXSolver (journal-based solving with run_journal.exe)
- optimization_engine.extractors.* (pyNastran OP2-based result extraction)

NX Expression Names (confirmed via binary introspection — CONTEXT.md):
    Design Variables:
        - beam_half_core_thickness  (mm, continuous)
        - beam_face_thickness       (mm, continuous)
        - holes_diameter            (mm, continuous)
        - hole_count                (integer, links to Pattern_p7)
    Outputs:
        - p173                      (mass in kg, body_property147.mass)
    Fixed:
        - beam_lenght               (⚠️ TYPO in NX — no 'h', 5000 mm)
        - beam_half_height          (250 mm)
        - beam_half_width           (150 mm)

References:
    CONTEXT.md — Full expression map
    OPTIMIZATION_STRATEGY.md §8.2 — Extractor requirements
"""

from __future__ import annotations

import logging
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Protocol

logger = logging.getLogger(__name__)

# Add Atomizer repo root to sys.path for imports
ATOMIZER_REPO_ROOT = Path("/home/papa/repos/Atomizer")
if str(ATOMIZER_REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(ATOMIZER_REPO_ROOT))


# ---------------------------------------------------------------------------
# Data types
# ---------------------------------------------------------------------------
@dataclass(frozen=True)
class TrialInput:
    """Design variable values for a single trial."""

    beam_half_core_thickness: float  # mm — DV1
    beam_face_thickness: float  # mm — DV2
    holes_diameter: float  # mm — DV3
    hole_count: int  # — DV4


@dataclass
class TrialResult:
    """Results extracted from NX after a trial solve.

    All values populated after a successful SOL 101 solve.
    On failure, success=False and error_message explains the failure.
    """

    success: bool
    mass: float = float("nan")  # kg — from expression `p173`
    tip_displacement: float = float("nan")  # mm — from SOL 101 results
    max_von_mises: float = float("nan")  # MPa — from SOL 101 results
    error_message: str = ""


# ---------------------------------------------------------------------------
# NX expression name constants
# ---------------------------------------------------------------------------
# ⚠️ These are EXACT NX expression names from binary introspection.
# Do NOT change spelling — `beam_lenght` has a typo (no 'h') in NX.
EXPR_HALF_CORE_THICKNESS = "beam_half_core_thickness"
EXPR_FACE_THICKNESS = "beam_face_thickness"
EXPR_HOLES_DIAMETER = "holes_diameter"
EXPR_HOLE_COUNT = "hole_count"
EXPR_MASS = "p173"  # body_property147.mass, kg
EXPR_BEAM_LENGTH = "beam_lenght"  # ⚠️ TYPO IN NX — intentional


# ---------------------------------------------------------------------------
# Interface protocol
# ---------------------------------------------------------------------------
class NXSolverInterface(Protocol):
    """Protocol for NX solver backends.

    Implementors must provide the full pipeline:
    1. Update expressions → 2. Rebuild model → 3. Solve SOL 101 → 4. Extract results
    """

    def evaluate(self, trial_input: TrialInput) -> TrialResult:
        """Run a full NX evaluation cycle for one trial.

        Args:
            trial_input: Design variable values.

        Returns:
            TrialResult with extracted outputs or failure info.
        """
        ...

    def close(self) -> None:
        """Clean up NX session resources.

        ⚠️ LAC CRITICAL: NEVER kill NX processes directly.
        Use NXSessionManager.close_nx_if_allowed() only.
        """
        ...


# ---------------------------------------------------------------------------
# Stub implementation (for development/testing without NX)
# ---------------------------------------------------------------------------
class NXStubSolver:
    """Stub NX solver for development and testing.

    Returns synthetic results based on simple analytical approximations
    of the beam behavior. NOT physically accurate — use only for
    testing the optimization pipeline.

    The stub uses rough scaling relationships:
        - Mass ∝ (core + face) and inversely with hole area
        - Displacement ∝ 1/I where I depends on core and face thickness
        - Stress ∝ M*y/I (bending stress approximation)
    """

    def __init__(self) -> None:
        """Initialize stub solver."""
        logger.warning(
            "Using NX STUB solver — results are synthetic approximations. "
            "Replace with NXOpenSolver for real evaluations."
        )

    def evaluate(self, trial_input: TrialInput) -> TrialResult:
        """Return synthetic results based on simplified beam mechanics.

        Args:
            trial_input: Design variable values.

        Returns:
            TrialResult with approximate values.
        """
        try:
            return self._compute_approximate(trial_input)
        except Exception as e:
            logger.error("Stub evaluation failed: %s", e)
            return TrialResult(
                success=False,
                error_message=f"Stub evaluation error: {e}",
            )

    def _compute_approximate(self, inp: TrialInput) -> TrialResult:
        """Simple analytical approximation of beam response.

        This is a ROUGH approximation for pipeline testing only.
        Real physics requires NX Nastran SOL 101.
        """
        import math

        # Geometry
        L = 5000.0  # mm — beam length
        b = 300.0  # mm — beam width (2 × beam_half_width)
        h_core = inp.beam_half_core_thickness  # mm — half core
        t_face = inp.beam_face_thickness  # mm — face thickness
        d_hole = inp.holes_diameter  # mm
        n_holes = inp.hole_count

        # Total height and section properties (simplified I-beam)
        h_total = 500.0  # mm — 2 × beam_half_height (fixed)

        # Approximate second moment of area (sandwich beam)
        # I ≈ b*h_total^3/12 - b*(h_total-2*t_face)^3/12 + web contribution
        h_inner = h_total - 2.0 * t_face
        I_section = (b * h_total**3 / 12.0) - (b * max(h_inner, 0.0) ** 3 / 12.0)

        # Add core contribution
        I_section += 2.0 * h_core * h_total**2 / 4.0  # approximate

        # Hole area reduction (mass)
        hole_area = n_holes * math.pi * (d_hole / 2.0) ** 2  # mm²

        # Approximate mass (steel: 7.3 g/cm³ = 7.3e-6 kg/mm³)
        rho = 7.3e-6  # kg/mm³
        # Gross cross-section area (very simplified)
        A_gross = 2.0 * b * t_face + 2.0 * h_core * h_total
        # Remove holes from web
        web_thickness = 2.0 * h_core  # approximate web thickness
        A_holes = n_holes * math.pi * (d_hole / 2.0) ** 2
        V_solid = A_gross * L
        V_holes = A_holes * web_thickness
        mass = rho * (V_solid - min(V_holes, V_solid * 0.8))

        # Approximate tip displacement (cantilever, point load)
        # δ = PL³/(3EI)
        P = 10000.0 * 9.81  # 10,000 kgf → N
        E = 200000.0  # MPa (steel)
        if I_section > 0:
            displacement = P * L**3 / (3.0 * E * I_section)
        else:
            displacement = 9999.0

        # Approximate max bending stress
        # σ = M*y/I where M = P*L, y = h_total/2
        M_max = P * L  # N·mm
        y_max = h_total / 2.0
        if I_section > 0:
            stress = M_max * y_max / I_section  # MPa
        else:
            stress = 9999.0

        return TrialResult(
            success=True,
            mass=mass,
            tip_displacement=displacement,
            max_von_mises=stress,
        )

    def close(self) -> None:
        """No-op for stub solver."""
        logger.info("Stub solver closed.")


# ---------------------------------------------------------------------------
# NXOpen implementation using existing Atomizer engine
# ---------------------------------------------------------------------------
class NXOpenSolver:
    """Real NX solver using existing Atomizer optimization engine.

    Uses these Atomizer components:
    - optimization_engine.nx.solver.NXSolver (journal-based solving with run_journal.exe)
      → handles iteration folders, expression import via .exp, and NX solve
    - optimization_engine.extractors.extract_displacement.extract_displacement()
    - optimization_engine.extractors.extract_von_mises_stress.extract_solid_stress()
    - optimization_engine.extractors.extract_mass_from_expression.extract_mass_from_expression()

    Pipeline per trial (HEEDS-style iteration folder pattern):
        1. NXSolver.create_iteration_folder() — copies model files + writes params.exp
        2. NXSolver.run_simulation() — runs solve_simulation.py journal via run_journal.exe
           → The journal imports params.exp, rebuilds geometry, updates FEM, solves, extracts mass
        3. extract_displacement(op2) — max displacement from SOL 101
        4. extract_solid_stress(op2) — max von Mises (auto-detect element type)
        5. extract_mass_from_expression(prt) — reads _temp_mass.txt written by journal

    Files required in model_dir:
    - Beam.prt (part file with expressions)
    - Beam_sim1.sim (simulation file)
    - Expected OP2 output: beam_sim1-solution_1.op2

    Expression names (confirmed from binary introspection):
    - DVs: beam_half_core_thickness, beam_face_thickness, holes_diameter, hole_count
    - Mass output: p173 (body_property147.mass, kg)

    References:
    - M1_Mirror/SAT3_Trajectory_V7/run_optimization.py — FEARunner pattern
    - optimization_engine/nx/solver.py — NXSolver API
    - optimization_engine/nx/solve_simulation.py — Journal internals
    """

    # SIM filename and solution name for this model
    SIM_FILENAME = "Beam_sim1.sim"
    PRT_FILENAME = "Beam.prt"
    SOLUTION_NAME = "Solution 1"
    # Expected OP2: <sim_stem>-<solution_name_lower_underscored>.op2
    # = beam_sim1-solution_1.op2

    def __init__(
        self,
        model_dir: str | Path,
        nx_install_dir: str | Path | None = None,
        timeout: int = 600,
        nastran_version: str = "2412",
    ) -> None:
        """Initialize NXOpen solver using Atomizer engine.

        Args:
            model_dir: Path to directory containing Beam.prt, Beam_sim1.sim, etc.
                       This is the "master model" directory — files are copied per iteration.
            nx_install_dir: Path to NX installation (auto-detected if None).
            timeout: Timeout per trial in seconds (default: 600s = 10 min).
            nastran_version: NX version string (e.g., "2412", "2506", "2512").
        """
        import time as _time  # avoid repeated __import__
        self._time = _time

        self.model_dir = Path(model_dir)
        self.timeout = timeout
        self.nastran_version = nastran_version
        self.nx_install_dir = str(nx_install_dir) if nx_install_dir else None

        if not self.model_dir.exists():
            raise FileNotFoundError(f"Model directory not found: {self.model_dir}")

        # Validate required files
        self.prt_file = self.model_dir / self.PRT_FILENAME
        self.sim_file = self.model_dir / self.SIM_FILENAME

        for f in (self.prt_file, self.sim_file):
            if not f.exists():
                raise FileNotFoundError(f"Required file not found: {f}")

        # Iterations output directory (sibling to model_dir per study convention)
        # Layout: studies/01_doe_landscape/
        #   1_setup/model/   ← model_dir (master)
        #   2_iterations/    ← iteration folders
        #   3_results/       ← final outputs
        self.iterations_dir = self.model_dir.parent.parent / "2_iterations"
        self.iterations_dir.mkdir(parents=True, exist_ok=True)

        # Import Atomizer components at init time (fail-fast on missing engine)
        try:
            from optimization_engine.nx.solver import NXSolver
            from optimization_engine.extractors.extract_displacement import (
                extract_displacement,
            )
            from optimization_engine.extractors.extract_von_mises_stress import (
                extract_solid_stress,
            )
            from optimization_engine.extractors.extract_mass_from_expression import (
                extract_mass_from_expression,
            )

            self._NXSolver = NXSolver
            self._extract_displacement = staticmethod(extract_displacement)
            self._extract_stress = staticmethod(extract_solid_stress)
            self._extract_mass = staticmethod(extract_mass_from_expression)

        except ImportError as e:
            raise ImportError(
                f"Failed to import Atomizer optimization engine: {e}\n"
                f"Ensure {ATOMIZER_REPO_ROOT} is accessible and contains optimization_engine/"
            ) from e

        # Lazy-init solver on first evaluate() call
        self._solver: object | None = None
        self._trial_counter: int = 0

        logger.info(
            "NXOpenSolver initialized — model_dir=%s, timeout=%ds, nastran=%s",
            self.model_dir,
            self.timeout,
            self.nastran_version,
        )

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def evaluate(self, trial_input: TrialInput) -> TrialResult:
        """Full NX evaluation pipeline for one trial.

        Pipeline (mirrors M1_Mirror/SAT3_Trajectory_V7 FEARunner.run_fea):
            1. create_iteration_folder → copies model + writes params.exp
            2. run_simulation → journal updates expressions, rebuilds, solves
            3. extract displacement, stress, mass from results

        Args:
            trial_input: Design variable values.

        Returns:
            TrialResult with extracted outputs or failure info.
        """
        self._trial_counter += 1
        trial_num = self._trial_counter
        t_start = self._time.time()

        logger.info(
            "Trial %d — DVs: core=%.3f mm, face=%.3f mm, hole_d=%.3f mm, n_holes=%d",
            trial_num,
            trial_input.beam_half_core_thickness,
            trial_input.beam_face_thickness,
            trial_input.holes_diameter,
            trial_input.hole_count,
        )

        try:
            # 0. Lazy-init solver
            if self._solver is None:
                self._init_solver()

            expressions = self._build_expression_dict(trial_input)

            # 1. Create iteration folder with fresh model copies + params.exp
            iter_folder = self._solver.create_iteration_folder(
                iterations_base_dir=self.iterations_dir,
                iteration_number=trial_num,
                expression_updates=expressions,
            )
            logger.info("  Iteration folder: %s", iter_folder)

            working_sim = iter_folder / self.SIM_FILENAME
            working_prt = iter_folder / self.PRT_FILENAME

            if not working_sim.exists():
                return TrialResult(
                    success=False,
                    error_message=f"SIM file missing in iteration folder: {working_sim}",
                )

            # 2. Solve — journal handles expression import + geometry rebuild + FEM update + solve
            #    expression_updates are passed as argv to the journal (key=value pairs)
            logger.info("  Solving: %s", working_sim.name)
            solve_result = self._solver.run_simulation(
                sim_file=working_sim,
                working_dir=iter_folder,
                cleanup=False,  # keep OP2/F06 for extraction
                expression_updates=expressions,
                solution_name=self.SOLUTION_NAME,
            )

            if not solve_result["success"]:
                errors = solve_result.get("errors", ["Unknown error"])
                rc = solve_result.get("return_code", "?")
                msg = f"NX solve failed (rc={rc}): {'; '.join(errors)}"
                logger.error("  %s", msg)
                return TrialResult(success=False, error_message=msg)

            # 3. Locate OP2
            op2_file = solve_result.get("op2_file")
            if op2_file is None or not Path(op2_file).exists():
                # Fallback: try the expected naming convention
                op2_file = iter_folder / "beam_sim1-solution_1.op2"
                if not op2_file.exists():
                    return TrialResult(
                        success=False,
                        error_message=f"OP2 not found. Expected: {op2_file}",
                    )
            else:
                op2_file = Path(op2_file)

            logger.info("  OP2: %s", op2_file.name)

            # 4a. Extract displacement
            try:
                disp_result = self._extract_displacement(op2_file)
                tip_displacement = disp_result["max_displacement"]  # mm
            except Exception as e:
                logger.error("  Displacement extraction failed: %s", e)
                return TrialResult(
                    success=False,
                    error_message=f"Displacement extraction failed: {e}",
                )

            # 4b. Extract stress — auto-detect element type (solid or shell)
            #     Pass element_type=None so it checks CTETRA, CHEXA, CPENTA, CPYRAM
            try:
                stress_result = self._extract_stress(
                    op2_file,
                    element_type=None,  # auto-detect from OP2 contents
                    convert_to_mpa=True,  # NX kg-mm-s → kPa, convert to MPa
                )
                max_von_mises = stress_result["max_von_mises"]  # MPa
            except Exception as e:
                # Fallback: try shell elements if solid extraction failed
                logger.warning("  Solid stress extraction failed, trying shell: %s", e)
                try:
                    stress_result = self._extract_stress(
                        op2_file,
                        element_type="cquad4",
                        convert_to_mpa=True,
                    )
                    max_von_mises = stress_result["max_von_mises"]
                except Exception as e2:
                    logger.error("  Stress extraction failed (all types): %s", e2)
                    return TrialResult(
                        success=False,
                        error_message=f"Stress extraction failed: {e}; shell fallback: {e2}",
                    )

            # 4c. Extract mass — reads _temp_mass.txt written by solve_simulation.py journal
            try:
                mass = self._extract_mass(working_prt, expression_name=EXPR_MASS)  # kg
            except FileNotFoundError:
                # _temp_mass.txt not found — journal may not have written it for single-part models
                # Fallback: try reading from _temp_part_properties.json
                logger.warning("  _temp_mass.txt not found, trying _temp_part_properties.json")
                mass = self._extract_mass_fallback(iter_folder)
                if mass is None:
                    return TrialResult(
                        success=False,
                        error_message="Mass extraction failed: _temp_mass.txt not found",
                    )
            except Exception as e:
                logger.error("  Mass extraction failed: %s", e)
                return TrialResult(
                    success=False,
                    error_message=f"Mass extraction failed: {e}",
                )

            elapsed = self._time.time() - t_start
            logger.info(
                "  Trial %d OK (%.1fs) — mass=%.4f kg, disp=%.4f mm, σ_vm=%.2f MPa",
                trial_num,
                elapsed,
                mass,
                tip_displacement,
                max_von_mises,
            )

            return TrialResult(
                success=True,
                mass=mass,
                tip_displacement=tip_displacement,
                max_von_mises=max_von_mises,
            )

        except Exception as e:
            elapsed = self._time.time() - t_start
            logger.error("  Trial %d FAILED (%.1fs): %s", trial_num, elapsed, e)
            return TrialResult(
                success=False,
                error_message=f"Unexpected error in trial {trial_num}: {e}",
            )

    def close(self) -> None:
        """Clean up NX session resources.

        ⚠️ LAC CRITICAL: NEVER kill NX processes directly.
        Uses NXSessionManager for safe lock cleanup only.
        """
        if self._solver is not None:
            sm = getattr(self._solver, "session_manager", None)
            if sm is not None:
                logger.info("Cleaning up NX session locks via session manager")
                try:
                    sm.cleanup_stale_locks()
                except Exception as e:
                    logger.warning("Session lock cleanup warning: %s", e)

        self._solver = None
        logger.info("NXOpenSolver closed.")

    # ------------------------------------------------------------------
    # Private helpers
    # ------------------------------------------------------------------

    def _init_solver(self) -> None:
        """Lazy-initialize NXSolver (matches SAT3_V7 FEARunner.setup pattern)."""
        logger.info("Initializing NXSolver (nastran=%s, timeout=%ds)", self.nastran_version, self.timeout)

        kwargs: dict = {
            "nastran_version": self.nastran_version,
            "timeout": self.timeout,
            "use_journal": True,
            "enable_session_management": True,
            "study_name": "hydrotech_beam_doe",
            "use_iteration_folders": True,
            "master_model_dir": str(self.model_dir),
        }
        if self.nx_install_dir:
            kwargs["nx_install_dir"] = self.nx_install_dir

        self._solver = self._NXSolver(**kwargs)
        logger.info("NXSolver ready")

    def _build_expression_dict(self, trial_input: TrialInput) -> dict[str, float]:
        """Build NX expression name→value dict for the solver.

        These are passed to:
        - create_iteration_folder() → writes params.exp (unit defaulting to mm)
        - run_simulation(expression_updates=...) → passed as argv to solve journal
        """
        return {
            EXPR_HALF_CORE_THICKNESS: trial_input.beam_half_core_thickness,
            EXPR_FACE_THICKNESS: trial_input.beam_face_thickness,
            EXPR_HOLES_DIAMETER: trial_input.holes_diameter,
            EXPR_HOLE_COUNT: float(trial_input.hole_count),  # NX expressions are float
        }

    @staticmethod
    def _extract_mass_fallback(iter_folder: Path) -> float | None:
        """Try to read mass from _temp_part_properties.json (backup path)."""
        import json as _json

        props_file = iter_folder / "_temp_part_properties.json"
        if not props_file.exists():
            return None
        try:
            with open(props_file) as f:
                props = _json.load(f)
            mass = props.get("mass_kg", 0.0)
            if mass > 0:
                logger.info("  Mass from _temp_part_properties.json: %.6f kg", mass)
                return mass
            return None
        except Exception as e:
            logger.warning("  Failed to read %s: %s", props_file, e)
            return None


# ---------------------------------------------------------------------------
# Factory
# ---------------------------------------------------------------------------
def create_solver(
    backend: str = "stub",
    model_dir: str = "",
    nx_install_dir: str | None = None,
    timeout: int = 600,
    nastran_version: str = "2412",
) -> NXStubSolver | NXOpenSolver:
    """Create an NX solver instance.

    Args:
        backend: "stub" for development, "nxopen" for real NX (Windows/dalidou only).
        model_dir: Path to NX model directory (required for nxopen backend).
        nx_install_dir: Path to NX installation (auto-detected if None).
        timeout: Timeout per trial in seconds (default: 600s = 10 min).
        nastran_version: NX version (e.g., "2412", "2506", "2512").

    Returns:
        Solver instance implementing the NXSolverInterface protocol.

    Raises:
        ValueError: If backend is unknown or model_dir missing for nxopen.
    """
    if backend == "stub":
        return NXStubSolver()
    elif backend == "nxopen":
        if not model_dir:
            raise ValueError("model_dir required for nxopen backend")
        return NXOpenSolver(
            model_dir=model_dir,
            nx_install_dir=nx_install_dir,
            timeout=timeout,
            nastran_version=nastran_version,
        )
    else:
        raise ValueError(f"Unknown backend: {backend!r}. Use 'stub' or 'nxopen'.")