Петля print-in-place

"""Shared CAD primitives, rigid poses, mesh export and geometric checks. Units: mm."""

import FreeCAD as App

# isort: split
import argparse
import json
import math
from collections.abc import Callable
from dataclasses import dataclass
from functools import lru_cache
from pathlib import Path
from tempfile import TemporaryDirectory

import MeshPart
import Part

V = App.Vector
WHITE, ORANGE, METAL = [0.88, 0.89, 0.86], [1, 0.42, 0.09], [0.35, 0.36, 0.34]


def box(x, y, z, dx, dy, dz):
    return Part.makeBox(dx, dy, dz, V(x, y, z))


def cyl(r, h, at=(0, 0, 0), axis=(0, 0, 1)):
    return Part.makeCylinder(r, h, V(*at), V(*axis))


def cone(r1, r2, height, at=None):
    at = V() if at is None else at
    return (
        Part.makeCylinder(r1, height, at)
        if abs(r1 - r2) < 1e-9
        else Part.makeCone(r1, r2, height, at)
    )


def polygon(points, z, height):
    pts = [V(x, y, z) for x, y in points]
    return Part.Face(Part.makePolygon(pts + pts[:1])).extrude(V(0, 0, height))


def section(points, y, width):
    pts = [V(x, y, z) for x, z in points]
    return Part.Face(Part.makePolygon(pts + pts[:1])).extrude(V(0, width, 0))


def join(*shapes):
    return (
        shapes[0].multiFuse(list(shapes[1:])).removeSplitter()
        if len(shapes) > 1
        else shapes[0]
    )


def ring(outer, inner, height, z=0):
    return cyl(outer, height, (0, 0, z)).cut(cyl(inner, height + 2, (0, 0, z - 1)))


def sector_ring(outer, inner, start, end, height, z=0):
    # Polygonal arcs give the split lips reproducible printable geometry.
    angles = [math.radians(start + (end - start) * i / 40) for i in range(41)]
    pts = [(outer * math.cos(a), outer * math.sin(a)) for a in angles]
    pts += [(inner * math.cos(a), inner * math.sin(a)) for a in reversed(angles)]
    return polygon(pts, z, height)


def hexagon(radius, height, z=0):
    return polygon(
        [
            (radius * math.cos(i * math.pi / 3), radius * math.sin(i * math.pi / 3))
            for i in range(6)
        ],
        z,
        height,
    )


@lru_cache(maxsize=48)
def thread(radius, pitch, length, extra=0):
    if extra:
        matrix = App.Matrix()
        matrix.A11 = matrix.A22 = 1 + extra / radius
        return thread(radius, pitch, length, 0).transformGeometry(matrix)
    root, major = radius - pitch * 0.42, radius
    if root <= 0:
        raise ValueError("Thread root radius must be positive")
    flat = min(0.09, pitch * 0.08)
    inner = root - flat
    half = flat + (major - inner) / math.sqrt(3)
    start = -pitch
    points = [
        V(inner, 0, start - half),
        V(major, 0, start - flat),
        V(major, 0, start + flat),
        V(inner, 0, start + half),
    ]
    profile = Part.makePolygon(points + points[:1])
    helix = Part.Wire(Part.makeHelix(pitch, length + 2 * pitch, root).Edges).translated(
        V(0, 0, start)
    )
    rib = helix.makePipeShell([profile], True, True)
    stock = cyl(root, length + 4 * pitch, (0, 0, -2 * pitch)).fuse(rib)
    return stock.common(cyl(major + 0.1, length)).removeSplitter()


def phase(t, begin=0, end=1):
    return max(0, min(1, (t - begin) / (end - begin)))


@dataclass
class Pose:
    offset: tuple = (0, 0, 0)
    angle: float = 0
    axis: tuple = (0, 0, 1)
    pivot: tuple = (0, 0, 0)

    def placement(self):
        rotation = App.Rotation(V(*self.axis), self.angle)
        pivot = V(*self.pivot)
        return App.Placement(
            V(*self.offset) + pivot - rotation.multVec(pivot), rotation
        )

    def data(self, t):
        return {
            "t": t,
            "translation": list(self.offset),
            "angle": self.angle,
            "axis": list(self.axis),
            "pivot": list(self.pivot),
        }


