0705530001

#!/usr/bin/env python3
"""Pose a manufacturer STEP per the Adom standard by rotation + translation only, keeping its XCAF content
(names, any colours) untouched: the original free shapes are wrapped in a new assembly carrying the location.

usage: pose_vendor.py <series> <in.stp> <footprint.kicad_mod> <out.step> <out.json>

Rotation per series (Molex axes as delivered, checked by sections of each file):
  70553, 74164 : Molex +z (toward the mating face) -> +X, +y (up) -> +Z, +x -> +Y
  87833        : Molex +z -> +X, -y (up) -> +Z, -x -> +Y
  90130        : Molex +y (toward the mating face) -> -X, -z (up) -> +Z, -x -> +Y
Translation: seat (lowest housing level: bisection for the first thin slab wider than 3 mm) to z = 0; the pad-1 tail (outer row, max Y end)
onto the origin. Then every tail found 1 mm below the seat is matched to its footprint pad and the offsets
reported.
"""
import json, math, re, sys
import cadquery as cq
from OCP.STEPCAFControl import STEPCAFControl_Reader, STEPCAFControl_Writer
from OCP.STEPControl import STEPControl_AsIs
from OCP.TDocStd import TDocStd_Document
from OCP.TCollection import TCollection_ExtendedString
from OCP.XCAFDoc import XCAFDoc_DocumentTool
from OCP.TDF import TDF_ChildIterator
from OCP.gp import gp_Trsf, gp_Mat, gp_XYZ
from OCP.TopLoc import TopLoc_Location

ROT = {  # rows: new X, Y, Z expressed in Molex axes
    '70553': ((0, 0, 1), (1, 0, 0), (0, 1, 0)),
    '74164': ((0, 0, 1), (1, 0, 0), (0, 1, 0)),
    '87833': ((0, 0, 1), (-1, 0, 0), (0, -1, 0)),
    '90130': ((0, -1, 0), (-1, 0, 0), (0, 0, -1)),
}
FWD = {'70553': 1, '74164': 1, '87833': 1, '90130': -1}   # mating face toward +X or -X


def rot_trsf(R, t=(0, 0, 0)):
    tr = gp_Trsf()
    tr.SetValues(R[0][0], R[0][1], R[0][2], t[0], R[1][0], R[1][1], R[1][2], t[1], R[2][0], R[2][1], R[2][2], t[2])
    return tr


def pads_of(kmod):
    txt = open(kmod).read()
    return [(n, float(x), float(y)) for n, x, y in re.findall(
        r'\(pad "([^"]+)" thru_hole \w+\s*\(at ([-\d.]+) ([-\d.]+)\)', txt)]


def slab_solids(shape, z, dz=0.02):
    bb = shape.BoundingBox()
    box = cq.Solid.makeBox(bb.xlen + 10, bb.ylen + 10, dz, cq.Vector(bb.xmin - 5, bb.ymin - 5, z - dz / 2))
    return shape.intersect(box).Solids()


def main(series, src, kmod, out, rep):
    R = ROT[series]
    s0 = cq.importers.importStep(src).val()
    from OCP.BRepBuilderAPI import BRepBuilderAPI_Transform
    rs = cq.Shape.cast(BRepBuilderAPI_Transform(s0.wrapped, rot_trsf(R), True).Shape())
    bb = rs.BoundingBox()
    # seat: lowest z at which a thin slab cuts something wider than 3 mm (below it only tails and pegs exist);
    # bisection between the tail tips and the top
    def housing_at(z):
        return any(max(b.xlen, b.ylen) > 3.0 for b in (x.BoundingBox() for x in slab_solids(rs, z, 0.002)))
    lo, hi = bb.zmin + 0.3, bb.zmax - 0.5
    assert not housing_at(lo) and housing_at(hi)
    while hi - lo > 0.0005:
        mid = (lo + hi) / 2
        if housing_at(mid): hi = mid
        else: lo = mid
    seat = (lo + hi) / 2
    tails = [x.BoundingBox() for x in slab_solids(rs, seat - 1.0)]
    tails = [b for b in tails if b.xlen < 1.0 and b.ylen < 1.0]
    cs = [((b.xmin + b.xmax) / 2, (b.ymin + b.ymax) / 2) for b in tails]
    f = FWD[series]
    xs = sorted({round(c[0], 2) for c in cs})
    outer_x = min(xs) if f > 0 else max(xs)                 # row farther from the mating face
    row = [c for c in cs if abs(c[0] - outer_x) < 0.2]
    p1 = max(row, key=lambda c: c[1])
    t = (-p1[0], -p1[1], -seat)
    # write the posed STEP: wrap the original free shapes in an assembly carrying the location
    doc = TDocStd_Document(TCollection_ExtendedString('doc'))
    r = STEPCAFControl_Reader(); r.SetColorMode(True); r.SetNameMode(True); r.SetLayerMode(True)
    r.ReadFile(src); r.Transfer(doc)
    st = XCAFDoc_DocumentTool.ShapeTool_s(doc.Main())
    labs = []
    it = TDF_ChildIterator(st.Label(), False)
    while it.More():
        l = it.Value()
        if st.IsFree_s(l) and st.IsTopLevel(l): labs.append(l)
        it.Next()
    tr = rot_trsf(R, t)
    asm = st.NewShape()
    for lab in labs:
        st.AddComponent(asm, lab, TopLoc_Location(tr))
    st.UpdateAssemblies()
    w = STEPCAFControl_Writer(); w.SetColorMode(True); w.SetNameMode(True); w.SetLayerMode(True)
    w.Transfer(doc, STEPControl_AsIs); w.Write(out)
    # check: tails onto pads
    ps = cq.importers.importStep(out)
    sh = cq.Compound.makeCompound(ps.solids().vals())
    pb = sh.BoundingBox()
    tl = [x.BoundingBox() for x in slab_solids(sh, -1.0)]
    tl = [((b.xmin + b.xmax) / 2, (b.ymin + b.ymax) / 2, b.xlen, b.ylen) for b in tl if b.xlen < 1.0 and b.ylen < 1.0]
    pads = pads_of(kmod)
    match = []
    for n, px, py in pads:
        best = min(tl, key=lambda c: (c[0] - px) ** 2 + (c[1] + py) ** 2)
        match.append(dict(pad=n, pad_xy=[px, py], tail_xy=[round(best[0], 3), round(-best[1], 3)],
                          off=round(math.hypot(best[0] - px, best[1] + py), 3), tail_size=[round(best[2], 3), round(best[3], 3)]))
    res = dict(series=series, source=src, rotation_rows_in_source_axes=R, translate_after_rotation_mm=[round(v, 4) for v in t],
               seat_source_frame_after_rotation=round(seat, 4), posed_bbox=[round(v, 3) for v in (pb.xmin, pb.xmax, pb.ymin, pb.ymax, pb.zmin, pb.zmax)],
               tails_found=len(tl), pads=len(pads), max_tail_offset_mm=max(m['off'] for m in match), match=match)
    json.dump(res, open(rep, 'w'), indent=1)
    print(json.dumps({k: v for k, v in res.items() if k != 'match'}))


if __name__ == '__main__':
    main(*sys.argv[1:6])