create/src/Mod/Assembly/App/AssemblyObjectPyImp.cpp

// SPDX-License-Identifier: LGPL-2.1-or-later
/****************************************************************************
 *                                                                          *
 *   Copyright (c) 2024 Ondsel <development@ondsel.com>                     *
 *                                                                          *
 *   This file is part of FreeCAD.                                          *
 *                                                                          *
 *   FreeCAD is free software: you can redistribute it and/or modify it     *
 *   under the terms of the GNU Lesser General Public License as            *
 *   published by the Free Software Foundation, either version 2.1 of the   *
 *   License, or (at your option) any later version.                        *
 *                                                                          *
 *   FreeCAD is distributed in the hope that it will be useful, but         *
 *   WITHOUT ANY WARRANTY; without even the implied warranty of             *
 *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU       *
 *   Lesser General Public License for more details.                        *
 *                                                                          *
 *   You should have received a copy of the GNU Lesser General Public       *
 *   License along with FreeCAD. If not, see                                *
 *   <https://www.gnu.org/licenses/>.                                       *
 *                                                                          *
 ***************************************************************************/

 
 // inclusion of the generated files (generated out of AssemblyObject.xml)
#include "AssemblyObjectPy.h"
#include "AssemblyObjectPy.cpp"

#include <Mod/Assembly/Solver/SolverRegistry.h>

using namespace Assembly;