@dataclass
class Piece:
    id: str
    name: str
    shape: object
    motion: Callable = lambda t: Pose()
    deform: Callable | None = None
    color: list | None = None
    role: str = "printed"
    note: str = ""
    fabrication_shape: object | None = None

    def at(self, t):
        shape = (self.deform(t) if self.deform else self.shape).copy()
        shape.Placement = self.motion(t).placement().multiply(shape.Placement)
        return shape


def params(defaults, supplied=None):
    supplied = {} if supplied is None else supplied
    if not isinstance(supplied, dict):
        raise TypeError("Parameters must be a JSON object")
    unknown = set(supplied) - set(defaults)
    if unknown:
        raise ValueError(f"Unknown parameters: {sorted(unknown)}")
    result = defaults | supplied
    for key, value in result.items():
        if isinstance(value, bool) or not isinstance(value, (int, float)):
            raise TypeError(f"{key} must be numeric")
        if not math.isfinite(value) or value < 0:
            raise ValueError(f"{key} must be finite and non-negative")
        if value == 0 and key not in {"clearance", "interference", "gap", "preload"}:
            raise ValueError(f"{key} must be positive")
    return result


def require(condition, message):
    if not condition:
        raise ValueError(message)


def mesh_data(shape):
    vertices, triangles = shape.tessellate(0.10)
    return {
        "vertices": [[round(v.x, 6), round(v.y, 6), round(v.z, 6)] for v in vertices],
        "triangles": [list(tri) for tri in triangles],
        "edges": [
            [
                [round(v.x, 6), round(v.y, 6), round(v.z, 6)]
                for v in edge.discretize(Deflection=0.12)
            ]
            for edge in shape.Edges
        ],
        "volume": shape.Volume,
    }


def check(parts, mode="rigid"):
    final = []
    for part in parts:
        shape = part.at(1)
        assert shape.isValid() and len(shape.Solids) == 1 and shape.Volume > 1e-7, (
            part.id
        )
        final.append(shape)
    overlaps = []
    motion_overlaps = []
    maximum = 0.0
    for i, left in enumerate(parts):
        for j in range(i + 1, len(parts)):
            right = parts[j]
            volume = final[i].common(final[j]).Volume
            if volume > 1e-5:
                overlaps.append(
                    {"parts": [left.id, right.id], "volume_mm3": round(volume, 6)}
                )
            pair_max, pair_time = 0, 0
            for k in range(9):
                volume = left.at(k / 8).common(right.at(k / 8)).Volume
                if volume > pair_max:
                    pair_max, pair_time = volume, k / 8
            maximum = max(maximum, pair_max)
            if pair_max > 1e-5:
                motion_overlaps.append(
                    {
                        "parts": [left.id, right.id],
                        "volume_mm3": round(pair_max, 6),
                        "t": pair_time,
                    }
                )
    expected = {
        "interference": {frozenset(("receiver", "pin"))},
        "thermal-fit": {frozenset(("receiver", "insert"))},
        "thread-forming": {frozenset(("receiver", "screw"))},
    }.get(mode, set())
    unexpected = [
        item for item in motion_overlaps if frozenset(item["parts"]) not in expected
    ]
    require(not unexpected, f"Детали пересекаются при сборке: {unexpected}")
    return {
        "valid_solids": len(parts),
        "final_overlaps": overlaps,
        "motion_sample_count": 9,
        "maximum_motion_overlap_mm3": round(maximum, 6),
        "mode": mode,
        "motion_overlaps": motion_overlaps,
        "physical_test": "pending",
    }


