#include <pybind11/pybind11.h>
#include <pybind11/eigen.h>
#include <pybind11/stl.h>

#include "robot_dynamics.hpp"

namespace py = pybind11;
using reforge::dynamics::RobotDynamics;

/*
Summary:
  Register Python bindings for the native URDF-based robot dynamics model.
Args:
  m: Pybind11 module object provided by `PYBIND11_MODULE`.
Returns:
  None.
Side Effects:
  Defines Python-visible `RobotDynamics` kinematics and dynamics methods.
Raises:
  None.
Preconditions:
  Native dynamics library and pybind11 runtime are correctly linked.
*/
PYBIND11_MODULE(_robot_dynamics_native, m) {
    m.doc() = "URDF-native kinematics model (non-DH)";

    py::class_<RobotDynamics>(m, "RobotDynamics")
        .def(py::init<std::vector<Eigen::Vector3d>,
                      std::vector<Eigen::Vector3d>,
                      std::vector<Eigen::Vector3d>,
                      std::vector<double>,
                      std::vector<Eigen::Vector3d>,
                      std::vector<Eigen::Matrix3d>,
                      const Eigen::Matrix4d&,
                      const Eigen::Matrix4d&>(),
             py::arg("origins"),
             py::arg("rpy"),
             py::arg("axes"),
             py::arg("masses"),
             py::arg("coms"),
             py::arg("inertias"),
             py::arg("base_prefix"),
             py::arg("tool_offset"))
        .def("tcp_transform", &RobotDynamics::tcpTransform, py::arg("joint_angles"))
        .def("tcp_rotation", &RobotDynamics::tcpRotation, py::arg("joint_angles"))
        .def("tcp_quaternion", &RobotDynamics::tcpQuaternion, py::arg("joint_angles"))
        .def("individual_transforms", &RobotDynamics::individualTransforms, py::arg("joint_angles"))
        .def("individual_rotations", &RobotDynamics::individualRotations, py::arg("joint_angles"))
        .def("jacobian_matrix", &RobotDynamics::jacobianMatrix, py::arg("joint_angles"))
        .def("mass_matrix", &RobotDynamics::massMatrix, py::arg("joint_angles"))
        .def("gravity_torque", &RobotDynamics::gravityTorque, py::arg("joint_angles"), py::arg("g_world"))
        .def_property_readonly("dof", &RobotDynamics::dof);
}
