#pragma once

#include <Eigen/Dense>
#include <tuple>
#include <vector>

namespace reforge {
namespace dynamics {

/**
 * URDF-native kinematics model (no DH reduction).
 *
 * Uses per-joint origin (xyz + extrinsic RPY) and joint axis in parent frame.
 * Applies a base_prefix transform up front and a tool_offset at the end.
 */
class RobotDynamics {
public:
    RobotDynamics(std::vector<Eigen::Vector3d> origins,
                     std::vector<Eigen::Vector3d> rpy,
                     std::vector<Eigen::Vector3d> axes,
                     std::vector<double> masses,
                     std::vector<Eigen::Vector3d> coms,
                     std::vector<Eigen::Matrix3d> inertias,
                     const Eigen::Matrix4d& base_prefix,
                     const Eigen::Matrix4d& tool_offset);

    Eigen::Matrix4d tcpTransform(const Eigen::VectorXd& joint_angles) const;
    Eigen::Matrix3d tcpRotation(const Eigen::VectorXd& joint_angles) const;
    Eigen::Vector4d tcpQuaternion(const Eigen::VectorXd& joint_angles) const;

    std::vector<Eigen::Matrix4d> individualTransforms(const Eigen::VectorXd& joint_angles) const;
    std::vector<Eigen::Matrix3d> individualRotations(const Eigen::VectorXd& joint_angles) const;

    Eigen::MatrixXd jacobianMatrix(const Eigen::VectorXd& joint_angles) const;
    Eigen::MatrixXd massMatrix(const Eigen::VectorXd& joint_angles) const;
    Eigen::VectorXd gravityTorque(const Eigen::VectorXd& joint_angles,
                                  const Eigen::Vector3d& g_world) const;

    std::size_t dof() const { return axes_.size(); }

private:
    Eigen::Matrix3d rotationFromRPY(const Eigen::Vector3d& rpy) const;
    Eigen::Matrix3d axisAngle(const Eigen::Vector3d& axis, double angle) const;
    Eigen::Vector4d rotationMatrixToQuaternion(const Eigen::Matrix3d& R) const;

    std::vector<Eigen::Vector3d> origins_;
    std::vector<Eigen::Vector3d> rpy_;
    std::vector<Eigen::Vector3d> axes_;
    std::vector<double> masses_;
    std::vector<Eigen::Vector3d> coms_;
    std::vector<Eigen::Matrix3d> inertias_;
    Eigen::Matrix4d base_prefix_;
    Eigen::Matrix4d tool_offset_;
};

} // namespace dynamics
} // namespace reforge
