#include "robot_dynamics.hpp"

#include <cmath>

namespace reforge {
namespace dynamics {

RobotDynamics::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)
    : origins_(std::move(origins)),
      rpy_(std::move(rpy)),
      axes_(std::move(axes)),
      masses_(std::move(masses)),
      coms_(std::move(coms)),
      inertias_(std::move(inertias)),
      base_prefix_(base_prefix),
      tool_offset_(tool_offset) {
    const auto n = origins_.size();
    if (n != rpy_.size() || n != axes_.size()) {
        throw std::invalid_argument("RobotDynamics: origins, rpy, and axes must be the same length");
    }
    if (masses_.size() != n || coms_.size() != n || inertias_.size() != n) {
        throw std::invalid_argument("RobotDynamics: inertia data must match number of joints");
    }
}

/*
Summary:
  Convert roll-pitch-yaw angles to a rotation matrix using `Rz * Ry * Rx`.
Args:
  rpy: Euler angles `[roll, pitch, yaw]` [rad].
Returns:
  `Eigen::Matrix3d` rotation matrix.
Side Effects:
  None.
Raises:
  None.
Preconditions:
  Angles are finite.
*/
Eigen::Matrix3d RobotDynamics::rotationFromRPY(const Eigen::Vector3d& rpy) const {
    // Extrinsic roll (x), then pitch (y), then yaw (z): Rz * Ry * Rx
    const double roll = rpy[0];
    const double pitch = rpy[1];
    const double yaw = rpy[2];

    const double cr = std::cos(roll);
    const double sr = std::sin(roll);
    const double cp = std::cos(pitch);
    const double sp = std::sin(pitch);
    const double cy = std::cos(yaw);
    const double sy = std::sin(yaw);

    Eigen::Matrix3d Rx;
    Rx << 1, 0, 0,
          0, cr, -sr,
          0, sr, cr;

    Eigen::Matrix3d Ry;
    Ry << cp, 0, sp,
          0, 1, 0,
          -sp, 0, cp;

    Eigen::Matrix3d Rz;
    Rz << cy, -sy, 0,
          sy,  cy, 0,
          0,    0,  1;

    return Rz * Ry * Rx;
}

/*
Summary:
  Convert axis-angle representation to a rotation matrix.
Args:
  axis: Rotation axis.
  angle: Rotation angle [rad].
Returns:
  `Eigen::Matrix3d` rotation matrix.
Side Effects:
  None.
Raises:
  None.
Preconditions:
  `axis` is non-zero.
*/
Eigen::Matrix3d RobotDynamics::axisAngle(const Eigen::Vector3d& axis, double angle) const {
    Eigen::Vector3d a = axis.normalized();
    const double x = a[0], y = a[1], z = a[2];
    const double c = std::cos(angle);
    const double s = std::sin(angle);
    Eigen::Matrix3d R;
    R << c + x * x * (1 - c),     x * y * (1 - c) - z * s, x * z * (1 - c) + y * s,
         y * x * (1 - c) + z * s, c + y * y * (1 - c),     y * z * (1 - c) - x * s,
         z * x * (1 - c) - y * s, z * y * (1 - c) + x * s, c + z * z * (1 - c);
    return R;
}

/*
Summary:
  Convert a rotation matrix to quaternion in `[x, y, z, w]` order.
Args:
  R: Rotation matrix.
Returns:
  `Eigen::Vector4d` quaternion `[x, y, z, w]`.
Side Effects:
  None.
Raises:
  None.
Preconditions:
  `R` is approximately orthonormal.
*/
Eigen::Vector4d RobotDynamics::rotationMatrixToQuaternion(const Eigen::Matrix3d& R) const {
    double trace = R.trace();
    double qw, qx, qy, qz;
    if (trace > 0.0) {
        double S = std::sqrt(trace + 1.0) * 2.0;
        qw = 0.25 * S;
        qx = (R(2, 1) - R(1, 2)) / S;
        qy = (R(0, 2) - R(2, 0)) / S;
        qz = (R(1, 0) - R(0, 1)) / S;
    } else if ((R(0, 0) > R(1, 1)) && (R(0, 0) > R(2, 2))) {
        double S = std::sqrt(1.0 + R(0, 0) - R(1, 1) - R(2, 2)) * 2.0;
        qw = (R(2, 1) - R(1, 2)) / S;
        qx = 0.25 * S;
        qy = (R(0, 1) + R(1, 0)) / S;
        qz = (R(0, 2) + R(2, 0)) / S;
    } else if (R(1, 1) > R(2, 2)) {
        double S = std::sqrt(1.0 + R(1, 1) - R(0, 0) - R(2, 2)) * 2.0;
        qw = (R(0, 2) - R(2, 0)) / S;
        qx = (R(0, 1) + R(1, 0)) / S;
        qy = 0.25 * S;
        qz = (R(1, 2) + R(2, 1)) / S;
    } else {
        double S = std::sqrt(1.0 + R(2, 2) - R(0, 0) - R(1, 1)) * 2.0;
        qw = (R(1, 0) - R(0, 1)) / S;
        qx = (R(0, 2) + R(2, 0)) / S;
        qy = (R(1, 2) + R(2, 1)) / S;
        qz = 0.25 * S;
    }
    Eigen::Vector4d q;
    q << qx, qy, qz, qw;
    return q;
}

/*
Summary:
  Compute TCP homogeneous transform for a joint configuration.
Args:
  joint_angles: Joint angle vector [rad].
Returns:
  `Eigen::Matrix4d` TCP transform in world frame.
Side Effects:
  None.
Raises:
  std::invalid_argument: If joint vector length mismatches model DOF.
Preconditions:
  Model arrays (`origins_`, `rpy_`, `axes_`) are dimensionally consistent.
*/
Eigen::Matrix4d RobotDynamics::tcpTransform(const Eigen::VectorXd& joint_angles) const {
    const int dof = static_cast<int>(axes_.size());
    if (joint_angles.size() != dof) {
        throw std::invalid_argument("RobotDynamics: joint angle size mismatch");
    }

    Eigen::Matrix4d T = base_prefix_;
    // Each iteration applies the fixed joint origin transform, then the
    // revolute/prismatic joint motion transform about the declared axis.
    for (int i = 0; i < dof; ++i) {
        Eigen::Matrix4d T_origin = Eigen::Matrix4d::Identity();
        T_origin.block<3, 3>(0, 0) = rotationFromRPY(rpy_[static_cast<std::size_t>(i)]);
        T_origin.block<3, 1>(0, 3) = origins_[static_cast<std::size_t>(i)];

        Eigen::Matrix4d T_joint = Eigen::Matrix4d::Identity();
        T_joint.block<3, 3>(0, 0) = axisAngle(axes_[static_cast<std::size_t>(i)], joint_angles[i]);

        T = T * T_origin * T_joint;
    }
    return T * tool_offset_;
}

/*
Summary:
  Compute TCP rotation matrix for a joint configuration.
Args:
  joint_angles: Joint angle vector [rad].
Returns:
  `Eigen::Matrix3d` TCP rotation.
Side Effects:
  None.
Raises:
  std::invalid_argument: Propagated from `tcpTransform` on size mismatch.
Preconditions:
  `joint_angles` length equals model DOF.
*/
Eigen::Matrix3d RobotDynamics::tcpRotation(const Eigen::VectorXd& joint_angles) const {
    return tcpTransform(joint_angles).block<3, 3>(0, 0);
}

/*
Summary:
  Compute TCP quaternion for a joint configuration.
Args:
  joint_angles: Joint angle vector [rad].
Returns:
  `Eigen::Vector4d` quaternion `[x, y, z, w]`.
Side Effects:
  None.
Raises:
  std::invalid_argument: Propagated from transform/rotation helpers.
Preconditions:
  `joint_angles` length equals model DOF.
*/
Eigen::Vector4d RobotDynamics::tcpQuaternion(const Eigen::VectorXd& joint_angles) const {
    return rotationMatrixToQuaternion(tcpRotation(joint_angles));
}

/*
Summary:
  Compute cumulative transform after each joint in chain order.
Args:
  joint_angles: Joint angle vector [rad].
Returns:
  `std::vector<Eigen::Matrix4d>` per-joint world transforms.
Side Effects:
  None.
Raises:
  std::invalid_argument: If joint vector length mismatches model DOF.
Preconditions:
  `joint_angles` length equals model DOF.
*/
std::vector<Eigen::Matrix4d> RobotDynamics::individualTransforms(const Eigen::VectorXd& joint_angles) const {
    const int dof = static_cast<int>(axes_.size());
    if (joint_angles.size() != dof) {
        throw std::invalid_argument("RobotDynamics: joint angle size mismatch");
    }
    std::vector<Eigen::Matrix4d> transforms(static_cast<std::size_t>(dof));

    Eigen::Matrix4d T = base_prefix_;
    // Store cumulative transform after each joint so Jacobian and diagnostics
    // can reference per-joint world frames directly.
    for (int i = 0; i < dof; ++i) {
        Eigen::Matrix4d T_origin = Eigen::Matrix4d::Identity();
        T_origin.block<3, 3>(0, 0) = rotationFromRPY(rpy_[static_cast<std::size_t>(i)]);
        T_origin.block<3, 1>(0, 3) = origins_[static_cast<std::size_t>(i)];

        Eigen::Matrix4d T_joint = Eigen::Matrix4d::Identity();
        T_joint.block<3, 3>(0, 0) = axisAngle(axes_[static_cast<std::size_t>(i)], joint_angles[i]);

        T = T * T_origin * T_joint;
        transforms[static_cast<std::size_t>(i)] = T;
    }

    // Append tool_offset to the last transform to match tcpTransform
    transforms.back() = transforms.back() * tool_offset_;
    return transforms;
}

/*
Summary:
  Compute per-joint rotation matrices from cumulative transforms.
Args:
  joint_angles: Joint angle vector [rad].
Returns:
  `std::vector<Eigen::Matrix3d>` per-joint rotations.
Side Effects:
  None.
Raises:
  std::invalid_argument: Propagated from `individualTransforms`.
Preconditions:
  `joint_angles` length equals model DOF.
*/
std::vector<Eigen::Matrix3d> RobotDynamics::individualRotations(const Eigen::VectorXd& joint_angles) const {
    auto transforms = individualTransforms(joint_angles);
    std::vector<Eigen::Matrix3d> rotations;
    rotations.reserve(transforms.size());
    for (const auto& T : transforms) {
        rotations.push_back(T.block<3, 3>(0, 0));
    }
    return rotations;
}

/*
Summary:
  Compute 6xN geometric Jacobian at TCP.
Args:
  joint_angles: Joint angle vector [rad].
Returns:
  `Eigen::MatrixXd` Jacobian with linear rows `[0:3]` and angular rows `[3:6]`.
Side Effects:
  None.
Raises:
  std::invalid_argument: If joint vector length mismatches model DOF.
Preconditions:
  Kinematic chain description is valid.
*/
Eigen::MatrixXd RobotDynamics::jacobianMatrix(const Eigen::VectorXd& joint_angles) const {
    const int dof = static_cast<int>(axes_.size());
    if (joint_angles.size() != dof) {
        throw std::invalid_argument("RobotDynamics: joint angle size mismatch");
    }

    Eigen::MatrixXd J(6, dof);
    std::vector<Eigen::Vector3d> origins_world(static_cast<std::size_t>(dof));
    std::vector<Eigen::Vector3d> axes_world(static_cast<std::size_t>(dof));

    Eigen::Matrix4d T = base_prefix_;
    // First pass: world-space joint origins/axes at current configuration.
    for (int i = 0; i < dof; ++i) {
        Eigen::Matrix4d T_origin = Eigen::Matrix4d::Identity();
        T_origin.block<3, 3>(0, 0) = rotationFromRPY(rpy_[static_cast<std::size_t>(i)]);
        T_origin.block<3, 1>(0, 3) = origins_[static_cast<std::size_t>(i)];

        Eigen::Matrix4d T_joint = Eigen::Matrix4d::Identity();
        T_joint.block<3, 3>(0, 0) = axisAngle(axes_[static_cast<std::size_t>(i)], joint_angles[i]);

        // Express joint origin/axis after applying the fixed URDF origin transform.
        Eigen::Matrix4d T_before_joint = T * T_origin;
        origins_world[static_cast<std::size_t>(i)] = T_before_joint.block<3, 1>(0, 3);
        axes_world[static_cast<std::size_t>(i)] =
            T_before_joint.block<3, 3>(0, 0) * axes_[static_cast<std::size_t>(i)];

        T = T_before_joint * T_joint;
    }

    // TCP position including tool offset
    Eigen::Vector3d p_tcp = (T * tool_offset_).block<3, 1>(0, 3);

    // Second pass: assemble geometric Jacobian columns [Jv; Jw].
    for (int i = 0; i < dof; ++i) {
        const Eigen::Vector3d& z = axes_world[static_cast<std::size_t>(i)];
        const Eigen::Vector3d& o = origins_world[static_cast<std::size_t>(i)];
        Eigen::Vector3d Jv = z.cross(p_tcp - o);
        Eigen::Vector3d Jw = z;
        J.block<3,1>(0,i) = Jv;
        J.block<3,1>(3,i) = Jw;
    }
    return J;
}

/*
Summary:
  Compute joint-space mass matrix from link masses/inertias and Jacobians.
Args:
  joint_angles: Joint angle vector [rad].
Returns:
  `Eigen::MatrixXd` mass matrix `[N x N]`.
Side Effects:
  None.
Raises:
  std::invalid_argument: If joint vector length mismatches model DOF.
Preconditions:
  Mass, COM, and inertia data are defined per link.
*/
Eigen::MatrixXd RobotDynamics::massMatrix(const Eigen::VectorXd& joint_angles) const {
    const int dof = static_cast<int>(axes_.size());
    if (joint_angles.size() != dof) {
        throw std::invalid_argument("RobotDynamics: joint angle size mismatch");
    }

    Eigen::MatrixXd M = Eigen::MatrixXd::Zero(dof, dof);
    std::vector<Eigen::Vector3d> origins_world(static_cast<std::size_t>(dof));
    std::vector<Eigen::Vector3d> axes_world(static_cast<std::size_t>(dof));
    std::vector<Eigen::Vector3d> com_positions(static_cast<std::size_t>(dof));
    std::vector<Eigen::Matrix3d> rotation_stack(static_cast<std::size_t>(dof));

    Eigen::Matrix4d T = base_prefix_;
    // Build link COM positions and world-space joint axes needed by
    // the composite rigid-body accumulation below.
    for (int i = 0; i < dof; ++i) {
        Eigen::Matrix4d T_origin = Eigen::Matrix4d::Identity();
        T_origin.block<3, 3>(0, 0) = rotationFromRPY(rpy_[static_cast<std::size_t>(i)]);
        T_origin.block<3, 1>(0, 3) = origins_[static_cast<std::size_t>(i)];

        Eigen::Matrix4d T_joint = Eigen::Matrix4d::Identity();
        T_joint.block<3, 3>(0, 0) = axisAngle(axes_[static_cast<std::size_t>(i)], joint_angles[i]);

        Eigen::Matrix4d T_before_joint = T * T_origin;
        origins_world[static_cast<std::size_t>(i)] = T_before_joint.block<3, 1>(0, 3);
        axes_world[static_cast<std::size_t>(i)] =
            T_before_joint.block<3, 3>(0, 0) * axes_[static_cast<std::size_t>(i)];

        T = T_before_joint * T_joint;
        rotation_stack[static_cast<std::size_t>(i)] = T.block<3, 3>(0, 0);
        com_positions[static_cast<std::size_t>(i)] =
            T.block<3, 1>(0, 3) + rotation_stack[static_cast<std::size_t>(i)] * coms_[static_cast<std::size_t>(i)];
    }

    // Sum each link's translational and rotational inertia contribution:
    // M += Jv^T m Jv + Jw^T R I_body R^T Jw.
    for (int link = 0; link < dof; ++link) {
        Eigen::MatrixXd Jv = Eigen::MatrixXd::Zero(3, dof);
        Eigen::MatrixXd Jw = Eigen::MatrixXd::Zero(3, dof);

        for (int joint = 0; joint <= link; ++joint) {
            const Eigen::Vector3d& axis = axes_world[static_cast<std::size_t>(joint)];
            const Eigen::Vector3d& origin = origins_world[static_cast<std::size_t>(joint)];
            Jw.col(joint) = axis;
            Jv.col(joint) = axis.cross(com_positions[static_cast<std::size_t>(link)] - origin);
        }

        const double mass = masses_[static_cast<std::size_t>(link)];
        const Eigen::Matrix3d& inertia_local = inertias_[static_cast<std::size_t>(link)];
        const Eigen::Matrix3d& R = rotation_stack[static_cast<std::size_t>(link)];

        M += Jv.transpose() * (mass * Eigen::Matrix3d::Identity()) * Jv +
             Jw.transpose() * R * inertia_local * R.transpose() * Jw;
    }
    return M;
}

Eigen::VectorXd RobotDynamics::gravityTorque(const Eigen::VectorXd& joint_angles,
                                                const Eigen::Vector3d& g_world) const {
    const int dof = static_cast<int>(axes_.size());
    if (joint_angles.size() != dof) {
        throw std::invalid_argument("RobotDynamics: joint angle size mismatch");
    }

    Eigen::VectorXd tau = Eigen::VectorXd::Zero(dof);
    std::vector<Eigen::Vector3d> origins_world(static_cast<std::size_t>(dof));
    std::vector<Eigen::Vector3d> axes_world(static_cast<std::size_t>(dof));
    std::vector<Eigen::Vector3d> com_positions(static_cast<std::size_t>(dof));

    Eigen::Matrix4d T = base_prefix_;
    // Build world-space joint frames and COM positions for potential energy
    // gradient via tau = Jv^T * m * g.
    for (int i = 0; i < dof; ++i) {
        Eigen::Matrix4d T_origin = Eigen::Matrix4d::Identity();
        T_origin.block<3, 3>(0, 0) = rotationFromRPY(rpy_[static_cast<std::size_t>(i)]);
        T_origin.block<3, 1>(0, 3) = origins_[static_cast<std::size_t>(i)];

        Eigen::Matrix4d T_joint = Eigen::Matrix4d::Identity();
        T_joint.block<3, 3>(0, 0) = axisAngle(axes_[static_cast<std::size_t>(i)], joint_angles[i]);

        Eigen::Matrix4d T_before_joint = T * T_origin;
        origins_world[static_cast<std::size_t>(i)] = T_before_joint.block<3, 1>(0, 3);
        axes_world[static_cast<std::size_t>(i)] =
            T_before_joint.block<3, 3>(0, 0) * axes_[static_cast<std::size_t>(i)];

        T = T_before_joint * T_joint;
        const Eigen::Vector3d com_world =
            T.block<3, 1>(0, 3) + T.block<3, 3>(0, 0) * coms_[static_cast<std::size_t>(i)];
        com_positions[static_cast<std::size_t>(i)] = com_world;
    }

    for (int link = 0; link < dof; ++link) {
        Eigen::MatrixXd Jv = Eigen::MatrixXd::Zero(3, dof);
        for (int joint = 0; joint <= link; ++joint) {
            const Eigen::Vector3d& axis = axes_world[static_cast<std::size_t>(joint)];
            const Eigen::Vector3d& origin = origins_world[static_cast<std::size_t>(joint)];
            Jv.col(joint) = axis.cross(com_positions[static_cast<std::size_t>(link)] - origin);
        }
        tau += Jv.transpose() * (masses_[static_cast<std::size_t>(link)] * g_world);
    }
    return tau;
}

} // namespace dynamics
} // namespace reforge
