pythonocc-step-editor/scripts/verify_prismatic_feature.py

from __future__ import annotations

from pathlib import Path
import sys
import tempfile

from OCC.Core.BRepAlgoAPI import BRepAlgoAPI_Cut, BRepAlgoAPI_Fuse
from OCC.Core.BRepPrimAPI import BRepPrimAPI_MakeBox
from OCC.Core.gp import 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.step_io import _write_step
from step_editor.model import StepModel
from step_editor.window_state import WindowStateMixin, _feature_dimension_keys


DEFAULT_MODEL = PROJECT_ROOT / "assets" / "models" / "cube_10mm.step"


def close_to(value: object, expected: float, tolerance: float = 1e-6) -> bool:
    try:
        return abs(float(value) - expected) <= tolerance
    except (TypeError, ValueError):
        return False


def find_feature(model: StepModel, feature_type: str) -> tuple[int, dict[str, object]]:
    matches = [
        (face_id, model.feature_info(face_id))
        for face_id in range(len(model.faces))
        if model.feature_info(face_id).get("feature_type") == feature_type
    ]
    if not matches:
        observed = sorted({str(model.feature_info(face_id).get("feature_type")) for face_id in range(len(model.faces))})
        raise AssertionError(f"expected {feature_type}, observed {observed}")
    return matches[0]


def find_prismatic_semantics(model: StepModel, semantics: str) -> tuple[int, dict[str, object]]:
    matches = [
        (face_id, model.feature_info(face_id))
        for face_id in range(len(model.faces))
        if model.feature_info(face_id).get("prismatic_feature_semantics") == semantics
        and model.feature_info(face_id).get("prismatic_profile_status") == "candidate"
    ]
    if not matches:
        observed = sorted(
            {
                str(model.feature_info(face_id).get("prismatic_feature_semantics"))
                for face_id in range(len(model.faces))
            }
        )
        raise AssertionError(f"expected prismatic semantics {semantics}, observed {observed}")
    return matches[0]


def assert_prismatic_sizes(
    info: dict[str, object], length: float, width: float, depth: float
) -> None:
    expected = {
        "prismatic_length": length,
        "prismatic_width": width,
        "prismatic_extrusion_estimate": depth,
    }
    for key, value in expected.items():
        if not close_to(info.get(key), value):
            raise AssertionError(f"{key}={info.get(key)!r}, expected {value}")


def assert_vector_close(
    value: object,
    expected: tuple[float, float, float],
    tolerance: float = 2e-4,
) -> None:
    if value is None:
        raise AssertionError(f"missing vector, expected {expected}")
    try:
        actual = tuple(float(item) for item in value)  # type: ignore[iteration-over-annotation]
    except (TypeError, ValueError):
        raise AssertionError(f"invalid vector {value!r}, expected {expected}") from None
    if len(actual) != 3:
        raise AssertionError(f"invalid vector {value!r}, expected {expected}")
    errors = tuple(abs(actual_item - expected_item) for actual_item, expected_item in zip(actual, expected))
    if max(errors) > tolerance:
        raise AssertionError(f"vector={actual!r}, expected={expected!r}, errors={errors!r}")


def write_rectangular_pocket_model(path: Path) -> None:
    base = BRepPrimAPI_MakeBox(30.0, 20.0, 5.0).Shape()
    pocket_tool = BRepPrimAPI_MakeBox(gp_Pnt(10.0, 6.0, 2.0), 10.0, 8.0, 4.0).Shape()
    _write_step(BRepAlgoAPI_Cut(base, pocket_tool).Shape(), path)


def write_rectangular_boss_model(path: Path) -> None:
    base = BRepPrimAPI_MakeBox(30.0, 20.0, 5.0).Shape()
    boss_tool = BRepPrimAPI_MakeBox(gp_Pnt(10.0, 6.0, 5.0), 10.0, 8.0, 3.0).Shape()
    _write_step(BRepAlgoAPI_Fuse(base, boss_tool).Shape(), path)


def write_rectangular_multistep_boss_model(path: Path) -> None:
    base = BRepPrimAPI_MakeBox(30.0, 20.0, 5.0).Shape()
    lower = BRepPrimAPI_MakeBox(gp_Pnt(8.0, 5.0, 5.0), 14.0, 10.0, 3.0).Shape()
    upper = BRepPrimAPI_MakeBox(gp_Pnt(12.0, 8.0, 8.0), 6.0, 4.0, 2.0).Shape()
    _write_step(BRepAlgoAPI_Fuse(BRepAlgoAPI_Fuse(base, lower).Shape(), upper).Shape(), path)


