Atomizer/tools/adaptive-isogrid/src/nx/import_profile.py

"""
NXOpen script — Import rib profile into NX as a sketch.

Reads `rib_profile_<sandbox_id>.json` (output from Python Brain) and creates
an NX sketch on the sandbox plane containing:
  - Outer boundary polyline
  - All pocket cutout polylines

The sketch is placed in the idealized part. Antoine extrudes manually the first
time; subsequent iterations only update the sketch and the extrude regenerates.

Usage (NX Journal):
    File > Execute > NX Journal > import_profile.py

Expects rib_profile JSON files in the same `adaptive_isogrid_data/` folder
created by extract_sandbox.py (next to the idealized part).

Author: Atomizer / Adaptive Isogrid
Created: 2026-02-16
"""

from __future__ import annotations

import json
import math
import os
from typing import Any, Dict, List, Tuple

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------

SKETCH_NAME_PREFIX = "ISOGRID_RIB_"  # e.g., ISOGRID_RIB_sandbox_1


# ---------------------------------------------------------------------------
# Geometry helpers
# ---------------------------------------------------------------------------

def unproject_to_3d(
    points2d: List[List[float]],
    transform: Dict[str, List[float]],
) -> List[Tuple[float, float, float]]:
    """Convert 2D profile points back to 3D using the extraction transform."""
    ox, oy, oz = transform["origin"]
    xx, xy, xz = transform["x_axis"]
    yx, yy, yz = transform["y_axis"]

    pts3d = []
    for x, y in points2d:
        px = ox + x * xx + y * yx
        py = oy + x * xy + y * yy
        pz = oz + x * xz + y * yz
        pts3d.append((px, py, pz))
    return pts3d


# ---------------------------------------------------------------------------
# NX sketch creation
# ---------------------------------------------------------------------------

def _find_or_create_sketch(
    part: Any,
    sketch_name: str,
    transform: Dict[str, List[float]],
    lister: Any,
) -> Any:
    """
    Find existing sketch by name, or create a new one on the sandbox plane.
    If found, delete all existing geometry in it (for update).
    Returns the Sketch object.
    """
    import NXOpen

    # Try to find existing sketch by name
    existing_sketch = None
    try:
        for feat in part.Features:
            fname = ""
            try:
                fname = feat.Name
            except Exception:
                continue
            if fname == sketch_name:
                # Get the sketch from the feature
                try:
                    existing_sketch = feat.GetEntities()[0]
                    lister.WriteLine(f"[import] Found existing sketch: {sketch_name}")
                except Exception:
                    pass
                break
    except Exception:
        pass

    if existing_sketch is not None:
        # Clear existing geometry for update
        try:
            existing_sketch.Activate(NXOpen.Sketch.ViewReorient.DoNotOrientView)
            all_geom = existing_sketch.GetAllGeometry()
            if all_geom:
                existing_sketch.DeleteObjects(list(all_geom))
                lister.WriteLine(f"[import] Cleared {len(all_geom)} objects from existing sketch")
            existing_sketch.Deactivate(
                NXOpen.Sketch.ViewReorient.DoNotOrientView,
                NXOpen.Sketch.UpdateLevel.Model,
            )
        except Exception as exc:
            lister.WriteLine(f"[import] Warning clearing sketch: {exc}")
        return existing_sketch

    # Create new sketch on the sandbox plane
    lister.WriteLine(f"[import] Creating new sketch: {sketch_name}")

    origin = transform["origin"]
    normal = transform["normal"]
    x_axis = transform["x_axis"]

    origin_pt = NXOpen.Point3d(origin[0], origin[1], origin[2])
    normal_vec = NXOpen.Vector3d(normal[0], normal[1], normal[2])
    x_vec = NXOpen.Vector3d(x_axis[0], x_axis[1], x_axis[2])
    y_axis = transform["y_axis"]
    y_vec = NXOpen.Vector3d(y_axis[0], y_axis[1], y_axis[2])

    # Build Matrix3x3 from x_axis, y_axis, normal (row-major: Xx,Xy,Xz, Yx,Yy,Yz, Zx,Zy,Zz)
    mtx = NXOpen.Matrix3x3()
    mtx.Xx = x_axis[0]; mtx.Xy = x_axis[1]; mtx.Xz = x_axis[2]
    mtx.Yx = y_axis[0]; mtx.Yy = y_axis[1]; mtx.Yz = y_axis[2]
    mtx.Zx = normal[0]; mtx.Zy = normal[1]; mtx.Zz = normal[2]

    # Create fixed datum plane
    datum_plane = part.Datums.CreateFixedDatumPlane(origin_pt, mtx)
    lister.WriteLine(f"[import] Created datum plane: {datum_plane}")

    # Create sketch-in-place builder
    sketch_builder = part.Sketches.CreateSketchInPlaceBuilder2(NXOpen.Sketch.Null)

    # Set plane — PlaneReference is a property setter
    sketch_builder.PlaneReference = datum_plane

    # Set sketch origin
    origin_point = part.Points.CreatePoint(origin_pt)
    sketch_builder.SketchOrigin = origin_point

    # Set axis reference
    axis_dir = part.Directions.CreateDirection(origin_pt, x_vec, NXOpen.SmartObject.UpdateOption.WithinModeling)
    sketch_builder.AxisReference = axis_dir

    # Commit to create the sketch
    sketch_feature = sketch_builder.CommitFeature()
    sketch_builder.Destroy()

    # Get the sketch object from the feature
    sketch = None
    try:
        entities = sketch_feature.GetEntities()
        for e in entities:
            if isinstance(e, NXOpen.Sketch):
                sketch = e
                break
        if sketch is None and entities:
            sketch = entities[0]
    except Exception:
        # Try alternate: find sketch by feature
        try:
            sketch = sketch_feature.Sketch
        except Exception:
            pass

    if sketch is None:
        raise RuntimeError("Could not get Sketch object from feature")

    # Rename the feature
    try:
        sketch_feature.Name = sketch_name
    except Exception:
        pass

    lister.WriteLine(f"[import] Created sketch: {sketch_name}")
    return sketch


