use std::vec;
#[cfg(not(target_arch = "wasm32"))]
use pyo3::*;
#[cfg(target_arch = "wasm32")]
use wasm_bindgen::prelude::*;
use nalgebra::DMatrix;
use serde::{Serialize, Deserialize};
use crate::robot_modules::robot_kinematics_module::{RobotKinematicsModule, RobotFKResult, FloatingLinkInput, JacobianEndPoint, JacobianMode, RobotFKResultLinkEntry};
use crate::robot_set_modules::robot_set_configuration_module::RobotSetConfigurationModule;
use crate::robot_set_modules::robot_set_joint_state_module::{RobotSetJointState, RobotSetJointStateModule, RobotSetJointStateType};
use crate::utils::utils_console::{optima_print, PrintColor, PrintMode};
use crate::utils::utils_errors::OptimaError;
use crate::utils::utils_files::optima_path::load_object_from_json_string;
use crate::utils::utils_nalgebra::conversions::NalgebraConversions;
use crate::utils::utils_se3::optima_se3_pose::{OptimaSE3Pose, OptimaSE3PoseType};
#[cfg(not(target_arch = "wasm32"))]
use crate::utils::utils_se3::optima_se3_pose::{OptimaSE3PosePy};
#[cfg(target_arch = "wasm32")]
use crate::utils::utils_se3::optima_se3_pose::{OptimaSE3PoseWASM};
use crate::utils::utils_traits::{SaveAndLoadable, ToAndFromRonString};
#[cfg(target_arch = "wasm32")]
use crate::utils::utils_wasm::JsMatrix;
#[cfg_attr(not(target_arch = "wasm32"), pyclass, derive(Clone, Debug, Serialize, Deserialize))]
#[cfg_attr(target_arch = "wasm32", wasm_bindgen, derive(Clone, Debug, Serialize, Deserialize))]
pub struct RobotSetKinematicsModule {
robot_set_joint_state_module: RobotSetJointStateModule,
robot_kinematics_modules: Vec<RobotKinematicsModule>
}
impl RobotSetKinematicsModule {
pub fn new(robot_set_configuration_module: &RobotSetConfigurationModule) -> Self {
let mut robot_fk_modules = vec![];
for r in robot_set_configuration_module.robot_configuration_modules() {
let robot_fk_module = RobotKinematicsModule::new(r.clone());
robot_fk_modules.push(robot_fk_module);
}
Self {
robot_set_joint_state_module: RobotSetJointStateModule::new(robot_set_configuration_module),
robot_kinematics_modules: robot_fk_modules
}
}
pub fn new_from_set_name(set_name: &str) -> Result<Self, OptimaError> {
let robot_set_configuration_module = RobotSetConfigurationModule::new_from_set_name(set_name)?;
return Ok(Self::new(&robot_set_configuration_module));
}
pub fn compute_fk(&self, set_joint_state: &RobotSetJointState, t: &OptimaSE3PoseType) -> Result<RobotSetFKResult, OptimaError> {
let mut out_vec = vec![];
let joint_states = &self.robot_set_joint_state_module.split_robot_set_joint_state_into_robot_joint_states(set_joint_state)?;
for (i, joint_state) in joint_states.iter().enumerate() {
let fk_res = self.robot_kinematics_modules.get(i).unwrap().compute_fk(joint_state, t)?;
out_vec.push(fk_res);
}
Ok(RobotSetFKResult {
robot_fk_results: out_vec
})
}
pub fn compute_fk_on_subset_of_robots(&self, robot_idxs: Vec<usize>, set_joint_state: &RobotSetJointState, t: &OptimaSE3PoseType) -> Result<RobotSetFKResult, OptimaError> {
let mut out_vec = vec![];
let num_robots = self.robot_kinematics_modules.len();
let joint_states = &self.robot_set_joint_state_module.split_robot_set_joint_state_into_robot_joint_states(set_joint_state)?;
for robot_idx in 0..num_robots {
if robot_idxs.contains(&robot_idx) {
let fk_res = self.robot_kinematics_modules.get(robot_idx).unwrap().compute_fk(joint_states.get(robot_idx).unwrap(), t)?;
out_vec.push(fk_res);
} else {
out_vec.push(RobotFKResult::new_empty(self.robot_kinematics_modules.get(robot_idx).unwrap()))
}
}
Ok(RobotSetFKResult {
robot_fk_results: out_vec
})
}
pub fn compute_fk_floating_chain(&self, set_joint_state: &RobotSetJointState, t: &OptimaSE3PoseType, floating_link_inputs: Vec<Option<FloatingLinkInput>>) -> Result<RobotSetFKResult, OptimaError> {
let mut out_vec = vec![];
let joint_states = &self.robot_set_joint_state_module.split_robot_set_joint_state_into_robot_joint_states(set_joint_state)?;
for (i, joint_state) in joint_states.iter().enumerate() {
let floating_link_input = floating_link_inputs.get(i).unwrap();
match floating_link_input {
None => {
out_vec.push(RobotFKResult::new_empty(self.robot_kinematics_modules.get(i).unwrap()));
}
Some(floating_link_input) => {
let fk_res = self.robot_kinematics_modules.get(i).unwrap().compute_fk_floating_chain(joint_state, t, floating_link_input)?;
out_vec.push(fk_res);
}
}
}
Ok(RobotSetFKResult {
robot_fk_results: out_vec
})
}
pub fn compute_fk_dof_perturbations(&self, joint_state: &RobotSetJointState, t: &OptimaSE3PoseType, perturbation: Option<f64>) -> Result<RobotSetFKDOFPerturbationsResult, OptimaError> {
let perturbation = match perturbation {
None => { 0.00001 }
Some(p) => { p }
};
let dof_joint_state = self.robot_set_joint_state_module.convert_state_to_dof_state(joint_state)?;
let central_fk_result = self.compute_fk(&dof_joint_state, t)?;
let mut fk_dof_perturbation_results = vec![];
let len = dof_joint_state.concatenated_state().len();
for i in 0..len {
let mut dof_joint_state_copy = dof_joint_state.clone();
dof_joint_state_copy[i] += perturbation;
let fk_res = self.compute_fk(&dof_joint_state_copy, t)?;
fk_dof_perturbation_results.push(fk_res);
}
Ok(RobotSetFKDOFPerturbationsResult {
perturbation,
central_fk_result,
fk_dof_perturbation_results
})
}
pub fn compute_jacobian(&self,
joint_state: &RobotSetJointState,
robot_idx_in_set: usize,
start_link_idx: Option<usize>,
end_link_idx: usize,
robot_jacobian_end_point: &JacobianEndPoint,
start_link_pose: Option<OptimaSE3Pose>,
jacobian_mode: JacobianMode) -> Result<DMatrix<f64>, OptimaError> {
if robot_idx_in_set >= self.robot_kinematics_modules.len() {
return Err(OptimaError::new_idx_out_of_bound_error(robot_idx_in_set, self.robot_kinematics_modules.len(), file!(), line!()));
}
let num_dofs = self.robot_set_joint_state_module.num_dofs();
let num_rows = match jacobian_mode {
JacobianMode::Full => { 6 }
JacobianMode::Translational => { 3 }
JacobianMode::Rotational => { 3 }
};
let mut jacobian = DMatrix::zeros(num_rows, num_dofs);
let mut start_idx = 0;
for i in 0..robot_idx_in_set {
start_idx += self.robot_set_joint_state_module.robot_joint_state_modules().get(i).unwrap().num_dofs();
}
let robot_states = self.robot_set_joint_state_module.split_robot_set_joint_state_into_robot_joint_states(joint_state)?;
let robot_state = robot_states.get(robot_idx_in_set).unwrap();
let robot_jacobian = self.robot_kinematics_modules.get(robot_idx_in_set).unwrap().compute_jacobian(robot_state, start_link_idx, end_link_idx, robot_jacobian_end_point, start_link_pose, jacobian_mode)?;
let num_columns = robot_state.len();
for i in 0..num_columns {
let col = i + start_idx;
for row in 0..num_rows {
jacobian[(row, col)] = robot_jacobian[(row, i)];
}
}
Ok(jacobian)
}
pub fn compute_reverse_fk(&self, input: &RobotSetFKResult) -> Result<RobotSetJointState, OptimaError> {
let mut out_state = self.robot_set_joint_state_module.spawn_zeros_robot_set_joint_state(RobotSetJointStateType::Full);
let mut idx = 0;
for (robot_idx_in_set, robot_fk_result) in input.robot_fk_results.iter().enumerate() {
let res = self.robot_kinematics_modules[robot_idx_in_set].compute_reverse_fk(robot_fk_result)?;
let len = res.len();
for i in 0..len {
out_state[idx] = res[i];
idx += 1;
}
}
Ok(out_state)
}
}
impl SaveAndLoadable for RobotSetKinematicsModule {
type SaveType = (String, String);
fn get_save_serialization_object(&self) -> Self::SaveType {
(self.robot_set_joint_state_module.get_serialization_string(), self.robot_kinematics_modules.get_serialization_string())
}
fn load_from_json_string(json_str: &str) -> Result<Self, OptimaError> where Self: Sized {
let load: Self::SaveType = load_object_from_json_string(json_str)?;
let robot_set_joint_state_module = RobotSetJointStateModule::load_from_json_string(&load.0)?;
let robot_kinematics_modules = Vec::load_from_json_string(&load.1)?;
Ok(Self {
robot_set_joint_state_module,
robot_kinematics_modules
})
}
}
#[cfg(not(target_arch = "wasm32"))]
#[pymethods]
impl RobotSetKinematicsModule {
#[new]
pub fn new_from_set_name_py(set_name: &str) -> Self {
Self::new_from_set_name(set_name).expect("error")
}
#[staticmethod]
pub fn new_py(robot_set_configuration_module: &RobotSetConfigurationModule) -> Self {
Self::new(robot_set_configuration_module)
}
#[args(pose_type = "\"ImplicitDualQuaternion\"")]
pub fn compute_fk_py(&self, joint_state: Vec<f64>, pose_type: &str) -> RobotSetFKResult {
let robot_joint_state = self.robot_set_joint_state_module.spawn_robot_set_joint_state_try_auto_type(NalgebraConversions::vec_to_dvector(&joint_state)).expect("error");
return self.compute_fk(&robot_joint_state, &OptimaSE3PoseType::from_ron_string(pose_type).expect("error")).expect("error");
}
#[args(robot_jacobian_end_point = "\"Link\"", jacobian_mode = "\"Full\"")]
pub fn compute_jacobian_py(&self, joint_state: Vec<f64>, robot_idx_in_set: usize, end_link_idx: usize, start_link_idx: Option<usize>, start_link_pose: Option<OptimaSE3PosePy>, robot_jacobian_end_point: &str, jacobian_mode: &str) -> Vec<Vec<f64>> {
let robot_joint_state = self.robot_set_joint_state_module.spawn_robot_set_joint_state_try_auto_type(NalgebraConversions::vec_to_dvector(&joint_state)).expect("error");
let start_link_pose = match start_link_pose {
None => { None }
Some(p) => { Some(p.pose().clone()) }
};
let jac = self.compute_jacobian(&robot_joint_state,
robot_idx_in_set,
start_link_idx,
end_link_idx,
&JacobianEndPoint::from_ron_string(robot_jacobian_end_point).expect("error"),
start_link_pose,
JacobianMode::from_ron_string(jacobian_mode).expect("error")).expect("error");
let jac_vecs = NalgebraConversions::dmatrix_to_vecs(&jac);
return jac_vecs;
}
pub fn compute_reverse_fk_py(&self, v: Vec<Vec<Option<OptimaSE3PosePy>>>) -> Vec<f64> {
let mut out_vec = vec![];
for (robot_idx_in_set, vec) in v.iter().enumerate() {
let res = self.robot_kinematics_modules[robot_idx_in_set].compute_reverse_fk_py(vec.clone());
for r in res { out_vec.push(r); }
}
return out_vec;
}
pub fn robot_kinematics_modules_py(&self) -> Vec<RobotKinematicsModule> {
self.robot_kinematics_modules.clone()
}
}
#[cfg(target_arch = "wasm32")]
#[wasm_bindgen]
impl RobotSetKinematicsModule {
pub fn new_from_set_name_wasm(set_name: &str) -> Self {
Self::new_from_set_name(set_name).expect("error")
}
pub fn compute_fk_wasm(&self, joint_state: Vec<f64>, pose_type: &str) -> JsValue {
let robot_joint_state = self.robot_set_joint_state_module.spawn_robot_set_joint_state_try_auto_type(NalgebraConversions::vec_to_dvector(&joint_state)).expect("error");
let res = self.compute_fk(&robot_joint_state, &OptimaSE3PoseType::from_ron_string(pose_type).expect("error")).expect("error");
return JsValue::from_serde(&res).unwrap();
}
pub fn compute_jacobian_wasm(&self, joint_state: Vec<f64>, robot_idx_in_set: usize, end_link_idx: usize, start_link_idx: Option<usize>, start_link_pose: Option<OptimaSE3PoseWASM>, robot_jacobian_end_point: &str, jacobian_mode: &str) -> JsValue {
let robot_joint_state = self.robot_set_joint_state_module.spawn_robot_set_joint_state_try_auto_type(NalgebraConversions::vec_to_dvector(&joint_state)).expect("error");
let start_link_pose = match start_link_pose {
None => { None }
Some(p) => { Some(p.pose().clone()) }
};
let jac = self.compute_jacobian(&robot_joint_state,
robot_idx_in_set,
start_link_idx,
end_link_idx,
&JacobianEndPoint::from_ron_string(robot_jacobian_end_point).expect("error"),
start_link_pose,
JacobianMode::from_ron_string(jacobian_mode).expect("error")).expect("error");
let jac_vecs = NalgebraConversions::dmatrix_to_vecs(&jac);
let jac_vecs_js = JsMatrix::new(jac_vecs);
return JsValue::from_serde(&jac_vecs_js).unwrap();
}
}
#[cfg_attr(not(target_arch = "wasm32"), pyclass, derive(Clone, Debug, Serialize, Deserialize))]
#[cfg_attr(target_arch = "wasm32", derive(Clone, Debug, Serialize, Deserialize))]
pub struct RobotSetFKResult {
robot_fk_results: Vec<RobotFKResult>
}
impl RobotSetFKResult {
pub fn get_pose_from_idxs(&self, robot_idx_in_set: usize, link_idx_in_robot: usize) ->&OptimaSE3Pose {
let res = &self.robot_fk_results[robot_idx_in_set].link_entries()[link_idx_in_robot].pose().as_ref().unwrap();
return res;
}
pub fn get_pose_option_from_idxs(&self, robot_idx_in_set: usize, link_idx_in_robot: usize) -> Option<&OptimaSE3Pose> {
let res = self.robot_fk_results[robot_idx_in_set].link_entries()[link_idx_in_robot].pose().as_ref();
return res;
}
pub fn robot_fk_results(&self) -> &Vec<RobotFKResult> {
&self.robot_fk_results
}
pub fn robot_fk_result(&self, robot_idx_in_set: usize) -> Result<&RobotFKResult, OptimaError> {
OptimaError::new_check_for_idx_out_of_bound_error(robot_idx_in_set, self.robot_fk_results.len(), file!(), line!())?;
return Ok(&self.robot_fk_results[robot_idx_in_set]);
}
pub fn print_summary(&self) {
for (i, robot_fk_result) in self.robot_fk_results.iter().enumerate() {
optima_print(&format!("Robot {} ---> ", i), PrintMode::Println, PrintColor::Cyan, true, 0, None, vec![]);
robot_fk_result.print_summary();
}
}
}
#[cfg(not(target_arch = "wasm32"))]
#[pymethods]
impl RobotSetFKResult {
pub fn get_fk_result(&self, idx: usize) -> RobotFKResult {
self.robot_fk_results.get(idx).unwrap().clone()
}
pub fn robot_fk_results_py(&self) -> Vec<RobotFKResult> {
self.robot_fk_results.clone()
}
pub fn robot_fk_result_link_entries(&self) -> Vec<Vec<RobotFKResultLinkEntry>> {
let mut out_vec = vec![];
for fk_res in &self.robot_fk_results {
out_vec.push(fk_res.link_entries().clone())
}
out_vec
}
pub fn print_summary_py(&self) {
self.print_summary();
}
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct RobotSetFKDOFPerturbationsResult {
perturbation: f64,
central_fk_result: RobotSetFKResult,
fk_dof_perturbation_results: Vec<RobotSetFKResult>
}
impl RobotSetFKDOFPerturbationsResult {
pub fn perturbation(&self) -> f64 {
self.perturbation
}
pub fn central_fk_result(&self) -> &RobotSetFKResult {
&self.central_fk_result
}
pub fn fk_dof_perturbation_results(&self) -> &Vec<RobotSetFKResult> {
&self.fk_dof_perturbation_results
}
}