pythonocc-step-editor/scripts/verify_boss_resize.py

from __future__ import annotations

import argparse
import math
from pathlib import Path
import sys
import tempfile

from OCC.Core.BRepAlgoAPI import BRepAlgoAPI_Fuse
from OCC.Core.BRepPrimAPI import BRepPrimAPI_MakeBox, BRepPrimAPI_MakeCylinder
from OCC.Core.gp import gp_Ax2, gp_Dir, gp_Pnt

PROJECT_ROOT = Path(__file__).resolve().parent.parent
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from step_editor.geometry_utils import _finalize_boolean_result
from step_editor.model import StepModel
from step_editor.step_io import _write_step


def _write_boss_model(path: Path) -> None:
    base = BRepPrimAPI_MakeBox(30.0, 20.0, 6.0).Shape()
    axis = gp_Ax2(gp_Pnt(15.0, 10.0, 6.0), gp_Dir(0.0, 0.0, 1.0))
    boss = BRepPrimAPI_MakeCylinder(axis, 3.0, 5.0).Shape()
    fuse = BRepAlgoAPI_Fuse(base, boss)
    shape = _finalize_boolean_result(fuse, "verify boss model fuse", use_glue=False)
    _write_step(shape, path)


def _distance(left: tuple[float, float, float], right: tuple[float, float, float]) -> float:
    return math.sqrt(
        (left[0] - right[0]) ** 2
        + (left[1] - right[1]) ** 2
        + (left[2] - right[2]) ** 2
    )


def _boss_face_ids(model: StepModel) -> list[int]:
    face_ids: list[int] = []
    for face_id in range(len(model.faces)):
        info = model.face_info(face_id)
        if info.get("surface") != "cylinder":
            continue
        if str(info.get("feature_guess", "")) != "boss/outer-round candidate":
            continue
        span = float(info.get("angular_span") or 0.0)
        if span < math.tau * 0.92:
            continue
        face_ids.append(face_id)
    return face_ids


def _first_boss_face(model: StepModel) -> int:
    candidates = _boss_face_ids(model)
    if not candidates:
        raise SystemExit("no near-full cylindrical boss face was recognized")
    return candidates[0]


def _first_circle_edge_adjacent_to_face(model: StepModel, face_id: int) -> int:
    for edge_id in range(len(model.edges)):
        info = model.edge_info(edge_id)
        if info.get("curve") != "circle":
            continue
        adjacent_face_ids = tuple(int(item) for item in info.get("adjacent_face_ids", ()) or ())
        if face_id in adjacent_face_ids:
            return edge_id
    raise SystemExit(f"no circular edge adjacent to Face {face_id} was recognized")


def _axis_center(model: StepModel, face_id: int) -> tuple[float, float, float]:
    info = model.face_info(face_id)
    axis_point = info.get("axis_point")
    axis = info.get("axis")
    v_range = info.get("same_domain_v_range") or info.get("v_range")
    if not isinstance(axis_point, tuple) or not isinstance(axis, tuple) or not isinstance(v_range, tuple):
        raise SystemExit(f"axis center data is missing on Face {face_id}")
    v_mid = (float(v_range[0]) + float(v_range[1])) * 0.5
    return (
        float(axis_point[0]) + float(axis[0]) * v_mid,
        float(axis_point[1]) + float(axis[1]) * v_mid,
        float(axis_point[2]) + float(axis[2]) * v_mid,
    )


def _diameter(model: StepModel, face_id: int) -> float:
    value = model.face_info(face_id).get("diameter")
    if value is None:
        raise SystemExit(f"diameter is missing on Face {face_id}")
    return float(value)


def _height(model: StepModel, face_id: int) -> float:
    feature = model.feature_info(face_id)
    value = feature.get("same_domain_height_estimate")
    if value is None:
        value = model.face_info(face_id).get("height_estimate")
    if value is None:
        raise SystemExit(f"height is missing on Face {face_id}")
    return float(value)


def _nearest_boss(
    model: StepModel,
    target_diameter: float,
    target_center: tuple[float, float, float] | None = None,
) -> tuple[int, float, float, tuple[float, float, float], float]:
    best: tuple[int, float, float, tuple[float, float, float], float] | None = None
    for face_id in _boss_face_ids(model):
        diameter = _diameter(model, face_id)
        height = _height(model, face_id)
        center = _axis_center(model, face_id)
        diameter_error = abs(diameter - target_diameter)
        center_error = _distance(center, target_center) if target_center is not None else 0.0
        score = diameter_error + center_error
        if best is None or score < best[4]:
            best = (face_id, diameter, height, center, score)
    if best is None:
        raise SystemExit("no cylindrical boss face remained after edit")
    return best


