pythonocc-step-editor/step_editor/transforms.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


class TransformMixin:
    def translate_part_plan(self, part_id: int, vector: tuple[float, float, float]) -> dict[str, object]:
        part = self.part_by_id(part_id)
        if part is None:
            raise ValueError(f"Unknown part id {part_id}")
        readiness = _translation_readiness(vector, part.shape)
        return {
            "status": readiness["translate_status"],
            "risk": readiness["translate_risk"],
            "message": readiness["translate_note"],
            "warnings": readiness["translate_warnings"],
            "blockers": readiness["translate_blockers"],
            "target_kind": "part",
            "part_id": part.id,
            "name": part.name,
            "translation_vector": vector,
            "translation_distance": _vector_length(vector),
            "bbox_diagonal": _shape_diagonal(part.shape),
        }

    def translate_solid_plan(self, solid_id: int, vector: tuple[float, float, float]) -> dict[str, object]:
        if solid_id < 0 or solid_id >= len(self.solids):
            raise ValueError(f"Unknown solid id {solid_id}")
        part_id, solid = self.solids[solid_id]
        readiness = _translation_readiness(vector, solid)
        part = self.part_by_id(part_id)
        part_solid_count = len(_explore(part.shape, TopAbs_SOLID)) if part is not None else 0
        warnings = readiness["translate_warnings"]
        risk = readiness["translate_risk"]
        status = readiness["translate_status"]
        if part_solid_count <= 1 and status != "blocked":
            warnings = _join_nonempty(warnings, "当前零件只有一个Solid，平移Solid实际会移动整个零件 shape。")
            if risk == "low":
                risk = "medium"
                status = "caution"
        return {
            "status": status,
            "risk": risk,
            "message": _join_nonempty(readiness["translate_note"], warnings),
            "warnings": warnings,
            "blockers": readiness["translate_blockers"],
            "target_kind": "solid",
            "part_id": part_id,
            "solid_id": solid_id,
            "part_solid_count": part_solid_count,
            "translation_vector": vector,
            "translation_distance": _vector_length(vector),
            "bbox_diagonal": _shape_diagonal(solid),
        }

    def translate_part(self, part_id: int, vector: tuple[float, float, float]) -> str:
        plan = self.translate_part_plan(part_id, vector)
        if plan["status"] == "blocked":
            raise ValueError(str(plan["message"]))
        part = self.part_by_id(part_id)
        if part is None:
            raise ValueError(f"未知零件 ID {part_id}")
        previous_logical_ids = tuple(getattr(self, "face_logical_ids", ()))
        part.shape = _translated_shape_by_vector(part.shape, vector)
        _ensure_valid_shape(part.shape)
        self.refresh_topology()
        self._restore_face_logical_ids_if_count_matches(previous_logical_ids)
        return (
            f"零件已平移: 零件 {part_id}, vector={_format_tuple(vector)}, "
            f"distance={float(plan['translation_distance']):g}, risk={plan['risk']}."
        )

    def translate_solid(self, solid_id: int, vector: tuple[float, float, float]) -> str:
        plan = self.translate_solid_plan(solid_id, vector)
        if plan["status"] == "blocked":
            raise ValueError(str(plan["message"]))
        part_id, solid = self.solids[solid_id]
        part = self.part_by_id(part_id)
        if part is None:
            raise ValueError(f"Unknown part id {part_id}")

        previous_logical_ids = tuple(getattr(self, "face_logical_ids", ()))
        part_solids = _explore(part.shape, TopAbs_SOLID)
        if len(part_solids) <= 1:
            part.shape = _translated_shape_by_vector(part.shape, vector)
        else:
            translated = _translated_shape_by_vector(solid, vector)
            replaced = False
            shapes: list[TopoDS_Shape] = []
            for item in part_solids:
                if not replaced and _same_shape(item, solid):
                    shapes.append(translated)
                    replaced = True
                else:
                    shapes.append(item)
            if not replaced:
                raise RuntimeError(f"Could not locate solid {solid_id} inside part {part_id}.")
            part.shape = _compound_from_shapes(shapes)

        _ensure_valid_shape(part.shape)
        self.refresh_topology()
        self._restore_face_logical_ids_if_count_matches(previous_logical_ids)
        return (
            f"Solid translated: solid {solid_id}, part {part_id}, vector={_format_tuple(vector)}, "
            f"distance={float(plan['translation_distance']):g}, risk={plan['risk']}."
        )

    def rotate_part_plan(self, part_id: int, axis: str, angle_degrees: float) -> dict[str, object]:
        part = self.part_by_id(part_id)
        if part is None:
            raise ValueError(f"Unknown part id {part_id}")
        readiness = _rotation_readiness(axis, angle_degrees)
        return {
            "status": readiness["rotate_status"],
            "risk": readiness["rotate_risk"],
            "message": readiness["rotate_note"],
            "warnings": readiness["rotate_warnings"],
            "blockers": readiness["rotate_blockers"],
            "target_kind": "part",
            "part_id": part.id,
            "name": part.name,
            "rotation_axis": axis.upper(),
            "rotation_angle_degrees": angle_degrees,
            "rotation_center": _shape_center(part.shape),
            "bbox_diagonal": _shape_diagonal(part.shape),
        }

    def rotate_solid_plan(self, solid_id: int, axis: str, angle_degrees: float) -> dict[str, object]:
        if solid_id < 0 or solid_id >= len(self.solids):
            raise ValueError(f"Unknown solid id {solid_id}")
        part_id, solid = self.solids[solid_id]
        readiness = _rotation_readiness(axis, angle_degrees)
        part = self.part_by_id(part_id)
        part_solid_count = len(_explore(part.shape, TopAbs_SOLID)) if part is not None else 0
        warnings = readiness["rotate_warnings"]
        risk = readiness["rotate_risk"]
        status = readiness["rotate_status"]
        if part_solid_count <= 1 and status != "blocked":
            warnings = _join_nonempty(warnings, "当前零件只有一个Solid，旋转Solid实际会旋转整个零件 shape。")
            if risk == "low":
                risk = "medium"
                status = "caution"
        return {
            "status": status,
            "risk": risk,
            "message": _join_nonempty(readiness["rotate_note"], warnings),
            "warnings": warnings,
            "blockers": readiness["rotate_blockers"],
            "target_kind": "solid",
            "part_id": part_id,
            "solid_id": solid_id,
            "part_solid_count": part_solid_count,
            "rotation_axis": axis.upper(),
            "rotation_angle_degrees": angle_degrees,
            "rotation_center": _shape_center(solid),
            "bbox_diagonal": _shape_diagonal(solid),
        }

    def rotate_part(self, part_id: int, axis: str, angle_degrees: float) -> str:
        plan = self.rotate_part_plan(part_id, axis, angle_degrees)
        if plan["status"] == "blocked":
            raise ValueError(str(plan["message"]))
        part = self.part_by_id(part_id)
        if part is None:
            raise ValueError(f"未知零件 ID {part_id}")
        part.shape = _rotated_shape(part.shape, str(plan["rotation_axis"]), float(plan["rotation_angle_degrees"]), plan["rotation_center"])
        _ensure_valid_shape(part.shape)
        self.refresh_topology()
        return (
            f"零件已旋转: 零件 {part_id}, axis={plan['rotation_axis']}, "
            f"angle={float(plan['rotation_angle_degrees']):g}, center={_format_tuple(plan['rotation_center'])}, "
            f"risk={plan['risk']}."
        )

    def rotate_solid(self, solid_id: int, axis: str, angle_degrees: float) -> str:
        plan = self.rotate_solid_plan(solid_id, axis, angle_degrees)
        if plan["status"] == "blocked":
            raise ValueError(str(plan["message"]))
        part_id, solid = self.solids[solid_id]
        part = self.part_by_id(part_id)
        if part is None:
            raise ValueError(f"Unknown part id {part_id}")

        part_solids = _explore(part.shape, TopAbs_SOLID)
        if len(part_solids) <= 1:
            part.shape = _rotated_shape(part.shape, str(plan["rotation_axis"]), float(plan["rotation_angle_degrees"]), plan["rotation_center"])
        else:
            rotated = _rotated_shape(solid, str(plan["rotation_axis"]), float(plan["rotation_angle_degrees"]), plan["rotation_center"])
            replaced = False
            shapes: list[TopoDS_Shape] = []
            for item in part_solids:
                if not replaced and _same_shape(item, solid):
                    shapes.append(rotated)
                    replaced = True
                else:
                    shapes.append(item)
            if not replaced:
                raise RuntimeError(f"Could not locate solid {solid_id} inside part {part_id}.")
            part.shape = _compound_from_shapes(shapes)

        _ensure_valid_shape(part.shape)
        self.refresh_topology()
        return (
            f"Solid rotated: solid {solid_id}, part {part_id}, axis={plan['rotation_axis']}, "
            f"angle={float(plan['rotation_angle_degrees']):g}, center={_format_tuple(plan['rotation_center'])}, "
            f"risk={plan['risk']}."
        )

    def scale_part_plan(self, part_id: int, target_diagonal: float) -> dict[str, object]:
        part = self.part_by_id(part_id)
        if part is None:
            raise ValueError(f"Unknown part id {part_id}")
        return self._scale_shape_plan(
            target_kind="part",
            part_id=part.id,
            solid_id=-1,
            shape=part.shape,
            target_diagonal=target_diagonal,
            part_solid_count=len(_explore(part.shape, TopAbs_SOLID)),
        )

    def scale_solid_plan(self, solid_id: int, target_diagonal: float) -> dict[str, object]:
        if solid_id < 0 or solid_id >= len(self.solids):
            raise ValueError(f"Unknown solid id {solid_id}")
        part_id, solid = self.solids[solid_id]
        part = self.part_by_id(part_id)
        part_solid_count = len(_explore(part.shape, TopAbs_SOLID)) if part is not None else 0
        plan = self._scale_shape_plan(
            target_kind="solid",
            part_id=part_id,
            solid_id=solid_id,
            shape=solid,
            target_diagonal=target_diagonal,
            part_solid_count=part_solid_count,
        )
        if part_solid_count <= 1 and plan["status"] != "blocked":
            plan["warnings"] = _join_nonempty(
                plan.get("warnings", ""),
                "当前零件只有一个Solid，缩放Solid实际会缩放整个零件 shape。",
            )
            if plan["risk"] == "low":
                plan["risk"] = "medium"
                plan["status"] = "caution"
            plan["message"] = _join_nonempty(plan.get("message", ""), plan.get("warnings", ""))
        return plan

    def scale_part_axis_plan(self, part_id: int, axis: str, target_size: float) -> dict[str, object]:
        part = self.part_by_id(part_id)
        if part is None:
            raise ValueError(f"Unknown part id {part_id}")
        return self._axis_scale_shape_plan(
            target_kind="part",
            part_id=part.id,
            solid_id=-1,
            shape=part.shape,
            axis=axis,
            target_size=target_size,
            part_solid_count=len(_explore(part.shape, TopAbs_SOLID)),
        )

    def scale_solid_axis_plan(self, solid_id: int, axis: str, target_size: float) -> dict[str, object]:
        if solid_id < 0 or solid_id >= len(self.solids):
            raise ValueError(f"Unknown solid id {solid_id}")
        part_id, solid = self.solids[solid_id]
        part = self.part_by_id(part_id)
        part_solid_count = len(_explore(part.shape, TopAbs_SOLID)) if part is not None else 0
        plan = self._axis_scale_shape_plan(
            target_kind="solid",
            part_id=part_id,
            solid_id=solid_id,
            shape=solid,
            axis=axis,
            target_size=target_size,
            part_solid_count=part_solid_count,
        )
        if part_solid_count <= 1 and plan["status"] != "blocked":
            plan["warnings"] = _join_nonempty(
                plan.get("warnings", ""),
                "当前零件只有一个Solid，按轴缩放Solid实际会缩放整个零件shape。",
            )
            if plan["risk"] == "low":
                plan["risk"] = "medium"
                plan["status"] = "caution"
            plan["message"] = _join_nonempty(plan.get("message", ""), plan.get("warnings", ""))
        return plan

    def _scale_shape_plan(
        self,
        *,
        target_kind: str,
        part_id: int,
        solid_id: int,
        shape: TopoDS_Shape,
        target_diagonal: float,
        part_solid_count: int,
    ) -> dict[str, object]:
        blockers: list[str] = []
        warnings: list[str] = []
        try:
            target_diagonal = float(target_diagonal)
        except (TypeError, ValueError):
            target_diagonal = 0.0
            blockers.append("目标包围盒对角线必须是数字。")
        current_diagonal = _shape_diagonal(shape)
        center = _shape_center(shape)
        if current_diagonal <= 1e-9:
            blockers.append("当前对象包围盒对角线无效，不能缩放。")
        if target_diagonal <= 1e-9:
            blockers.append("目标包围盒对角线必须大于 0。")
        scale = target_diagonal / max(current_diagonal, 1e-9)
        delta_ratio = abs(target_diagonal - current_diagonal) / max(current_diagonal, 1e-9)
        risk = "low"
        status = "ready"
        if abs(target_diagonal - current_diagonal) <= max(current_diagonal * 1e-6, 1e-6):
            blockers.append("目标包围盒对角线与当前值几乎相同，不需要缩放。")
        elif delta_ratio > 1.0:
            risk = "high"
            status = "caution"
            warnings.append("目标尺寸变化超过当前包围盒对角线的 100%，请确认单位和目标。")
        elif delta_ratio > 0.35:
            risk = "medium"
            status = "caution"
            warnings.append("目标尺寸变化超过当前包围盒对角线的 35%，请确认缩放结果。")
        if blockers:
            risk = "blocked"
            status = "blocked"
        return {
            "status": status,
            "risk": risk,
            "message": _join_nonempty(" ".join(blockers), " ".join(warnings)) or "可以按目标包围盒对角线等比缩放当前对象。",
            "warnings": "；".join(warnings),
            "blockers": "；".join(blockers),
            "target_kind": target_kind,
            "part_id": part_id,
            "solid_id": solid_id,
            "part_solid_count": part_solid_count,
            "current_bbox_diagonal": current_diagonal,
            "target_bbox_diagonal": target_diagonal,
            "scale_factor": scale,
            "scale_delta_ratio": delta_ratio,
            "scale_center": center,
            "resize_strategy": "uniform-scale-selected-shape-to-bbox-diagonal",
        }

    def _axis_scale_shape_plan(
        self,
        *,
        target_kind: str,
        part_id: int,
        solid_id: int,
        shape: TopoDS_Shape,
        axis: str,
        target_size: float,
        part_solid_count: int,
    ) -> dict[str, object]:
        axis = str(axis or "X").strip().upper()
        axis_index = {"X": 0, "Y": 1, "Z": 2}.get(axis)
        blockers: list[str] = []
        warnings: list[str] = [
            "按单轴尺寸缩放是 B-Rep 仿射编辑，可能把解析几何转换成 B-spline 或影响同一对象上的其它特征。"
        ]
        try:
            target_size = float(target_size)
        except (TypeError, ValueError):
            target_size = 0.0
            blockers.append("目标轴向尺寸必须是数字。")
        bounds = _shape_bounds_info(shape)
        bbox_size = tuple(bounds.get("bbox_size", (0.0, 0.0, 0.0)))
        current_size = float(bbox_size[axis_index]) if axis_index is not None else 0.0
        center = _shape_center(shape)
        if axis_index is None:
            blockers.append("缩放轴必须是 X、Y 或 Z。")
        if current_size <= 1e-9:
            blockers.append("当前对象在该轴向的包围盒尺寸无效，不能按轴缩放。")
        if target_size <= 1e-9:
            blockers.append("目标轴向尺寸必须大于 0。")
        scale = target_size / max(current_size, 1e-9)
        delta_ratio = abs(target_size - current_size) / max(current_size, 1e-9)
        risk = "medium"
        status = "caution"
        if abs(target_size - current_size) <= max(current_size * 1e-6, 1e-6):
            blockers.append("目标轴向尺寸与当前值几乎相同，不需要缩放。")
        elif delta_ratio > 1.0:
            risk = "high"
            warnings.append("目标轴向尺寸变化超过当前尺寸的 100%，请确认单位和目标。")
        elif delta_ratio <= 0.2:
            risk = "medium"
        if blockers:
            risk = "blocked"
            status = "blocked"
        return {
            "status": status,
            "risk": risk,
            "message": _join_nonempty(" ".join(blockers), " ".join(warnings)) or "可以按目标轴向尺寸缩放当前对象。",
            "warnings": "；".join(warnings),
            "blockers": "；".join(blockers),
            "target_kind": target_kind,
            "part_id": part_id,
            "solid_id": solid_id,
            "part_solid_count": part_solid_count,
            "scale_axis": axis,
            "current_axis_size": current_size,
            "target_axis_size": target_size,
            "axis_scale_factor": scale,
            "axis_scale_delta_ratio": delta_ratio,
            "scale_center": center,
            "bbox_size": bbox_size,
            "resize_strategy": "axis-scale-selected-shape-to-bbox-size",
        }

    def scale_part(self, part_id: int, target_diagonal: float) -> str:
        plan = self.scale_part_plan(part_id, target_diagonal)
        if plan["status"] == "blocked":
            raise ValueError(str(plan["message"]))
        part = self.part_by_id(part_id)
        if part is None:
            raise ValueError(f"Unknown part id {part_id}")
        part.shape = _scaled_shape(part.shape, float(plan["scale_factor"]), plan["scale_center"])
        _ensure_valid_shape(part.shape)
        self.refresh_topology()
        return (
            f"Part scaled: part {part_id}, "
            f"bbox_diagonal={float(plan['current_bbox_diagonal']):g}->{float(plan['target_bbox_diagonal']):g}, "
            f"scale={float(plan['scale_factor']):g}, risk={plan['risk']}."
        )

    def scale_part_axis(self, part_id: int, axis: str, target_size: float) -> str:
        plan = self.scale_part_axis_plan(part_id, axis, target_size)
        if plan["status"] == "blocked":
            raise ValueError(str(plan["message"]))
        part = self.part_by_id(part_id)
        if part is None:
            raise ValueError(f"Unknown part id {part_id}")
        part.shape = _axis_scaled_shape(part.shape, str(plan["scale_axis"]), float(plan["axis_scale_factor"]), plan["scale_center"])
        _ensure_valid_shape(part.shape)
        self.refresh_topology()
        return (
            f"Part axis-scaled: part {part_id}, axis={plan['scale_axis']}, "
            f"size={float(plan['current_axis_size']):g}->{float(plan['target_axis_size']):g}, "
            f"scale={float(plan['axis_scale_factor']):g}, risk={plan['risk']}."
        )

    def scale_solid(self, solid_id: int, target_diagonal: float) -> str:
        plan = self.scale_solid_plan(solid_id, target_diagonal)
        if plan["status"] == "blocked":
            raise ValueError(str(plan["message"]))
        part_id, solid = self.solids[solid_id]
        part = self.part_by_id(part_id)
        if part is None:
            raise ValueError(f"Unknown part id {part_id}")
        part_solids = _explore(part.shape, TopAbs_SOLID)
        if len(part_solids) <= 1:
            part.shape = _scaled_shape(part.shape, float(plan["scale_factor"]), plan["scale_center"])
        else:
            scaled = _scaled_shape(solid, float(plan["scale_factor"]), plan["scale_center"])
            replaced = False
            shapes: list[TopoDS_Shape] = []
            for item in part_solids:
                if not replaced and _same_shape(item, solid):
                    shapes.append(scaled)
                    replaced = True
                else:
                    shapes.append(item)
            if not replaced:
                raise RuntimeError(f"Could not locate solid {solid_id} inside part {part_id}.")
            part.shape = _compound_from_shapes(shapes)
        _ensure_valid_shape(part.shape)
        self.refresh_topology()
        return (
            f"Solid scaled: solid {solid_id}, part {part_id}, "
            f"bbox_diagonal={float(plan['current_bbox_diagonal']):g}->{float(plan['target_bbox_diagonal']):g}, "
            f"scale={float(plan['scale_factor']):g}, risk={plan['risk']}."
        )

    def scale_solid_axis(self, solid_id: int, axis: str, target_size: float) -> str:
        plan = self.scale_solid_axis_plan(solid_id, axis, target_size)
        if plan["status"] == "blocked":
            raise ValueError(str(plan["message"]))
        part_id, solid = self.solids[solid_id]
        part = self.part_by_id(part_id)
        if part is None:
            raise ValueError(f"Unknown part id {part_id}")
        part_solids = _explore(part.shape, TopAbs_SOLID)
        if len(part_solids) <= 1:
            part.shape = _axis_scaled_shape(part.shape, str(plan["scale_axis"]), float(plan["axis_scale_factor"]), plan["scale_center"])
        else:
            scaled = _axis_scaled_shape(solid, str(plan["scale_axis"]), float(plan["axis_scale_factor"]), plan["scale_center"])
            replaced = False
            shapes: list[TopoDS_Shape] = []
            for item in part_solids:
                if not replaced and _same_shape(item, solid):
                    shapes.append(scaled)
                    replaced = True
                else:
                    shapes.append(item)
            if not replaced:
                raise RuntimeError(f"Could not locate solid {solid_id} inside part {part_id}.")
            part.shape = _compound_from_shapes(shapes)
        _ensure_valid_shape(part.shape)
        self.refresh_topology()
        return (
            f"Solid axis-scaled: solid {solid_id}, part {part_id}, axis={plan['scale_axis']}, "
            f"size={float(plan['current_axis_size']):g}->{float(plan['target_axis_size']):g}, "
            f"scale={float(plan['axis_scale_factor']):g}, risk={plan['risk']}."
        )