namespace
{

// ── Enum-to-string tables for dict serialization ───────────────────
// String values must match kcsolve_py.cpp py::enum_ .value() names exactly.

const char* baseJointKindStr(KCSolve::BaseJointKind k)
{
    switch (k) {
        case KCSolve::BaseJointKind::Coincident:         return "Coincident";
        case KCSolve::BaseJointKind::PointOnLine:        return "PointOnLine";
        case KCSolve::BaseJointKind::PointInPlane:       return "PointInPlane";
        case KCSolve::BaseJointKind::Concentric:         return "Concentric";
        case KCSolve::BaseJointKind::Tangent:            return "Tangent";
        case KCSolve::BaseJointKind::Planar:             return "Planar";
        case KCSolve::BaseJointKind::LineInPlane:        return "LineInPlane";
        case KCSolve::BaseJointKind::Parallel:           return "Parallel";
        case KCSolve::BaseJointKind::Perpendicular:      return "Perpendicular";
        case KCSolve::BaseJointKind::Angle:              return "Angle";
        case KCSolve::BaseJointKind::Fixed:              return "Fixed";
        case KCSolve::BaseJointKind::Revolute:           return "Revolute";
        case KCSolve::BaseJointKind::Cylindrical:        return "Cylindrical";
        case KCSolve::BaseJointKind::Slider:             return "Slider";
        case KCSolve::BaseJointKind::Ball:               return "Ball";
        case KCSolve::BaseJointKind::Screw:              return "Screw";
        case KCSolve::BaseJointKind::Universal:          return "Universal";
        case KCSolve::BaseJointKind::Gear:               return "Gear";
        case KCSolve::BaseJointKind::RackPinion:         return "RackPinion";
        case KCSolve::BaseJointKind::Cam:                return "Cam";
        case KCSolve::BaseJointKind::Slot:               return "Slot";
        case KCSolve::BaseJointKind::DistancePointPoint: return "DistancePointPoint";
        case KCSolve::BaseJointKind::DistanceCylSph:     return "DistanceCylSph";
        case KCSolve::BaseJointKind::Custom:             return "Custom";
    }
    return "Custom";
}

const char* limitKindStr(KCSolve::Constraint::Limit::Kind k)
{
    switch (k) {
        case KCSolve::Constraint::Limit::Kind::TranslationMin: return "TranslationMin";
        case KCSolve::Constraint::Limit::Kind::TranslationMax: return "TranslationMax";
        case KCSolve::Constraint::Limit::Kind::RotationMin:    return "RotationMin";
        case KCSolve::Constraint::Limit::Kind::RotationMax:    return "RotationMax";
    }
    return "TranslationMin";
}

const char* motionKindStr(KCSolve::MotionDef::Kind k)
{
    switch (k) {
        case KCSolve::MotionDef::Kind::Rotational:    return "Rotational";
        case KCSolve::MotionDef::Kind::Translational: return "Translational";
        case KCSolve::MotionDef::Kind::General:       return "General";
    }
    return "Rotational";
}

// ── Python dict builders ───────────────────────────────────────────
// Layout matches solve_context_to_dict() in kcsolve_py.cpp exactly.

Py::Dict transformToDict(const KCSolve::Transform& t)
{
    Py::Dict d;
    d.setItem("position", Py::TupleN(
        Py::Float(t.position[0]),
        Py::Float(t.position[1]),
        Py::Float(t.position[2])));
    d.setItem("quaternion", Py::TupleN(
        Py::Float(t.quaternion[0]),
        Py::Float(t.quaternion[1]),
        Py::Float(t.quaternion[2]),
        Py::Float(t.quaternion[3])));
    return d;
}

Py::Dict partToDict(const KCSolve::Part& p)
{
    Py::Dict d;
    d.setItem("id", Py::String(p.id));
    d.setItem("placement", transformToDict(p.placement));
    d.setItem("mass", Py::Float(p.mass));
    d.setItem("grounded", Py::Boolean(p.grounded));
    return d;
}

Py::Dict limitToDict(const KCSolve::Constraint::Limit& lim)
{
    Py::Dict d;
    d.setItem("kind", Py::String(limitKindStr(lim.kind)));
    d.setItem("value", Py::Float(lim.value));
    d.setItem("tolerance", Py::Float(lim.tolerance));
    return d;
}

Py::Dict constraintToDict(const KCSolve::Constraint& c)
{
    Py::Dict d;
    d.setItem("id", Py::String(c.id));
    d.setItem("part_i", Py::String(c.part_i));
    d.setItem("marker_i", transformToDict(c.marker_i));
    d.setItem("part_j", Py::String(c.part_j));
    d.setItem("marker_j", transformToDict(c.marker_j));
    d.setItem("type", Py::String(baseJointKindStr(c.type)));

    Py::List params;
    for (double v : c.params) {
        params.append(Py::Float(v));
    }
    d.setItem("params", params);

    Py::List lims;
    for (const auto& l : c.limits) {
        lims.append(limitToDict(l));
    }
    d.setItem("limits", lims);
    d.setItem("activated", Py::Boolean(c.activated));
    return d;
}

Py::Dict motionToDict(const KCSolve::MotionDef& m)
{
    Py::Dict d;
    d.setItem("kind", Py::String(motionKindStr(m.kind)));
    d.setItem("joint_id", Py::String(m.joint_id));
    d.setItem("marker_i", Py::String(m.marker_i));
    d.setItem("marker_j", Py::String(m.marker_j));
    d.setItem("rotation_expr", Py::String(m.rotation_expr));
    d.setItem("translation_expr", Py::String(m.translation_expr));
    return d;
}

Py::Dict simToDict(const KCSolve::SimulationParams& s)
{
    Py::Dict d;
    d.setItem("t_start", Py::Float(s.t_start));
    d.setItem("t_end", Py::Float(s.t_end));
    d.setItem("h_out", Py::Float(s.h_out));
    d.setItem("h_min", Py::Float(s.h_min));
    d.setItem("h_max", Py::Float(s.h_max));
    d.setItem("error_tol", Py::Float(s.error_tol));
    return d;
}

}  // anonymous namespace

// returns a string which represents the object e.g. when printed in python
std::string AssemblyObjectPy::representation() const
{
    return {"<Assembly object>"};
}

PyObject* AssemblyObjectPy::getCustomAttributes(const char* /*attr*/) const
{
    return nullptr;
}

int AssemblyObjectPy::setCustomAttributes(const char* /*attr*/, PyObject* /*obj*/)
{
    return 0;
}