def _run_diameter_case(target: float, tolerance: float) -> None:
    with tempfile.TemporaryDirectory(prefix="geom_param_boss_diam_") as temp_dir:
        model_path = Path(temp_dir) / "boss.step"
        _write_boss_model(model_path)
        model = StepModel.load(model_path)
        face_id = _first_boss_face(model)
        before = model.stats()
        center = _axis_center(model, face_id)
        current_height = _height(model, face_id)
        plan = model.cylindrical_boss_resize_plan(face_id, target)
        if plan["status"] == "blocked":
            raise SystemExit(f"diameter plan was blocked: {plan['message']}")
        result = model.resize_cylindrical_boss(face_id, target)
        after = model.stats()
        verified_face, diameter, height, verified_center, _ = _nearest_boss(model, target, center)
        error = abs(diameter - target)
        height_error = abs(height - current_height)
        if after.solids != before.solids:
            raise SystemExit(f"diameter changed solid count: before={before.solids}, after={after.solids}")
        if error > tolerance:
            raise SystemExit(f"diameter verification failed: target={target:g}, value={diameter:g}, error={error:g}")
        if height_error > tolerance:
            raise SystemExit(
                f"diameter changed boss height: before={current_height:g}, after={height:g}, error={height_error:g}"
            )
        if "verified_face=" not in result:
            raise SystemExit(f"diameter result did not report verified target Face: {result}")
        if "first_level_topology_matched=True" not in result:
            raise SystemExit(f"diameter result did not verify first-level topology: {result}")
        print("mode=diameter")
        print(f"source_face={face_id}")
        print(f"strategy={plan.get('resize_strategy')}")
        print(f"before={before}")
        print(f"after={after}")
        print(f"verified_face={verified_face}")
        print(f"target={target:.6f} value={diameter:.6f} height={height:.6f} center={verified_center} error={error:.6g}")
        print(result.encode("ascii", "backslashreplace").decode("ascii"))


def _run_axis_center_case(offset: float, tolerance: float) -> None:
    with tempfile.TemporaryDirectory(prefix="geom_param_boss_axis_") as temp_dir:
        model_path = Path(temp_dir) / "boss.step"
        _write_boss_model(model_path)
        model = StepModel.load(model_path)
        face_id = _first_boss_face(model)
        before = model.stats()
        current_center = _axis_center(model, face_id)
        current_diameter = _diameter(model, face_id)
        current_height = _height(model, face_id)
        target_center = (current_center[0] + offset, current_center[1], current_center[2])
        plan = model.cylindrical_boss_axis_move_plan(face_id, target_center)
        if plan["status"] == "blocked":
            raise SystemExit(f"axis_center plan was blocked: {plan['message']}")
        result = model.move_cylindrical_boss_axis(face_id, target_center)
        after = model.stats()
        verified_face, diameter, height, center, _ = _nearest_boss(model, current_diameter, target_center)
        center_error = _distance(center, target_center)
        height_error = abs(height - current_height)
        if after.solids != before.solids:
            raise SystemExit(f"axis_center changed solid count: before={before.solids}, after={after.solids}")
        if center_error > tolerance or abs(diameter - current_diameter) > tolerance:
            raise SystemExit(
                f"axis_center verification failed: target={target_center}, center={center}, "
                f"center_error={center_error:g}, diameter={diameter:g}"
            )
        if height_error > tolerance:
            raise SystemExit(
                f"axis_center changed boss height: before={current_height:g}, after={height:g}, error={height_error:g}"
            )
        if "first_level_topology_matched=True" not in result:
            raise SystemExit(f"axis_center result did not verify first-level topology: {result}")
        print("mode=axis_center")
        print(f"source_face={face_id}")
        print(f"strategy={plan.get('resize_strategy')}")
        print(f"before={before}")
        print(f"after={after}")
        print(f"verified_face={verified_face}")
        print(f"target={target_center} value={center} diameter={diameter:.6f} height={height:.6f} error={center_error:.6g}")
        print(result.encode("ascii", "backslashreplace").decode("ascii"))


