Atomizer/optimization_engine/hooks/examples.py

"""
NX Open Hooks - Usage Examples
==============================

This file contains practical examples of using the NX Open hooks
for common optimization tasks.

Run examples:
    python -m optimization_engine.hooks.examples

Or import specific examples:
    from optimization_engine.hooks.examples import design_exploration_example
"""

import os
import sys
from pathlib import Path

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

from optimization_engine.hooks.nx_cad import (
    part_manager,
    expression_manager,
    geometry_query,
    feature_manager,
)


# =============================================================================
# Example 1: Basic Expression Query
# =============================================================================

def basic_expression_query(part_path: str):
    """
    Example: Query all expressions from an NX part.

    This is useful for discovering available design parameters
    before setting up an optimization study.
    """
    print("\n" + "=" * 60)
    print("Example 1: Basic Expression Query")
    print("=" * 60)

    result = expression_manager.get_expressions(part_path)

    if not result["success"]:
        print(f"ERROR: {result['error']}")
        return None

    data = result["data"]
    print(f"\nFound {data['count']} expressions:\n")

    # Print in a nice table format
    print(f"{'Name':<25} {'Value':>12} {'Units':<15} {'RHS'}")
    print("-" * 70)

    for name, expr in data["expressions"].items():
        units = expr.get("units") or ""
        rhs = expr.get("rhs", "")
        # Truncate RHS if it's a formula reference
        if len(rhs) > 20:
            rhs = rhs[:17] + "..."
        print(f"{name:<25} {expr['value']:>12.4f} {units:<15} {rhs}")

    return data["expressions"]


# =============================================================================
# Example 2: Mass Properties Extraction
# =============================================================================

def mass_properties_example(part_path: str):
    """
    Example: Extract mass properties from an NX part.

    This is useful for mass optimization objectives.
    """
    print("\n" + "=" * 60)
    print("Example 2: Mass Properties Extraction")
    print("=" * 60)

    result = geometry_query.get_mass_properties(part_path)

    if not result["success"]:
        print(f"ERROR: {result['error']}")
        return None

    data = result["data"]

    print(f"\nMass Properties:")
    print("-" * 40)
    print(f"  Mass:         {data['mass']:.6f} {data['mass_unit']}")
    print(f"  Volume:       {data['volume']:.2f} {data['volume_unit']}")
    print(f"  Surface Area: {data['surface_area']:.2f} {data['area_unit']}")
    print(f"  Material:     {data['material'] or 'Not assigned'}")

    centroid = data["centroid"]
    print(f"\nCentroid (mm):")
    print(f"  X: {centroid['x']:.4f}")
    print(f"  Y: {centroid['y']:.4f}")
    print(f"  Z: {centroid['z']:.4f}")

    if data.get("principal_moments"):
        pm = data["principal_moments"]
        print(f"\nPrincipal Moments of Inertia ({pm['unit']}):")
        print(f"  Ixx: {pm['Ixx']:.4f}")
        print(f"  Iyy: {pm['Iyy']:.4f}")
        print(f"  Izz: {pm['Izz']:.4f}")

    return data


# =============================================================================
# Example 3: Design Parameter Update
# =============================================================================

def design_update_example(part_path: str, dry_run: bool = True):
    """
    Example: Update design parameters in an NX part.

    This demonstrates the workflow for parametric optimization:
    1. Read current values
    2. Compute new values
    3. Update the model

    Args:
        part_path: Path to the NX part
        dry_run: If True, only shows what would be changed (default)
    """
    print("\n" + "=" * 60)
    print("Example 3: Design Parameter Update")
    print("=" * 60)

    # Step 1: Get current expressions
    result = expression_manager.get_expressions(part_path)
    if not result["success"]:
        print(f"ERROR: {result['error']}")
        return None

    expressions = result["data"]["expressions"]

    # Step 2: Find numeric expressions (potential design variables)
    design_vars = {}
    for name, expr in expressions.items():
        # Skip linked expressions (RHS contains another expression name)
        if expr.get("rhs") and not expr["rhs"].replace(".", "").replace("-", "").isdigit():
            continue
        # Only include length/angle expressions
        if expr.get("units") in ["MilliMeter", "Degrees", None]:
            design_vars[name] = expr["value"]

    print(f"\nIdentified {len(design_vars)} potential design variables:")
    for name, value in design_vars.items():
        print(f"  {name}: {value}")

    if dry_run:
        print("\n[DRY RUN] Would update expressions (no changes made)")

        # Example: increase all dimensions by 10%
        new_values = {name: value * 1.1 for name, value in design_vars.items()}

        print("\nProposed changes:")
        for name, new_val in new_values.items():
            old_val = design_vars[name]
            print(f"  {name}: {old_val:.4f} -> {new_val:.4f} (+10%)")

        return new_values

    else:
        # Actually update the model
        new_values = {name: value * 1.1 for name, value in design_vars.items()}

        print("\nUpdating expressions...")
        result = expression_manager.set_expressions(part_path, new_values)

        if result["success"]:
            print(f"SUCCESS: Updated {result['data']['update_count']} expressions")
            if result["data"].get("errors"):
                print(f"Warnings: {result['data']['errors']}")
        else:
            print(f"ERROR: {result['error']}")

        return result


# =============================================================================
# Example 4: Feature Exploration
# =============================================================================