PyObject* AssemblyObjectPy::solve(PyObject* args) const
{
    PyObject* enableUndoPy;
    bool enableUndo;

    if (!PyArg_ParseTuple(args, "O!", &PyBool_Type, &enableUndoPy)) {
        PyErr_Clear();
        if (!PyArg_ParseTuple(args, "")) {
            return nullptr;
        }
        else {
            enableUndo = false;
        }
    }
    else {
        enableUndo = Base::asBoolean(enableUndoPy);
    }

    int ret = this->getAssemblyObjectPtr()->solve(enableUndo);
    return Py_BuildValue("i", ret);
}

PyObject* AssemblyObjectPy::generateSimulation(PyObject* args) const
{
    PyObject* pyobj;

    if (!PyArg_ParseTuple(args, "O", &pyobj)) {
        return nullptr;
    }
    auto* obj = static_cast<App::DocumentObjectPy*>(pyobj)->getDocumentObjectPtr();
    int ret = this->getAssemblyObjectPtr()->generateSimulation(obj);
    return Py_BuildValue("i", ret);
}

PyObject* AssemblyObjectPy::ensureIdentityPlacements(PyObject* args) const
{
    if (!PyArg_ParseTuple(args, "")) {
        return nullptr;
    }
    this->getAssemblyObjectPtr()->ensureIdentityPlacements();
    Py_Return;
}

PyObject* AssemblyObjectPy::updateForFrame(PyObject* args) const
{
    unsigned long index {};

    if (!PyArg_ParseTuple(args, "k", &index)) {
        throw Py::RuntimeError("updateForFrame requires an integer index");
    }
    PY_TRY
    {
        this->getAssemblyObjectPtr()->updateForFrame(index);
    }
    PY_CATCH;

    Py_Return;
}

PyObject* AssemblyObjectPy::numberOfFrames(PyObject* args) const
{
    if (!PyArg_ParseTuple(args, "")) {
        return nullptr;
    }
    size_t ret = this->getAssemblyObjectPtr()->numberOfFrames();
    return Py_BuildValue("k", ret);
}

PyObject* AssemblyObjectPy::updateSolveStatus(PyObject* args) const
{
    if (!PyArg_ParseTuple(args, "")) {
        return nullptr;
    }

    this->getAssemblyObjectPtr()->updateSolveStatus();
    Py_Return;
}

PyObject* AssemblyObjectPy::undoSolve(PyObject* args) const
{
    if (!PyArg_ParseTuple(args, "")) {
        return nullptr;
    }
    this->getAssemblyObjectPtr()->undoSolve();
    Py_Return;
}

PyObject* AssemblyObjectPy::clearUndo(PyObject* args) const
{
    if (!PyArg_ParseTuple(args, "")) {
        return nullptr;
    }
    this->getAssemblyObjectPtr()->clearUndo();
    Py_Return;
}

PyObject* AssemblyObjectPy::isPartConnected(PyObject* args) const
{
    PyObject* pyobj;

    if (!PyArg_ParseTuple(args, "O!", &(App::DocumentObjectPy::Type), &pyobj)) {
        return nullptr;
    }
    auto* obj = static_cast<App::DocumentObjectPy*>(pyobj)->getDocumentObjectPtr();
    bool ok = this->getAssemblyObjectPtr()->isPartConnected(obj);
    return Py_BuildValue("O", (ok ? Py_True : Py_False));
}

PyObject* AssemblyObjectPy::isPartGrounded(PyObject* args) const
{
    PyObject* pyobj;

    if (!PyArg_ParseTuple(args, "O!", &(App::DocumentObjectPy::Type), &pyobj)) {
        return nullptr;
    }
    auto* obj = static_cast<App::DocumentObjectPy*>(pyobj)->getDocumentObjectPtr();
    bool ok = this->getAssemblyObjectPtr()->isPartGrounded(obj);
    return Py_BuildValue("O", (ok ? Py_True : Py_False));
}