def feature_dimension_rows(
    model: StepModel,
    face_id: int,
    info: dict[str, object],
) -> tuple[tuple[str, str, str], ...]:
    state = object.__new__(WindowStateMixin)
    state.model = model
    state.operation_in_progress = False
    state.scan_in_progress = False
    state.load_in_progress = False
    state.selected_face_id = face_id
    state.selected_edge_id = None
    state.selected_kind = "feature"
    state.selected_part_id = int(info["part_id"])
    state.selected_solid_id = int(info["solid_id"])
    specs, _used = state._editable_property_specs(info)
    filtered = state._feature_property_specs(specs, info)
    dimensions = [spec for spec in filtered if spec.get("parameter_role") == "dimension"]
    return tuple(
        (
            str(spec.get("key")),
            str(spec.get("label")),
            str(spec.get("current_text")),
        )
        for spec in dimensions
    )


def assert_rectangular_feature_ui(
    model: StepModel,
    face_id: int,
    info: dict[str, object],
    expected_rows: tuple[tuple[str, str, str], ...],
) -> None:
    actual = feature_dimension_rows(model, face_id, info)
    if actual != expected_rows:
        raise AssertionError(f"unexpected rectangular feature UI parameters: {actual}")


def assert_prismatic_feature_display_labels() -> None:
    state = object.__new__(WindowStateMixin)
    if state._feature_display_label("矩形口袋候选") != "矩形槽/口袋":
        raise AssertionError("rectangular pocket should be displayed as a rectangular slot/pocket")
    if state._feature_display_label("矩形凸台候选") != "矩形凸台":
        raise AssertionError("rectangular boss display label should not expose candidate wording")


def assert_shell_depth_edit(
    model_path: Path,
    feature_type: str,
    target_depth: float,
    tolerance: float = 2e-4,
) -> None:
    model = StepModel.load(model_path)
    face_id, info = find_feature(model, feature_type)
    before = model.stats()
    assert_prismatic_sizes(info, 10.0, 8.0, 3.0)
    plan = model.shell_thickness_plan(face_id, target_depth)
    if plan["status"] == "blocked":
        raise AssertionError(f"{feature_type} depth plan was blocked: {plan['message']}")
    result = model.resize_shell_thickness(face_id, target_depth)
    after = model.stats()
    if after.solids != before.solids:
        raise AssertionError(f"{feature_type} depth edit changed solid count: before={before}, after={after}")

    _verified_face_id, verified = find_feature(model, feature_type)
    assert_prismatic_sizes(verified, 10.0, 8.0, target_depth)
    error = abs(float(verified["prismatic_extrusion_estimate"]) - target_depth)
    if error > tolerance:
        raise AssertionError(
            f"{feature_type} depth verification failed: target={target_depth:g}, "
            f"value={verified['prismatic_extrusion_estimate']!r}, error={error:g}"
        )
    if "resize completed" not in result.lower():
        raise AssertionError(f"{feature_type} depth result did not report completion: {result}")


def assert_rectangular_size_edit(
    model_path: Path,
    semantics: str,
    axis: str,
    target_size: float,
    expected_length: float,
    expected_width: float,
    expected_depth: float = 3.0,
    tolerance: float = 2e-4,
    initial_length: float = 10.0,
    initial_width: float = 8.0,
    initial_depth: float = 3.0,
) -> None:
    model = StepModel.load(model_path)
    face_id, info = find_prismatic_semantics(model, semantics)
    before = model.stats()
    assert_prismatic_sizes(info, initial_length, initial_width, initial_depth)
    axis_key = "height" if axis == "width" else "width"
    plan = model.face_size_local_resize_plan(face_id, target_size, axis_key)
    if plan["status"] == "blocked":
        raise AssertionError(f"{semantics} {axis} plan was blocked: {plan['message']}")
    if plan.get("resize_strategy") != "rectangular-prismatic-profile-rebuild":
        raise AssertionError(f"unexpected {semantics} {axis} strategy: {plan.get('resize_strategy')}")
    result = model.resize_face_size_local(face_id, target_size, axis_key)
    after = model.stats()
    if after.solids != before.solids:
        raise AssertionError(f"{semantics} {axis} edit changed solid count: before={before}, after={after}")

    _verified_face_id, verified = find_prismatic_semantics(model, semantics)
    assert_prismatic_sizes(verified, expected_length, expected_width, expected_depth)
    length_error = abs(float(verified["prismatic_length"]) - expected_length)
    width_error = abs(float(verified["prismatic_width"]) - expected_width)
    if length_error > tolerance or width_error > tolerance:
        raise AssertionError(
            f"{semantics} {axis} verification failed: "
            f"length={verified['prismatic_length']!r}, width={verified['prismatic_width']!r}, "
            f"expected=({expected_length:g}, {expected_width:g})"
        )
    if "rectangular prismatic feature size resize completed" not in result.lower():
        raise AssertionError(f"{semantics} {axis} result did not report completion: {result}")


