pythonocc-step-editor/step_editor/polydata.py

from __future__ import annotations

import math
from pathlib import Path
from typing import Callable, Iterable

from OCC.Core.BRep import BRep_Tool
from OCC.Core.BRepAdaptor import BRepAdaptor_Curve, BRepAdaptor_Surface
from OCC.Core.BRepAlgoAPI import BRepAlgoAPI_Cut, BRepAlgoAPI_Defeaturing, BRepAlgoAPI_Fuse
from OCC.Core.BRepBndLib import brepbndlib
from OCC.Core.BOPAlgo import BOPAlgo_GlueFull
from OCC.Core.BRepBuilderAPI import BRepBuilderAPI_Transform
from OCC.Core.BRepCheck import BRepCheck_Analyzer
from OCC.Core.BRepClass3d import BRepClass3d_SolidClassifier
from OCC.Core.BRepFilletAPI import BRepFilletAPI_MakeChamfer, BRepFilletAPI_MakeFillet
from OCC.Core.BRepGProp import brepgprop
from OCC.Core.BRepMesh import BRepMesh_IncrementalMesh
from OCC.Core.BRepPrimAPI import BRepPrimAPI_MakeCylinder, BRepPrimAPI_MakePrism
from OCC.Core.Bnd import Bnd_Box
from OCC.Core.GeomAbs import (
    GeomAbs_BSplineCurve,
    GeomAbs_BSplineSurface,
    GeomAbs_BezierCurve,
    GeomAbs_BezierSurface,
    GeomAbs_Circle,
    GeomAbs_Cone,
    GeomAbs_Cylinder,
    GeomAbs_Ellipse,
    GeomAbs_Hyperbola,
    GeomAbs_Line,
    GeomAbs_OffsetSurface,
    GeomAbs_OtherCurve,
    GeomAbs_OtherSurface,
    GeomAbs_Parabola,
    GeomAbs_Plane,
    GeomAbs_Sphere,
    GeomAbs_SurfaceOfExtrusion,
    GeomAbs_SurfaceOfRevolution,
    GeomAbs_Torus,
)
from OCC.Core.GProp import GProp_GProps
from OCC.Core.ShapeFix import ShapeFix_Shape
from OCC.Core.ShapeUpgrade import ShapeUpgrade_UnifySameDomain
from OCC.Core.TopAbs import (
    TopAbs_EDGE,
    TopAbs_EXTERNAL,
    TopAbs_FACE,
    TopAbs_FORWARD,
    TopAbs_IN,
    TopAbs_INTERNAL,
    TopAbs_OUT,
    TopAbs_REVERSED,
    TopAbs_SOLID,
)
from OCC.Core.TopExp import TopExp_Explorer, topexp
from OCC.Core.TopLoc import TopLoc_Location
from OCC.Core.TopoDS import TopoDS_Compound, TopoDS_Shape, topods
from OCC.Core.TopTools import TopTools_IndexedDataMapOfShapeListOfShape, TopTools_IndexedMapOfShape
from OCC.Core.gp import gp_Ax1, gp_Ax2, gp_Dir, gp_Pnt, gp_Trsf, gp_Vec
from OCC.Extend.TopologyUtils import TopologyExplorer, discretize_edge

from .constants import CURVE_TYPES, SNAPSHOT_FACE_LOGICAL_IDS_KEY, SURFACE_TYPES
from .geometry_utils import *  # noqa: F403


def _polydata_id_key(values: Iterable[int] | None) -> tuple[int, ...] | None:
    if values is None:
        return None
    return tuple(sorted({int(value) for value in values}))