def feature_exploration_example(part_path: str):
    """
    Example: Explore and manipulate features.

    This is useful for topological optimization where features
    can be suppressed/unsuppressed to explore design space.
    """
    print("\n" + "=" * 60)
    print("Example 4: Feature Exploration")
    print("=" * 60)

    result = feature_manager.get_features(part_path)

    if not result["success"]:
        print(f"ERROR: {result['error']}")
        return None

    data = result["data"]

    print(f"\nFound {data['count']} features ({data['suppressed_count']} suppressed):\n")

    print(f"{'Name':<30} {'Type':<20} {'Status'}")
    print("-" * 60)

    for feat in data["features"]:
        status = "SUPPRESSED" if feat["is_suppressed"] else "Active"
        print(f"{feat['name']:<30} {feat['type']:<20} {status}")

    # Group by type
    print("\n\nFeatures by type:")
    print("-" * 40)
    type_counts = {}
    for feat in data["features"]:
        feat_type = feat["type"]
        type_counts[feat_type] = type_counts.get(feat_type, 0) + 1

    for feat_type, count in sorted(type_counts.items(), key=lambda x: -x[1]):
        print(f"  {feat_type}: {count}")

    return data


# =============================================================================
# Example 5: Optimization Objective Evaluation
# =============================================================================

def evaluate_design_point(part_path: str, parameters: dict) -> dict:
    """
    Example: Complete design evaluation workflow.

    This demonstrates how hooks integrate into an optimization loop:
    1. Update parameters
    2. Extract objectives (mass, volume)
    3. Return metrics

    Args:
        part_path: Path to the NX part
        parameters: Dict of parameter_name -> new_value

    Returns:
        Dict with mass_kg, volume_mm3, surface_area_mm2
    """
    print("\n" + "=" * 60)
    print("Example 5: Optimization Objective Evaluation")
    print("=" * 60)

    print(f"\nParameters to set:")
    for name, value in parameters.items():
        print(f"  {name} = {value}")

    # Step 1: Update parameters
    print("\n[1/2] Updating design parameters...")
    result = expression_manager.set_expressions(part_path, parameters)

    if not result["success"]:
        raise RuntimeError(f"Failed to set expressions: {result['error']}")

    print(f"  Updated {result['data']['update_count']} expressions")

    # Step 2: Extract objectives
    print("\n[2/2] Extracting mass properties...")
    result = geometry_query.get_mass_properties(part_path)

    if not result["success"]:
        raise RuntimeError(f"Failed to get mass properties: {result['error']}")

    data = result["data"]

    # Return metrics
    metrics = {
        "mass_kg": data["mass"],
        "volume_mm3": data["volume"],
        "surface_area_mm2": data["surface_area"],
        "material": data.get("material"),
    }

    print(f"\nObjective metrics:")
    print(f"  Mass:         {metrics['mass_kg']:.6f} kg")
    print(f"  Volume:       {metrics['volume_mm3']:.2f} mm^3")
    print(f"  Surface Area: {metrics['surface_area_mm2']:.2f} mm^2")

    return metrics


# =============================================================================
# Example 6: Batch Processing Multiple Parts
# =============================================================================

def batch_mass_extraction(part_paths: list) -> list:
    """
    Example: Extract mass from multiple parts.

    Useful for comparing variants or processing a design library.
    """
    print("\n" + "=" * 60)
    print("Example 6: Batch Processing Multiple Parts")
    print("=" * 60)

    results = []

    for i, part_path in enumerate(part_paths, 1):
        print(f"\n[{i}/{len(part_paths)}] Processing: {Path(part_path).name}")

        result = geometry_query.get_mass_properties(part_path)

        if result["success"]:
            data = result["data"]
            results.append({
                "part": Path(part_path).name,
                "mass_kg": data["mass"],
                "volume_mm3": data["volume"],
                "material": data.get("material"),
                "success": True,
            })
            print(f"  Mass: {data['mass']:.4f} kg, Material: {data.get('material')}")
        else:
            results.append({
                "part": Path(part_path).name,
                "error": result["error"],
                "success": False,
            })
            print(f"  ERROR: {result['error']}")

    # Summary
    print("\n" + "-" * 60)
    print("Summary:")
    successful = [r for r in results if r["success"]]
    print(f"  Processed: {len(successful)}/{len(part_paths)} parts")

    if successful:
        total_mass = sum(r["mass_kg"] for r in successful)
        print(f"  Total mass: {total_mass:.4f} kg")

    return results


# =============================================================================
# Main - Run All Examples
# =============================================================================

def main():
    """Run all examples with a test part."""

    # Default test part
    default_part = project_root / "studies/bracket_stiffness_optimization_V3/1_setup/model/Bracket.prt"

    if len(sys.argv) > 1:
        part_path = sys.argv[1]
    else:
        part_path = str(default_part)

    print("\n" + "=" * 60)
    print("NX OPEN HOOKS - EXAMPLES")
    print("=" * 60)
    print(f"\nUsing part: {Path(part_path).name}")

    if not os.path.exists(part_path):
        print(f"\nERROR: Part file not found: {part_path}")
        print("\nUsage: python -m optimization_engine.hooks.examples [part_path]")
        sys.exit(1)

    # Run examples
    try:
        # Example 1: Query expressions
        basic_expression_query(part_path)

        # Example 2: Get mass properties
        mass_properties_example(part_path)

        # Example 3: Design update (dry run)
        design_update_example(part_path, dry_run=True)

        # Example 4: Feature exploration
        feature_exploration_example(part_path)

        print("\n" + "=" * 60)
        print("ALL EXAMPLES COMPLETED SUCCESSFULLY!")
        print("=" * 60)

    except Exception as e:
        print(f"\nEXAMPLE FAILED: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()