def export(model_id, folder, values, parts, note, mode="rigid"):
    checks = check(parts, mode)
    folder = Path(folder)
    folder.mkdir(parents=True, exist_ok=True)
    with TemporaryDirectory(dir=folder) as temp:
        work = Path(temp)
        doc = App.newDocument("FastenerDemo")
        objects, payload = [], []
        for index, piece in enumerate(parts):
            shape = piece.at(1)
            obj = doc.addObject("Part::Feature", piece.id)
            obj.Label, obj.Shape = piece.name, shape
            obj.addProperty("App::PropertyString", "GeneratorParameters", "Generator")
            obj.GeneratorParameters = json.dumps(values)
            obj.addProperty("App::PropertyString", "Role", "Generator")
            obj.Role = piece.role
            objects.append(obj)
            Part.export([obj], str(work / f"{piece.id}.step"))
            printable = (
                piece.fabrication_shape.copy()
                if piece.fabrication_shape is not None
                else piece.shape.copy()
            )
            b = printable.BoundBox
            printable.translate(V(-b.XMin, -b.YMin, -b.ZMin))
            mesh = MeshPart.meshFromShape(
                Shape=printable, LinearDeflection=0.1, AngularDeflection=0.18
            )
            assert mesh.isSolid(), piece.id
            if piece.role == "printed":
                mesh.write(str(work / f"{piece.id}.stl"))
            entry = {
                "id": piece.id,
                "name": piece.name,
                "role": piece.role,
                "note": piece.note,
                "color": piece.color or (WHITE if index == 0 else ORANGE),
                "motion": [0, 0, 0],
                **mesh_data(piece.shape),
            }
            if piece.role != "printed" and piece.color is None:
                entry["color"] = METAL
            entry["keyframes"] = [piece.motion(i / 32).data(i / 32) for i in range(33)]
            if piece.deform:
                entry["geometryFrames"] = []
                for i in range(33):
                    deformed = piece.deform(i / 32)
                    assert deformed.isValid() and len(deformed.Solids) == 1, (
                        piece.id,
                        i,
                    )
                    entry["geometryFrames"].append({"t": i / 32, **mesh_data(deformed)})
            payload.append(entry)
        if mode == "print-in-place":
            assembly = Part.Compound(
                [p.shape for p in parts if p.role == "printed"]
            )
            assembly.translate(V(0, 0, -assembly.BoundBox.ZMin))
            mesh = MeshPart.meshFromShape(
                Shape=assembly, LinearDeflection=0.1, AngularDeflection=0.18
            )
            assert mesh.isSolid()
            mesh.write(str(work / "assembly.stl"))
        doc.recompute()
        doc.saveAs(str(work / "demo.FCStd"))
        Part.export(objects, str(work / "demo.step"))
        restored = Part.read(str(work / "demo.step"))
        assert restored.isValid() and len(restored.Solids) == len(parts)
        expected_volume = sum(p.at(1).Volume for p in parts)
        delta = abs(restored.Volume - expected_volume)
        checks["step_volume_relative_error"] = delta / expected_volume
        assert delta < max(1e-3, expected_volume * 1e-4), (
            restored.Volume,
            expected_volume,
            delta,
        )
        App.closeDocument(doc.Name)
        data = {
            "id": model_id,
            "units": "mm",
            "parameters": values,
            "parts": payload,
            "note": note,
            "mode": mode,
            "travel": 0,
            "initialProgress": 0.45,
        }
        (work / "model.js").write_text(
            "globalThis.CADModel = " + json.dumps(data, ensure_ascii=False) + ";\n"
        )
        (work / "checks.json").write_text(
            json.dumps(checks, ensure_ascii=False, indent=2) + "\n"
        )
        (work / "parameters.json").write_text(json.dumps(values, indent=2) + "\n")
        manifest = {
            "id": model_id,
            "parameters": values,
            "note": note,
            "mode": mode,
            "parts": [
                {"id": p.id, "name": p.name, "role": p.role, "note": p.note}
                for p in parts
            ],
        }
        (work / "manifest.json").write_text(
            json.dumps(manifest, ensure_ascii=False, indent=2) + "\n"
        )
        for file in sorted(work.iterdir(), key=lambda f: f.name == "model.js"):
            file.replace(folder / file.name)
    return checks


def main(build, defaults, model_id, location, note, mode="rigid"):
    parser = argparse.ArgumentParser(description=note)
    parser.add_argument("--params", type=Path)
    parser.add_argument(
        "--output", type=Path, default=Path(location).resolve().parent / "generated"
    )
    args = parser.parse_args()
    values = params(
        defaults, json.loads(args.params.read_text()) if args.params else {}
    )
    result = export(model_id, args.output, values, build(values), note, mode)
    print(json.dumps(result, ensure_ascii=False), flush=True)