class PolydataMixin:
    def build_face_polydata(
        self,
        face_ids: Iterable[int] | None = None,
        part_ids: Iterable[int] | None = None,
        deflection: float = 0.8,
    ):
        import vtk

        face_key = _polydata_id_key(face_ids)
        part_key = _polydata_id_key(part_ids)
        cache_key = ("faces", face_key, part_key, float(deflection))
        cached = self._polydata_cache_get("_face_polydata_cache", cache_key)
        if cached is not None:
            return cached

        selected_faces = set(face_key) if face_key is not None else None
        selected_parts = set(part_key) if part_key is not None else None
        self._ensure_mesh(deflection)

        points = vtk.vtkPoints()
        polys = vtk.vtkCellArray()
        face_arr = vtk.vtkIntArray()
        face_arr.SetName("face_id")
        part_arr = vtk.vtkIntArray()
        part_arr.SetName("part_id")
        solid_arr = vtk.vtkIntArray()
        solid_arr.SetName("solid_id")

        for face_id, face in enumerate(self.faces):
            part_id = self.face_part_ids[face_id]
            if selected_faces is not None and face_id not in selected_faces:
                continue
            if selected_parts is not None and part_id not in selected_parts:
                continue

            loc = TopLoc_Location()
            tri = BRep_Tool.Triangulation(topods.Face(face), loc)
            if tri is None:
                continue
            transform = loc.Transformation()
            node_offset = points.GetNumberOfPoints()
            for node_index in range(1, tri.NbNodes() + 1):
                pnt = tri.Node(node_index).Transformed(transform)
                points.InsertNextPoint(pnt.X(), pnt.Y(), pnt.Z())

            reversed_face = face.Orientation() == TopAbs_REVERSED
            for tri_index in range(1, tri.NbTriangles() + 1):
                n1, n2, n3 = tri.Triangle(tri_index).Get()
                if reversed_face:
                    n2, n3 = n3, n2
                vtk_tri = vtk.vtkTriangle()
                vtk_tri.GetPointIds().SetId(0, node_offset + n1 - 1)
                vtk_tri.GetPointIds().SetId(1, node_offset + n2 - 1)
                vtk_tri.GetPointIds().SetId(2, node_offset + n3 - 1)
                polys.InsertNextCell(vtk_tri)
                face_arr.InsertNextValue(face_id)
                part_arr.InsertNextValue(part_id)
                solid_arr.InsertNextValue(self.face_solid_ids[face_id])

        poly = vtk.vtkPolyData()
        poly.SetPoints(points)
        poly.SetPolys(polys)
        poly.GetCellData().AddArray(face_arr)
        poly.GetCellData().AddArray(part_arr)
        poly.GetCellData().AddArray(solid_arr)
        return self._polydata_cache_remember("_face_polydata_cache", cache_key, poly)

    def _ensure_mesh(self, deflection: float) -> None:
        requested = max(float(deflection), 1e-9)
        if self._mesh_deflection is None or requested < self._mesh_deflection * 0.999:
            if requested <= 0.04:
                BRepMesh_IncrementalMesh(
                    self.shape,
                    0.35,
                    False,
                    math.radians(12.0),
                    True,
                )
                for face in self.faces:
                    try:
                        surface_type = BRepAdaptor_Surface(face).GetType()
                        if surface_type in {GeomAbs_Cylinder, GeomAbs_Cone}:
                            BRepMesh_IncrementalMesh(face, requested, False, math.radians(1.2), True)
                        elif surface_type in {GeomAbs_Sphere, GeomAbs_Torus}:
                            BRepMesh_IncrementalMesh(face, max(requested, 0.1), False, math.radians(4.0), True)
                    except Exception:
                        continue
            else:
                angular_degrees = min(24.0, max(12.0, 12.0 * requested / 0.35))
                BRepMesh_IncrementalMesh(self.shape, requested, False, math.radians(angular_degrees), True)
            self._mesh_deflection = requested

    def build_snapshot_polydata(self, snapshot: dict[object, object], deflection: float = 0.8):
        shape = _compound_from_shapes(value for value in snapshot.values() if isinstance(value, TopoDS_Shape))
        BRepMesh_IncrementalMesh(shape, deflection)
        return _shape_faces_polydata(shape)

    def build_edge_polydata(
        self,
        edge_ids: Iterable[int] | None = None,
        part_ids: Iterable[int] | None = None,
        deflection: float = 0.8,
        show_same_domain_internal_edges: bool = False,
    ):
        import vtk

        edge_key = _polydata_id_key(edge_ids)
        part_key = _polydata_id_key(part_ids)
        cache_key = ("edges", edge_key, part_key, float(deflection), bool(show_same_domain_internal_edges))
        cached = self._polydata_cache_get("_edge_polydata_cache", cache_key)
        if cached is not None:
            return cached

        selected_edges = set(edge_key) if edge_key is not None else None
        selected_parts = set(part_key) if part_key is not None else None
        points = vtk.vtkPoints()
        lines = vtk.vtkCellArray()
        edge_arr = vtk.vtkIntArray()
        edge_arr.SetName("edge_id")
        part_arr = vtk.vtkIntArray()
        part_arr.SetName("part_id")
        hidden_edge_ids = (
            set()
            if selected_edges is not None or show_same_domain_internal_edges
            else self._same_domain_internal_edge_ids()
        )

        for edge_id, edge in enumerate(self.edges):
            part_id = self.edge_part_ids[edge_id]
            if selected_edges is not None and edge_id not in selected_edges:
                continue
            if selected_parts is not None and part_id not in selected_parts:
                continue
            if edge_id in hidden_edge_ids:
                continue
            samples = discretize_edge(edge, deflection)
            if len(samples) < 2:
                continue
            polyline = vtk.vtkPolyLine()
            polyline.GetPointIds().SetNumberOfIds(len(samples))
            for i, coords in enumerate(samples):
                point_id = points.InsertNextPoint(float(coords[0]), float(coords[1]), float(coords[2]))
                polyline.GetPointIds().SetId(i, point_id)
            lines.InsertNextCell(polyline)
            edge_arr.InsertNextValue(edge_id)
            part_arr.InsertNextValue(part_id)

        poly = vtk.vtkPolyData()
        poly.SetPoints(points)
        poly.SetLines(lines)
        poly.GetCellData().AddArray(edge_arr)
        poly.GetCellData().AddArray(part_arr)
        return self._polydata_cache_remember("_edge_polydata_cache", cache_key, poly)

    def _polydata_cache_get(self, cache_name: str, key: tuple[object, ...]):
        cache = getattr(self, cache_name, None)
        if not isinstance(cache, dict):
            return None
        return cache.get(key)

    def _polydata_cache_remember(self, cache_name: str, key: tuple[object, ...], polydata):
        cache = getattr(self, cache_name, None)
        if not isinstance(cache, dict):
            return polydata
        limit = max(int(getattr(self, "_polydata_cache_limit", 96)), 1)
        if len(cache) >= limit and key not in cache:
            try:
                cache.pop(next(iter(cache)))
            except StopIteration:
                pass
        cache[key] = polydata
        return polydata

    def _is_same_domain_internal_edge(self, edge_id: int) -> bool:
        return edge_id in self._same_domain_internal_edge_ids()

    def _same_domain_internal_edge_ids(self) -> set[int]:
        if self._same_domain_internal_edge_ids_cache is not None:
            return self._same_domain_internal_edge_ids_cache

        hidden_edge_ids: set[int] = set()
        tolerance = min(max(_shape_diagonal(self.shape) * 1e-7, 1e-6), 1e-3)
        for edge_id in range(len(self.edges)):
            if self._is_topological_same_domain_internal_edge(edge_id, tolerance):
                hidden_edge_ids.add(edge_id)
        hidden_edge_ids.update(self._same_domain_duplicate_edge_ids(tolerance))
        self._same_domain_internal_edge_ids_cache = hidden_edge_ids
        return self._same_domain_internal_edge_ids_cache

    def _is_topological_same_domain_internal_edge(self, edge_id: int, tolerance: float) -> bool:
        if edge_id < 0 or edge_id >= len(self.edges):
            return False
        face_ids = self._edge_adjacent_face_ids(edge_id)
        if len(face_ids) == 2:
            left_id, right_id = face_ids
            if self.face_solid_ids[left_id] == self.face_solid_ids[right_id]:
                left = BRepAdaptor_Surface(self.faces[left_id])
                right = BRepAdaptor_Surface(self.faces[right_id])
                if _surfaces_are_coplanar(left, right, tolerance):
                    return True
                if _surfaces_are_cocylindrical(left, right, tolerance):
                    return True

        return False

    def _same_domain_duplicate_edge_ids(self, tolerance: float) -> set[int]:
        if self._same_domain_duplicate_edge_ids_cache is not None:
            return self._same_domain_duplicate_edge_ids_cache

        duplicate_edge_ids: set[int] = set()
        for bucket_edge_ids in self._edge_duplicate_key_ids(tolerance).values():
            if len(bucket_edge_ids) <= 1:
                continue
            for index, left_edge_id in enumerate(bucket_edge_ids):
                left_face_ids = self._edge_adjacent_face_ids(left_edge_id)
                if not left_face_ids:
                    continue
                for right_edge_id in bucket_edge_ids[index + 1 :]:
                    right_face_ids = self._edge_adjacent_face_ids(right_edge_id)
                    if not right_face_ids:
                        continue
                    if self._edge_face_sets_share_same_domain(left_face_ids, right_face_ids):
                        duplicate_edge_ids.add(left_edge_id)
                        duplicate_edge_ids.add(right_edge_id)

        self._same_domain_duplicate_edge_ids_cache = set(duplicate_edge_ids)
        return self._same_domain_duplicate_edge_ids_cache

    def _edge_duplicate_key_ids(self, tolerance: float) -> dict[tuple[object, ...], list[int]]:
        if self._edge_duplicate_key_ids_cache is not None:
            return self._edge_duplicate_key_ids_cache
        key_tolerance = max(tolerance * 10.0, _shape_diagonal(self.shape) * 1e-7, 1e-6)
        buckets: dict[tuple[object, ...], list[int]] = {}
        for edge_id, edge in enumerate(self.edges):
            key = _edge_duplicate_key(edge, key_tolerance)
            if key is None:
                continue
            buckets.setdefault(key, []).append(edge_id)
        self._edge_duplicate_key_ids_cache = {key: list(value) for key, value in buckets.items() if len(value) > 1}
        return self._edge_duplicate_key_ids_cache

    def _edge_face_sets_share_same_domain(self, left_face_ids: list[int], right_face_ids: list[int]) -> bool:
        tolerance = min(max(_shape_diagonal(self.shape) * 1e-7, 1e-6), 1e-3)
        for left_face_id in left_face_ids:
            if left_face_id < 0 or left_face_id >= len(self.faces):
                continue
            left_solid_id = self.face_solid_ids[left_face_id]
            left_surface = BRepAdaptor_Surface(self.faces[left_face_id])
            for right_face_id in right_face_ids:
                if right_face_id == left_face_id or right_face_id < 0 or right_face_id >= len(self.faces):
                    continue
                if left_solid_id != self.face_solid_ids[right_face_id]:
                    continue
                right_surface = BRepAdaptor_Surface(self.faces[right_face_id])
                if _surfaces_are_coplanar(left_surface, right_surface, tolerance):
                    return True
                if _surfaces_are_cocylindrical(left_surface, right_surface, tolerance):
                    return True
        return False