Robot_Arm/arm/cad/helpers.py

# -*- coding: utf-8 -*-
"""CAD 放置与几何 helper：刚体变换、求交/包围盒、基础体、网格化与 STEP 读取。

实测接口数值在 arm/constants.py，仓库路径在 arm/paths.py——本模块不再持有它们。
"""
import math

import cadquery as cq
from OCP.gp import gp_Trsf, gp_Vec, gp_XYZ, gp_Mat, gp_Pnt, gp_Dir, gp_Ax1, gp_Ax2, gp_Quaternion
from OCP.BRepBuilderAPI import BRepBuilderAPI_Transform
from OCP.Bnd import Bnd_Box
from OCP.BRepBndLib import BRepBndLib
from OCP.BRepMesh import BRepMesh_IncrementalMesh
from OCP.BRep import BRep_Tool
from OCP.TopExp import TopExp_Explorer
from OCP.TopAbs import TopAbs_FACE, TopAbs_REVERSED, TopAbs_SOLID
from OCP.TopoDS import TopoDS

V = cq.Vector



# ---------------- geometry helpers --------------------------------------------------------------------

def _nrm(v):
    m = v.Magnitude()
    return gp_Vec(v.X() / m, v.Y() / m, v.Z() / m)


def _dot(a, b):
    return a.X() * b.X() + a.Y() * b.Y() + a.Z() * b.Z()


def _perp(e, z):
    d = _dot(e, z)
    return gp_Vec(e.X() - d * z.X(), e.Y() - d * z.Y(), e.Z() - d * z.Z())


def trs(pos=(0, 0, 0), xdir=(1, 0, 0), zdir=(0, 0, 1)):
    """Local->global rigid transform: local +Z -> zdir, local +X -> xdir, y = z x x.
    Uses SetValues (row-major 3x4) because gp_Quaternion / SetCols overloads in this
    OCP build do not map axes the way the docs suggest - verified against known points."""
    z = _nrm(gp_Vec(*zdir))
    x = _perp(gp_Vec(*xdir), z)
    if x.Magnitude() < 1e-9:
        x = _perp(gp_Vec(0, 0, 1), z)
    x = _nrm(x)
    y = z.Crossed(x)
    t = gp_Trsf()
    t.SetValues(x.X(), y.X(), z.X(), pos[0],
                x.Y(), y.Y(), z.Y(), pos[1],
                x.Z(), y.Z(), z.Z(), pos[2])
    return t


def map_pt(t, p):
    v = cq.Vertex.makeVertex(*p)
    w = cq.Shape(BRepBuilderAPI_Transform(v.wrapped, t, True).Shape())
    q = w.Vertices()[0].Center()
    return (q.x, q.y, q.z)


def place(shape, t):
    if shape is None:
        return None
    return cq.Shape(BRepBuilderAPI_Transform(shape.wrapped, t, True).Shape())


def bb(shape):
    b = Bnd_Box()
    BRepBndLib.Add_s(shape.wrapped, b, True)
    lo, hi = b.CornerMin(), b.CornerMax()
    return (lo.X(), lo.Y(), lo.Z()), (hi.X(), hi.Y(), hi.Z())


def bbox_str(shape):
    a, b = bb(shape)
    return "[%.1f %.1f %.1f .. %.1f %.1f %.1f] (%.1fx%.1fx%.1f)" % (
        a[0], a[1], a[2], b[0], b[1], b[2], b[0] - a[0], b[1] - a[1], b[2] - a[2])


def cyl(r, z0, z1, x=0.0, y=0.0):
    return cq.Solid.makeCylinder(r, z1 - z0, V(x, y, z0), V(0, 0, 1))


def ring(r0, r1, z0, z1):
    return cyl(r1, z0, z1).cut(cyl(r0, z0 - 0.001, z1 + 0.001))


def box(x0, x1, y0, y1, z0, z1):
    return cq.Solid.makeBox(x1 - x0, y1 - y0, z1 - z0, V(x0, y0, z0))


def rect_pocket(cx, cy, w, h, z0, z1, rot=0.0):
    s = box(cx - w / 2.0, cx + w / 2.0, cy - h / 2.0, cy + h / 2.0, z0 - 0.001, z1 + 0.001)
    if rot:
        s = s.rotated(V(0, 0, 0), V(0, 0, 1), rot)
    return s


def holes(shape, pts, d, z0, z1):
    """Through/ blind cylindrical cuts at the given (x, y) points."""
    cut = [cyl(d / 2.0, z0, z1, p[0], p[1]) for p in pts]
    return shape.cut(*cut) if cut else shape


def slot_holes(shape, pts, d, cb_d=None, cb_z0=None, cb_z1=None, z0=-1, z1=1):
    cuts = [cyl(d / 2.0, z0, z1, p[0], p[1]) for p in pts]
    if cb_d and cb_z0 is not None:
        cuts += [cyl(cb_d / 2.0, cb_z0, cb_z1, p[0], p[1]) for p in pts]
    return shape.cut(*cuts)


def polygon_pts(r, n, phase=0.0):
    return [(r * math.cos(math.radians(phase + 360.0 * k / n)),
             r * math.sin(math.radians(phase + 360.0 * k / n))) for k in range(n)]


def convex_prism(pts2d, z0, z1):
    """Solid from an ordered 2D outline (assumed convex or simple)."""
    wp = cq.Workplane("XY").workplane(centerSection=(0, 0, z0))
    w2 = [(p[0], p[1]) for p in pts2d]
    r = (cq.Workplane("XY", origin=(0, 0, z0))
         .polyline(w2 + [w2[0]]).close().extrude(z1 - z0, combine=False))
    return r.val()


def fillet_all_edges(shape, r):
    try:
        return shape.edges().fillet(r).val() if isinstance(shape, cq.Workplane) else \
            cq.Workplane(obj=shape).edges().fillet(r).val()
    except Exception:
        return shape


def mesh_of(shape, defl=0.08):
    BRepMesh_IncrementalMesh(shape.wrapped, defl, False, 0.55, True)
    verts, tris, base = [], [], 0
    e = TopExp_Explorer(shape.wrapped, TopAbs_FACE)
    while e.More():
        cur = e.Current()
        f = TopoDS.Face_s(cur)
        tri = BRep_Tool.Triangulation_s(f, cur.Location())
        if tri is not None:
            trsf = cur.Location().Transformation()
            for i in range(1, tri.NbNodes() + 1):
                p = tri.Node(i).Transformed(trsf)
                verts.append((round(p.X(), 2), round(p.Y(), 2), round(p.Z(), 2)))
            rev = f.Orientation() == TopAbs_REVERSED
            for i in range(1, tri.NbTriangles() + 1):
                t = tri.Triangle(i)
                a, b, c = t.Value(1) - 1, t.Value(2) - 1, t.Value(3) - 1
                tris.append((base + a, base + c, base + b) if rev else (base + a, base + b, base + c))
            base += tri.NbNodes()
        e.Next()
    return verts, tris


def solids_of(shape):
    out = []
    e = TopExp_Explorer(shape.wrapped, TopAbs_SOLID)
    while e.More():
        out.append(TopoDS.Solid_s(e.Current()))
        e.Next()
    return [cq.Shape(s) for s in out]


def read_step(p):
    return cq.importers.importStep(p)


def norm_angle(a):
    return (a + 180.0) % 360.0 - 180.0