PyObject* AssemblyObjectPy::isJointConnectingPartToGround(PyObject* args) const
{
    PyObject* pyobj;
    char* pname;

    if (!PyArg_ParseTuple(args, "O!s", &(App::DocumentObjectPy::Type), &pyobj, &pname)) {
        return nullptr;
    }
    auto* obj = static_cast<App::DocumentObjectPy*>(pyobj)->getDocumentObjectPtr();
    bool ok = this->getAssemblyObjectPtr()->isJointConnectingPartToGround(obj, pname);
    return Py_BuildValue("O", (ok ? Py_True : Py_False));
}

PyObject* AssemblyObjectPy::exportAsASMT(PyObject* args) const
{
    char* utf8Name;
    if (!PyArg_ParseTuple(args, "et", "utf-8", &utf8Name)) {
        return nullptr;
    }

    std::string fileName = utf8Name;
    PyMem_Free(utf8Name);

    if (fileName.empty()) {
        PyErr_SetString(PyExc_ValueError, "Passed string is empty");
        return nullptr;
    }

    this->getAssemblyObjectPtr()->exportAsASMT(fileName);

    Py_Return;
}

Py::List AssemblyObjectPy::getJoints() const
{
    Py::List ret;
    std::vector<App::DocumentObject*> list = getAssemblyObjectPtr()->getJoints(false);

    for (auto It : list) {
        ret.append(Py::Object(It->getPyObject(), true));
    }

    return ret;
}

PyObject* AssemblyObjectPy::getDownstreamParts(PyObject* args) const
{
    PyObject* pyPart;
    PyObject* pyJoint;

    // Parse the two arguments: a part object and a joint object
    if (!PyArg_ParseTuple(
            args,
            "O!O!",
            &(App::DocumentObjectPy::Type),
            &pyPart,
            &(App::DocumentObjectPy::Type),
            &pyJoint
        )) {
        return nullptr;
    }

    auto* part = static_cast<App::DocumentObjectPy*>(pyPart)->getDocumentObjectPtr();
    auto* joint = static_cast<App::DocumentObjectPy*>(pyJoint)->getDocumentObjectPtr();

    // Call the C++ method
    std::vector<Assembly::ObjRef> downstreamParts
        = this->getAssemblyObjectPtr()->getDownstreamParts(part, joint);

    // Convert the result into a Python list of DocumentObjects
    Py::List ret;
    for (const auto& objRef : downstreamParts) {
        if (objRef.obj) {
            ret.append(Py::Object(objRef.obj->getPyObject(), true));
        }
    }

    return Py::new_reference_to(ret);
}

PyObject* AssemblyObjectPy::getSolveContext(PyObject* args) const
{
    if (!PyArg_ParseTuple(args, "")) {
        return nullptr;
    }
    PY_TRY
    {
        KCSolve::SolveContext ctx = getAssemblyObjectPtr()->getSolveContext();

        // Empty context (no grounded parts) → return empty dict
        if (ctx.parts.empty()) {
            return Py::new_reference_to(Py::Dict());
        }

        Py::Dict d;
        d.setItem("api_version", Py::Long(KCSolve::API_VERSION_MAJOR));

        Py::List parts;
        for (const auto& p : ctx.parts) {
            parts.append(partToDict(p));
        }
        d.setItem("parts", parts);

        Py::List constraints;
        for (const auto& c : ctx.constraints) {
            constraints.append(constraintToDict(c));
        }
        d.setItem("constraints", constraints);

        Py::List motions;
        for (const auto& m : ctx.motions) {
            motions.append(motionToDict(m));
        }
        d.setItem("motions", motions);

        if (ctx.simulation.has_value()) {
            d.setItem("simulation", simToDict(*ctx.simulation));
        }
        else {
            d.setItem("simulation", Py::None());
        }

        d.setItem("bundle_fixed", Py::Boolean(ctx.bundle_fixed));

        return Py::new_reference_to(d);
    }
    PY_CATCH;
}