def assert_rectangular_center_edit(
    model_path: Path,
    semantics: str,
    target_center: tuple[float, float, float],
    expected_center: tuple[float, float, float],
    expected_length: float = 10.0,
    expected_width: float = 8.0,
    expected_depth: float = 3.0,
    initial_length: float = 10.0,
    initial_width: float = 8.0,
    initial_depth: float = 3.0,
) -> None:
    model = StepModel.load(model_path)
    face_id, info = find_prismatic_semantics(model, semantics)
    before = model.stats()
    assert_prismatic_sizes(info, initial_length, initial_width, initial_depth)
    plan = model.face_center_local_move_plan(face_id, target_center)
    if plan["status"] == "blocked":
        raise AssertionError(f"{semantics} center move plan was blocked: {plan['message']}")
    if plan.get("resize_strategy") != "rectangular-prismatic-center-rebuild":
        raise AssertionError(f"unexpected {semantics} center move strategy: {plan.get('resize_strategy')}")
    result = model.move_face_center_local(face_id, target_center)
    after = model.stats()
    if after.solids != before.solids:
        raise AssertionError(f"{semantics} center move changed solid count: before={before}, after={after}")

    _verified_face_id, verified = find_prismatic_semantics(model, semantics)
    assert_prismatic_sizes(verified, expected_length, expected_width, expected_depth)
    assert_vector_close(verified.get("area_center") or verified.get("bbox_center"), expected_center)
    if "rectangular prismatic feature center move completed" not in result.lower():
        raise AssertionError(f"{semantics} center move result did not report completion: {result}")


def assert_rectangular_depth_edit_by_semantics(
    model_path: Path,
    semantics: str,
    target_depth: float,
    expected_length: float,
    expected_width: float,
    initial_depth: float,
    tolerance: float = 2e-4,
) -> None:
    model = StepModel.load(model_path)
    face_id, info = find_prismatic_semantics(model, semantics)
    before = model.stats()
    assert_prismatic_sizes(info, expected_length, expected_width, initial_depth)
    plan = model.shell_thickness_plan(face_id, target_depth)
    if plan["status"] == "blocked":
        raise AssertionError(f"{semantics} depth plan was blocked: {plan['message']}")
    result = model.resize_shell_thickness(face_id, target_depth)
    after = model.stats()
    if after.solids != before.solids:
        raise AssertionError(f"{semantics} depth edit changed solid count: before={before}, after={after}")

    _verified_face_id, verified = find_prismatic_semantics(model, semantics)
    assert_prismatic_sizes(verified, expected_length, expected_width, target_depth)
    error = abs(float(verified["prismatic_extrusion_estimate"]) - target_depth)
    if error > tolerance:
        raise AssertionError(
            f"{semantics} depth verification failed: target={target_depth:g}, "
            f"value={verified['prismatic_extrusion_estimate']!r}, error={error:g}"
        )
    if "resize completed" not in result.lower():
        raise AssertionError(f"{semantics} depth result did not report completion: {result}")


def assert_rectangular_axial_center_move_blocked(
    model_path: Path,
    semantics: str,
    target_center: tuple[float, float, float],
) -> None:
    model = StepModel.load(model_path)
    face_id, _info = find_prismatic_semantics(model, semantics)
    plan = model.face_center_local_move_plan(face_id, target_center)
    if plan["status"] != "blocked":
        raise AssertionError(f"{semantics} axial center move should be blocked: {plan}")
    if plan.get("resize_strategy") != "rectangular-prismatic-center-rebuild":
        raise AssertionError(f"unexpected blocked {semantics} center strategy: {plan.get('resize_strategy')}")