def _run_circle_edge_axis_center_case(offset: float, tolerance: float) -> None:
    with tempfile.TemporaryDirectory(prefix="geom_param_boss_circle_edge_axis_") as temp_dir:
        model_path = Path(temp_dir) / "boss.step"
        _write_boss_model(model_path)
        model = StepModel.load(model_path)
        face_id = _first_boss_face(model)
        edge_id = _first_circle_edge_adjacent_to_face(model, face_id)
        before = model.stats()
        current_center = _axis_center(model, face_id)
        current_diameter = _diameter(model, face_id)
        current_height = _height(model, face_id)
        edge_info = model.edge_info(edge_id)
        edge_center = edge_info.get("center")
        if not isinstance(edge_center, tuple):
            raise SystemExit(f"circle edge center is missing on Edge {edge_id}")
        target_edge_center = (float(edge_center[0]) + offset, float(edge_center[1]), float(edge_center[2]))
        target_axis_center = (current_center[0] + offset, current_center[1], current_center[2])
        plan = model.circular_edge_axis_move_plan(edge_id, target_edge_center)
        if plan["status"] == "blocked":
            raise SystemExit(f"circle_edge_axis_center plan was blocked: {plan['message']}")
        result = model.move_circular_edge_axis_center(edge_id, target_edge_center)
        after = model.stats()
        verified_face, diameter, height, center, _ = _nearest_boss(model, current_diameter, target_axis_center)
        center_error = _distance(center, target_axis_center)
        height_error = abs(height - current_height)
        if after.solids != before.solids:
            raise SystemExit(f"circle_edge_axis_center changed solid count: before={before.solids}, after={after.solids}")
        if center_error > tolerance or abs(diameter - current_diameter) > tolerance:
            raise SystemExit(
                f"circle_edge_axis_center verification failed: target={target_axis_center}, center={center}, "
                f"center_error={center_error:g}, diameter={diameter:g}"
            )
        if height_error > tolerance:
            raise SystemExit(
                f"circle_edge_axis_center changed boss height: before={current_height:g}, after={height:g}, error={height_error:g}"
            )
        if "first_level_topology_matched=True" not in result:
            raise SystemExit(f"circle_edge_axis_center result did not verify first-level topology: {result}")
        print("mode=circle_edge_axis_center")
        print(f"source_edge={edge_id}")
        print(f"source_face={face_id}")
        print(f"strategy={plan.get('resize_strategy')}")
        print(f"delegated_mode={plan.get('circular_edge_cylinder_mode')}")
        print(f"before={before}")
        print(f"after={after}")
        print(f"verified_face={verified_face}")
        print(f"target_edge_center={target_edge_center}")
        print(f"target_axis_center={target_axis_center} value={center} diameter={diameter:.6f} height={height:.6f} error={center_error:.6g}")
        print(result.encode("ascii", "backslashreplace").decode("ascii"))


def _run_height_case(target: float, tolerance: float) -> None:
    with tempfile.TemporaryDirectory(prefix="geom_param_boss_height_") as temp_dir:
        model_path = Path(temp_dir) / "boss.step"
        _write_boss_model(model_path)
        model = StepModel.load(model_path)
        face_id = _first_boss_face(model)
        before = model.stats()
        current_diameter = _diameter(model, face_id)
        plan = model.cylindrical_boss_height_plan(face_id, target)
        if plan["status"] == "blocked":
            raise SystemExit(f"height plan was blocked: {plan['message']}")
        result = model.resize_cylindrical_boss_height(face_id, target)
        after = model.stats()
        verified_face, diameter, height, center, _ = _nearest_boss(model, current_diameter)
        error = abs(height - target)
        if after.solids != before.solids:
            raise SystemExit(f"height changed solid count: before={before.solids}, after={after.solids}")
        if error > tolerance:
            raise SystemExit(f"height verification failed: target={target:g}, value={height:g}, error={error:g}")
        if "first_level_topology_matched=True" not in result:
            raise SystemExit(f"height result did not verify first-level topology: {result}")
        print("mode=height")
        print(f"source_face={face_id}")
        print(f"strategy={plan.get('resize_strategy')}")
        print(f"before={before}")
        print(f"after={after}")
        print(f"verified_face={verified_face}")
        print(f"target={target:.6f} value={height:.6f} diameter={diameter:.6f} center={center} error={error:.6g}")
        print(result.encode("ascii", "backslashreplace").decode("ascii"))


def main() -> int:
    parser = argparse.ArgumentParser(description="Verify cylindrical boss edit operations.")
    parser.add_argument(
        "--mode",
        default="all",
        choices=["all", "diameter", "diameter_shrink", "height", "axis_center", "circle_edge_axis_center"],
        help="Boss edit mode to verify.",
    )
    parser.add_argument("--diameter", type=float, default=8.0)
    parser.add_argument("--diameter-shrink", type=float, default=4.0)
    parser.add_argument("--height", type=float, default=7.0)
    parser.add_argument("--axis-center", type=float, default=2.0, help="Axis-center X offset to verify.")
    parser.add_argument("--circle-edge-axis-center", type=float, default=2.0, help="Circle Edge center X offset to verify.")
    parser.add_argument("--tolerance", type=float, default=2e-4)
    args = parser.parse_args()

    cases = (
        [
            ("diameter", args.diameter),
            ("diameter_shrink", args.diameter_shrink),
            ("height", args.height),
            ("axis_center", args.axis_center),
            ("circle_edge_axis_center", args.circle_edge_axis_center),
        ]
        if args.mode == "all"
        else [(args.mode, getattr(args, args.mode.replace("-", "_")))]
    )
    for mode, target in cases:
        if mode in {"diameter", "diameter_shrink"}:
            _run_diameter_case(float(target), args.tolerance)
        elif mode == "height":
            _run_height_case(float(target), args.tolerance)
        elif mode == "axis_center":
            _run_axis_center_case(float(target), args.tolerance)
        elif mode == "circle_edge_axis_center":
            _run_circle_edge_axis_center_case(float(target), args.tolerance)
        else:
            raise SystemExit(f"unsupported mode: {mode}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())