def _draw_polyline_in_sketch(
    part: Any,
    sketch: Any,
    points_3d: List[Tuple[float, float, float]],
    lister: Any,
    close: bool = True,
) -> int:
    """
    Draw a closed polyline in the sketch using individual line segments.
    Returns number of lines created.
    """
    import NXOpen

    if len(points_3d) < 2:
        return 0

    lines_created = 0
    n = len(points_3d)

    # If last point == first point, don't double-close
    if close and len(points_3d) >= 3:
        d = math.sqrt(sum((a - b) ** 2 for a, b in zip(points_3d[0], points_3d[-1])))
        if d < 0.001:
            n = len(points_3d) - 1  # skip duplicate closing point

    segments = n if close else (n - 1)

    for i in range(segments):
        p1 = points_3d[i]
        p2 = points_3d[(i + 1) % n]

        try:
            start_pt = NXOpen.Point3d(p1[0], p1[1], p1[2])
            end_pt = NXOpen.Point3d(p2[0], p2[1], p2[2])

            line_builder = part.Sketches.CreateLineBuilder()
            line_builder.SetStartPoint(start_pt)
            line_builder.SetEndPoint(end_pt)
            line = line_builder.Commit()
            line_builder.Destroy()

            lines_created += 1
        except Exception:
            # Fallback: try creating a curve and adding to sketch
            try:
                start_obj = part.Points.CreatePoint(NXOpen.Point3d(p1[0], p1[1], p1[2]))
                end_obj = part.Points.CreatePoint(NXOpen.Point3d(p2[0], p2[1], p2[2]))
                line = part.Curves.CreateLine(start_obj, end_obj)
                sketch.AddGeometry(line, NXOpen.Sketch.InferConstraintsOption.DoNotInferConstraints)
                lines_created += 1
            except Exception as exc2:
                if lines_created == 0:
                    lister.WriteLine(f"[import] Line creation failed: {exc2}")

    return lines_created