def assert_rectangular_role_swap_blocked(model_path: Path, semantics: str) -> None:
    model = StepModel.load(model_path)
    face_id, _info = find_prismatic_semantics(model, semantics)
    plan = model.face_size_local_resize_plan(face_id, 7.0, "width")
    if plan["status"] != "blocked":
        raise AssertionError(f"{semantics} length/width role swap should be blocked: {plan}")


def main() -> int:
    model = StepModel.load(DEFAULT_MODEL)
    if len(model.faces) != 6:
        raise AssertionError(f"expected 6 cube faces, got {len(model.faces)}")

    for face_id in range(len(model.faces)):
        info = model.feature_info(face_id)
        if info.get("prismatic_profile_status") != "candidate":
            raise AssertionError(f"Face {face_id} was not recognized as a rectangular profile")
        if info.get("prismatic_extrusion_status") != "candidate":
            raise AssertionError(f"Face {face_id} was not recognized as a connected extrusion")
        for key in ("prismatic_length", "prismatic_width", "prismatic_extrusion_estimate"):
            if not close_to(info.get(key), 10.0):
                raise AssertionError(f"Face {face_id} {key}={info.get(key)!r}, expected 10")
        if len(tuple(info.get("prismatic_connected_side_face_ids", ()))) != 4:
            raise AssertionError(f"Face {face_id} does not have four connected side faces")

    face_id = 0
    info = model.feature_info(face_id)
    expected_keys = ("local_face_width", "local_face_height", "shell_thickness_estimate")
    if _feature_dimension_keys(info) != expected_keys:
        raise AssertionError(f"unexpected cube feature dimensions: {_feature_dimension_keys(info)}")
    expected = (
        ("local_face_width", "长度", "10"),
        ("local_face_height", "宽度", "10"),
        ("shell_thickness_estimate", "高度/深度", "10"),
    )
    assert_rectangular_feature_ui(model, face_id, info, expected)
    assert_prismatic_feature_display_labels()

    with tempfile.TemporaryDirectory(prefix="step-editor-prismatic-") as temp_dir:
        temp_path = Path(temp_dir)
        pocket_path = temp_path / "rectangular-pocket.step"
        write_rectangular_pocket_model(pocket_path)
        pocket_model = StepModel.load(pocket_path)
        pocket_face_id, pocket_info = find_feature(pocket_model, "矩形口袋候选")
        assert_prismatic_sizes(pocket_info, 10.0, 8.0, 3.0)
        if pocket_info.get("prismatic_reference_source") != "side-wall-topology":
            raise AssertionError(f"unexpected pocket reference: {pocket_info.get('prismatic_reference_source')}")
        assert_rectangular_feature_ui(
            pocket_model,
            pocket_face_id,
            pocket_info,
            (
                ("local_face_width", "长度", "10"),
                ("local_face_height", "宽度", "8"),
                ("shell_thickness_estimate", "高度/深度", "3"),
                ("face_center_position", "中心", "(15, 10, 2)"),
            ),
        )

        boss_path = temp_path / "rectangular-boss.step"
        write_rectangular_boss_model(boss_path)
        boss_model = StepModel.load(boss_path)
        boss_face_id, boss_info = find_feature(boss_model, "矩形凸台候选")
        assert_prismatic_sizes(boss_info, 10.0, 8.0, 3.0)
        if boss_info.get("prismatic_reference_source") not in {
            "side-wall-topology", "overlapping-plane"
        }:
            raise AssertionError(f"unexpected boss reference: {boss_info.get('prismatic_reference_source')}")
        assert_rectangular_feature_ui(
            boss_model,
            boss_face_id,
            boss_info,
            (
                ("local_face_width", "长度", "10"),
                ("local_face_height", "宽度", "8"),
                ("shell_thickness_estimate", "高度/深度", "3"),
                ("face_center_position", "中心", "(15, 10, 8)"),
            ),
        )

        boss_edit_path = temp_path / "rectangular-boss-edit.step"
        write_rectangular_boss_model(boss_edit_path)
        assert_shell_depth_edit(boss_edit_path, "矩形凸台候选", 4.0)

        pocket_edit_path = temp_path / "rectangular-pocket-edit.step"
        write_rectangular_pocket_model(pocket_edit_path)
        assert_shell_depth_edit(pocket_edit_path, "矩形口袋候选", 4.0)

        for name, writer, semantics in (
            ("boss", write_rectangular_boss_model, "additive-boss"),
            ("pocket", write_rectangular_pocket_model, "subtractive-pocket"),
        ):
            length_edit_path = temp_path / f"rectangular-{name}-length-edit.step"
            writer(length_edit_path)
            assert_rectangular_size_edit(length_edit_path, semantics, "length", 12.0, 12.0, 8.0)

            width_edit_path = temp_path / f"rectangular-{name}-width-edit.step"
            writer(width_edit_path)
            assert_rectangular_size_edit(width_edit_path, semantics, "width", 9.0, 10.0, 9.0)

            center_z = 8.0 if semantics == "additive-boss" else 2.0
            center_edit_path = temp_path / f"rectangular-{name}-center-edit.step"
            writer(center_edit_path)
            assert_rectangular_center_edit(
                center_edit_path,
                semantics,
                (17.0, 11.0, center_z),
                (17.0, 11.0, center_z),
            )

            axial_center_path = temp_path / f"rectangular-{name}-axial-center.step"
            writer(axial_center_path)
            axial_z = center_z + (1.0 if semantics == "additive-boss" else -1.0)
            assert_rectangular_axial_center_move_blocked(
                axial_center_path,
                semantics,
                (15.0, 10.0, axial_z),
            )

            blocked_path = temp_path / f"rectangular-{name}-role-swap.step"
            writer(blocked_path)
            assert_rectangular_role_swap_blocked(blocked_path, semantics)

        multistep_path = temp_path / "rectangular-multistep-boss.step"
        write_rectangular_multistep_boss_model(multistep_path)
        multistep_model = StepModel.load(multistep_path)
        multistep_face_id, multistep_info = find_prismatic_semantics(multistep_model, "additive-boss")
        assert_prismatic_sizes(multistep_info, 6.0, 4.0, 2.0)
        assert_rectangular_feature_ui(
            multistep_model,
            multistep_face_id,
            multistep_info,
            (
                ("local_face_width", "长度", "6"),
                ("local_face_height", "宽度", "4"),
                ("shell_thickness_estimate", "高度/深度", "2"),
                ("face_center_position", "中心", "(15, 10, 10)"),
            ),
        )

        multistep_length_path = temp_path / "rectangular-multistep-boss-length.step"
        write_rectangular_multistep_boss_model(multistep_length_path)
        assert_rectangular_size_edit(
            multistep_length_path,
            "additive-boss",
            "length",
            7.0,
            7.0,
            4.0,
            2.0,
            initial_length=6.0,
            initial_width=4.0,
            initial_depth=2.0,
        )

        multistep_width_path = temp_path / "rectangular-multistep-boss-width.step"
        write_rectangular_multistep_boss_model(multistep_width_path)
        assert_rectangular_size_edit(
            multistep_width_path,
            "additive-boss",
            "width",
            5.0,
            6.0,
            5.0,
            2.0,
            initial_length=6.0,
            initial_width=4.0,
            initial_depth=2.0,
        )

        multistep_depth_path = temp_path / "rectangular-multistep-boss-depth.step"
        write_rectangular_multistep_boss_model(multistep_depth_path)
        assert_rectangular_depth_edit_by_semantics(
            multistep_depth_path,
            "additive-boss",
            2.5,
            6.0,
            4.0,
            2.0,
        )

        multistep_center_path = temp_path / "rectangular-multistep-boss-center.step"
        write_rectangular_multistep_boss_model(multistep_center_path)
        assert_rectangular_center_edit(
            multistep_center_path,
            "additive-boss",
            (16.0, 10.5, 10.0),
            (16.0, 10.5, 10.0),
            expected_length=6.0,
            expected_width=4.0,
            expected_depth=2.0,
            initial_length=6.0,
            initial_width=4.0,
            initial_depth=2.0,
        )

        multistep_axial_path = temp_path / "rectangular-multistep-boss-axial-center.step"
        write_rectangular_multistep_boss_model(multistep_axial_path)
        assert_rectangular_axial_center_move_blocked(
            multistep_axial_path,
            "additive-boss",
            (15.0, 10.0, 10.5),
        )

    print("prismatic feature recognition and size/depth/center/multistep edit ok")
    return 0


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