def _draw_circle_in_sketch(
    part: Any,
    sketch: Any,
    center_3d: Tuple[float, float, float],
    radius: float,
    normal: List[float],
    lister: Any,
) -> bool:
    """Draw a circle in the sketch."""
    import NXOpen

    try:
        circle_builder = part.Sketches.CreateCircleBuilder()
        center_pt = NXOpen.Point3d(center_3d[0], center_3d[1], center_3d[2])
        circle_builder.SetCenterPoint(center_pt)
        # Size point = center + radius along x
        size_pt = NXOpen.Point3d(
            center_3d[0] + radius,
            center_3d[1],
            center_3d[2],
        )
        circle_builder.SetSizePoint(size_pt)
        circle_builder.Commit()
        circle_builder.Destroy()
        return True
    except Exception as exc:
        lister.WriteLine(f"[import] Circle creation failed: {exc}")
        return False


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    import NXOpen

    session = NXOpen.Session.GetSession()
    lister = session.ListingWindow
    lister.Open()

    lister.WriteLine("=" * 60)
    lister.WriteLine("  Adaptive Isogrid — Rib Profile Import")
    lister.WriteLine("=" * 60)

    # Navigate to idealized part
    work_part = session.Parts.Work
    part_name = work_part.Name if hasattr(work_part, "Name") else ""
    lister.WriteLine(f"[import] Work part: {part_name}")

    # If not in idealized part, find it
    if not part_name.endswith("_i"):
        for part in session.Parts:
            pname = part.Name if hasattr(part, "Name") else ""
            if pname.endswith("_i"):
                session.Parts.SetWork(part)
                work_part = part
                lister.WriteLine(f"[import] Switched to idealized part: {pname}")
                break

    # Find data directory
    try:
        part_dir = os.path.dirname(work_part.FullPath)
    except Exception:
        part_dir = os.getcwd()

    data_dir = os.path.join(part_dir, "adaptive_isogrid_data")

    if not os.path.isdir(data_dir):
        lister.WriteLine(f"[import] ERROR: Data directory not found: {data_dir}")
        return

    # Find all rib profile + geometry JSON pairs
    profile_files = sorted([
        f for f in os.listdir(data_dir)
        if f.startswith("rib_profile_") and f.endswith(".json")
    ])

    if not profile_files:
        # Also check for sandbox1_rib_profile.json pattern
        profile_files = sorted([
            f for f in os.listdir(data_dir)
            if "rib_profile" in f and f.endswith(".json")
        ])

    if not profile_files:
        lister.WriteLine(f"[import] ERROR: No rib_profile*.json found in {data_dir}")
        lister.WriteLine(f"[import] Files present: {os.listdir(data_dir)}")
        return

    lister.WriteLine(f"[import] Found {len(profile_files)} profile(s) to import")

    for profile_file in profile_files:
        profile_path = os.path.join(data_dir, profile_file)

        # Determine sandbox_id from filename
        # Expected: rib_profile_sandbox_1.json or sandbox1_rib_profile.json
        sandbox_id = profile_file.replace("rib_profile_", "").replace(".json", "")
        if not sandbox_id.startswith("sandbox"):
            sandbox_id = "sandbox_1"  # fallback

        # Load corresponding geometry JSON for the transform
        geom_path = os.path.join(data_dir, f"geometry_{sandbox_id}.json")
        if not os.path.exists(geom_path):
            # Try alternate names
            candidates = [f for f in os.listdir(data_dir) if f.startswith("geometry_") and f.endswith(".json")]
            if candidates:
                geom_path = os.path.join(data_dir, candidates[0])
            else:
                lister.WriteLine(f"[import] ERROR: No geometry JSON found for transform data")
                continue

        lister.WriteLine(f"\n--- Importing {profile_file} ---")
        lister.WriteLine(f"[import] Geometry (transform): {geom_path}")

        try:
            with open(profile_path, "r") as f:
                profile = json.load(f)
            with open(geom_path, "r") as f:
                geometry = json.load(f)
        except Exception as exc:
            lister.WriteLine(f"[import] ERROR reading JSON: {exc}")
            continue

        transform = geometry["transform"]
        sketch_name = SKETCH_NAME_PREFIX + sandbox_id

        # Find or create sketch
        try:
            sketch = _find_or_create_sketch(work_part, sketch_name, transform, lister)
        except Exception as exc:
            lister.WriteLine(f"[import] ERROR creating sketch: {exc}")
            import traceback
            lister.WriteLine(traceback.format_exc())
            continue

        # Activate sketch for drawing
        try:
            sketch.Activate(NXOpen.Sketch.ViewReorient.DoNotOrientView)
        except Exception as exc:
            lister.WriteLine(f"[import] ERROR activating sketch: {exc}")
            continue

        total_lines = 0

        # Draw outer boundary
        outer_2d = profile.get("outer_boundary", [])
        if outer_2d:
            outer_3d = unproject_to_3d(outer_2d, transform)
            n = _draw_polyline_in_sketch(work_part, sketch, outer_3d, lister, close=True)
            total_lines += n
            lister.WriteLine(f"[import] Outer boundary: {n} lines")

        # Draw pocket cutouts
        pockets = profile.get("pockets", [])
        lister.WriteLine(f"[import] Drawing {len(pockets)} pockets...")

        for pi, pocket_pts in enumerate(pockets):
            if len(pocket_pts) < 3:
                continue
            pocket_3d = unproject_to_3d(pocket_pts, transform)
            n = _draw_polyline_in_sketch(work_part, sketch, pocket_3d, lister, close=True)
            total_lines += n

            # Progress every 50 pockets
            if (pi + 1) % 50 == 0:
                lister.WriteLine(f"[import]   ... {pi + 1}/{len(pockets)} pockets drawn")

        # Deactivate sketch
        try:
            sketch.Deactivate(
                NXOpen.Sketch.ViewReorient.DoNotOrientView,
                NXOpen.Sketch.UpdateLevel.Model,
            )
        except Exception as exc:
            lister.WriteLine(f"[import] Warning deactivating: {exc}")

        lister.WriteLine(f"[import] Done: {total_lines} total line segments in sketch '{sketch_name}'")
        lister.WriteLine(f"[import]   Outer boundary: {len(outer_2d)} pts")
        lister.WriteLine(f"[import]   Pockets: {len(pockets)}")

    lister.WriteLine("\n" + "=" * 60)
    lister.WriteLine("  Import complete — extrude the sketch to rib thickness")
    lister.WriteLine("=" * 60)


main()