#include "RBDUtil.h"
#include <iostream>

void cRBDUtil::SolveInvDyna(const cRBDModel& model, const Eigen::VectorXd& acc, Eigen::VectorXd& out_tau)
{
	const Eigen::MatrixXd& joint_mat = model.GetJointMat();
	const Eigen::MatrixXd& body_defs = model.GetBodyDefs();
	const tVector& gravity = model.GetGravity();
	const Eigen::VectorXd& pose = model.GetPose();
	const Eigen::VectorXd& vel = model.GetVel();

	assert(joint_mat.rows() == body_defs.rows());
	assert(pose.rows() == vel.rows());
	assert(pose.rows() == acc.rows());
	assert(cKinTree::GetNumDof(joint_mat) == pose.rows());

	cSpAlg::tSpVec vel0 = cSpAlg::tSpVec::Zero();
	cSpAlg::tSpVec acc0 = cSpAlg::BuildSV(tVector::Zero(), -gravity);

	int num_joints = cKinTree::GetNumJoints(joint_mat);
	Eigen::MatrixXd vels = Eigen::MatrixXd(num_joints, cSpAlg::gSpVecSize);
	Eigen::MatrixXd accs = Eigen::MatrixXd(num_joints, cSpAlg::gSpVecSize);
	Eigen::MatrixXd fs = Eigen::MatrixXd(num_joints, cSpAlg::gSpVecSize);

	for (int j = 0; j < num_joints; ++j)
	{
		if (cKinTree::IsValidBody(body_defs, j))
		{
			cSpAlg::tSpTrans parent_child_trans = model.GetSpParentChildTrans(j);
			cSpAlg::tSpTrans world_child_trans = model.GetSpWorldJointTrans(j);

			const auto S = model.GetJointSubspace(j);
			Eigen::VectorXd q = cKinTree::GetJointParams(joint_mat, pose, j);
			Eigen::VectorXd dq = cKinTree::GetJointParams(joint_mat, vel, j);
			Eigen::VectorXd ddq = cKinTree::GetJointParams(joint_mat, acc, j);

			cSpAlg::tSpVec cj = BuildCj(joint_mat, dq, j);
			cSpAlg::tSpVec vj = S * dq;

			cSpAlg::tSpMat I = BuildInertiaSpatialMat(body_defs, j);

			cSpAlg::tSpVec vel_p;
			cSpAlg::tSpVec acc_p;
			if (cKinTree::HasParent(joint_mat, j))
			{
				int parent_id = cKinTree::GetParent(joint_mat, j);
				vel_p = vels.row(parent_id);
				acc_p = accs.row(parent_id);
			}
			else
			{
				vel_p = vel0;
				acc_p = acc0;
			}

			cSpAlg::tSpVec curr_vel = cSpAlg::ApplyTransM(parent_child_trans, vel_p) + vj;
			cSpAlg::tSpVec curr_acc = cSpAlg::ApplyTransM(parent_child_trans, acc_p) + S * ddq + cj + cSpAlg::CrossM(curr_vel, vj);
			cSpAlg::tSpVec curr_f = I * curr_acc + cSpAlg::CrossF(curr_vel, I * curr_vel);
			
			vels.row(j) = curr_vel;
			accs.row(j) = curr_acc;
			fs.row(j) = curr_f;
		}
	}

	out_tau = Eigen::VectorXd::Zero(pose.size());
	for (int j = num_joints - 1; j >= 0; --j)
	{
		if (cKinTree::IsValidBody(body_defs, j))
		{
			cSpAlg::tSpVec curr_f = fs.row(j);
			const auto S = model.GetJointSubspace(j);
			Eigen::VectorXd curr_tau = S.transpose() * curr_f;

			cKinTree::SetJointParams(joint_mat, j, curr_tau, out_tau);
			if (cKinTree::HasParent(joint_mat, j))
			{
				int parent_id = cKinTree::GetParent(joint_mat, j);
				cSpAlg::tSpTrans child_parent_trans = model.GetSpChildParentTrans(j);
				fs.row(parent_id) += cSpAlg::ApplyTransF(child_parent_trans, curr_f);
			}
		}
	}
}

void cRBDUtil::SolveForDyna(const cRBDModel& model, const Eigen::VectorXd& tau, Eigen::VectorXd& out_acc)
{
	Eigen::VectorXd total_force = Eigen::VectorXd::Zero(model.GetNumDof());
	SolveForDyna(model, tau, total_force, out_acc);
}

void cRBDUtil::SolveForDyna(const cRBDModel& model, const Eigen::VectorXd& tau, const Eigen::VectorXd& total_force, Eigen::VectorXd& out_acc)
{
	Eigen::VectorXd C;
	Eigen::MatrixXd H;
	BuildBiasForce(model, C);
	BuildMassMat(model, H);
	out_acc = H.ldlt().solve(tau + total_force - C);
}


void cRBDUtil::BuildMassMat(const cRBDModel& model, Eigen::MatrixXd& out_mass_mat)
{
	const int svs = cSpAlg::gSpVecSize;
	int num_joints = model.GetNumJoints();
	Eigen::MatrixXd Is = Eigen::MatrixXd::Zero(num_joints * svs, svs);
	BuildMassMat(model, Is, out_mass_mat);
}

void cRBDUtil::BuildMassMat(const cRBDModel& model, Eigen::MatrixXd& inertia_buffer, Eigen::MatrixXd& out_mass_mat)
{
	// use composite-rigid-body algorithm
	const Eigen::MatrixXd& joint_mat = model.GetJointMat();
	const Eigen::MatrixXd& body_defs = model.GetBodyDefs();
	const Eigen::VectorXd& pose = model.GetPose();
	Eigen::MatrixXd& H = out_mass_mat;

	int dim = model.GetNumDof();
	int num_joints = model.GetNumJoints();
	H.setZero(dim, dim);
	const int svs = cSpAlg::gSpVecSize;

	Eigen::MatrixXd child_parent_mats_F = Eigen::MatrixXd(svs * num_joints, svs);
	Eigen::MatrixXd parent_child_mats_M = Eigen::MatrixXd(svs * num_joints, svs);
	Eigen::MatrixXd& Is = inertia_buffer;
	for (int j = 0; j < num_joints; ++j)
	{
		if (cKinTree::IsValidBody(body_defs, j))
		{
			Is.block(j * svs, 0, svs, svs) = BuildInertiaSpatialMat(body_defs, j);
		}

		cSpAlg::tSpTrans child_parent_trans = model.GetSpChildParentTrans(j);
		cSpAlg::tSpMat child_parent_mat = cSpAlg::BuildSpatialMatF(child_parent_trans);
		cSpAlg::tSpMat parent_child_mat = cSpAlg::BuildSpatialMatM(cSpAlg::InvTrans(child_parent_trans));
		child_parent_mats_F.block(j * svs, 0, svs, svs) = child_parent_mat;
		parent_child_mats_M.block(j * svs, 0, svs, svs) = parent_child_mat;
	}

	for (int j = num_joints - 1; j >= 0; --j)
	{
		if (cKinTree::IsValidBody(body_defs, j))
		{
			const auto curr_I = Is.block(j * svs, 0, svs, svs);
			int parent_id = cKinTree::GetParent(joint_mat, j);
			if (cKinTree::HasParent(joint_mat, j))
			{
				cSpAlg::tSpTrans child_parent_trans = model.GetSpChildParentTrans(j);
				cSpAlg::tSpMat child_parent_mat = child_parent_mats_F.block(j * svs, 0, svs, svs);
				cSpAlg::tSpMat parent_child_mat = parent_child_mats_M.block(j * svs, 0, svs, svs);
				Is.block(parent_id * svs, 0, svs, svs) += child_parent_mat * curr_I * parent_child_mat;
			}

			const auto S = model.GetJointSubspace(j);
			int offset = cKinTree::GetParamOffset(joint_mat, j);
			int dim = cKinTree::GetParamSize(joint_mat, j);
			Eigen::MatrixXd F = curr_I * S;
			H.block(offset, offset, dim, dim) = S.transpose() * F;

			int curr_id = j;
			while (cKinTree::HasParent(joint_mat, curr_id))
			{
				cSpAlg::tSpMat child_parent_mat = child_parent_mats_F.block(curr_id * svs, 0, svs, svs);
				F = child_parent_mat * F;

				curr_id = cKinTree::GetParent(joint_mat, curr_id);
				int curr_offset = cKinTree::GetParamOffset(joint_mat, curr_id);
				int curr_dim = cKinTree::GetParamSize(joint_mat, curr_id);

				const auto curr_S = model.GetJointSubspace(curr_id);
				H.block(offset, curr_offset, dim, curr_dim) = F.transpose() * curr_S;
				H.block(curr_offset, offset, curr_dim, dim) = H.block(offset, curr_offset, dim, curr_dim).transpose();
			}
		}
	}
}

void cRBDUtil::BuildEndEffectorJacobian(const cRBDModel& model, int joint_id, Eigen::MatrixXd& out_J)
{
	// jacobian in world coordinates
	const Eigen::MatrixXd& joint_mat = model.GetJointMat();
	const Eigen::VectorXd& pose = model.GetPose();
	
	int num_dofs = cKinTree::GetNumDof(joint_mat);
	out_J = Eigen::MatrixXd::Zero(cSpAlg::gSpVecSize, num_dofs);

	int curr_id = joint_id;
	cSpAlg::tSpTrans curr_trans = cSpAlg::BuildTrans();
	while (curr_id != cKinTree::gInvalidJointID)
	{
		int offset = cKinTree::GetParamOffset(joint_mat, curr_id);
		int size = cKinTree::GetParamSize(joint_mat, curr_id);
		const Eigen::MatrixXd S = model.GetJointSubspace(curr_id);

		out_J.block(0, offset, cSpAlg::gSpVecSize, size) = cSpAlg::ApplyTransM(curr_trans, S);

		int parent_id = cKinTree::GetParent(joint_mat, curr_id);
		cSpAlg::tSpTrans parent_child_trans = model.GetSpParentChildTrans(curr_id);
		curr_trans = cSpAlg::CompTrans(curr_trans, parent_child_trans);
		curr_id = parent_id;
	}

	out_J = cSpAlg::ApplyInvTransM(curr_trans, out_J);
}

void cRBDUtil::BuildEndEffectorJacobian(const Eigen::MatrixXd& joint_mat, const Eigen::VectorXd& pose, int joint_id, Eigen::MatrixXd& out_J)
{
	// jacobian in world coordinates
	int num_dofs = cKinTree::GetNumDof(joint_mat);
	out_J = Eigen::MatrixXd::Zero(cSpAlg::gSpVecSize, num_dofs);

	int curr_id = joint_id;
	cSpAlg::tSpTrans curr_trans = cSpAlg::BuildTrans();
	while (curr_id != cKinTree::gInvalidJointID)
	{
		int offset = cKinTree::GetParamOffset(joint_mat, curr_id);
		int size = cKinTree::GetParamSize(joint_mat, curr_id);
		Eigen::MatrixXd S = BuildJointSubspace(joint_mat, pose, curr_id);

		S = cSpAlg::ApplyTransM(curr_trans, S);
		out_J.block(0, offset, cSpAlg::gSpVecSize, size) = S;

		int parent_id = cKinTree::GetParent(joint_mat, curr_id);
		cSpAlg::tSpTrans parent_child_trans = BuildParentChildTransform(joint_mat, pose, curr_id);
		curr_trans = cSpAlg::CompTrans(curr_trans, parent_child_trans);
		curr_id = parent_id;
	}

	out_J = cSpAlg::ApplyInvTransM(curr_trans, out_J);
}

Eigen::MatrixXd cRBDUtil::MultJacobianEndEff(const Eigen::MatrixXd& joint_mat, const Eigen::VectorXd& q, const Eigen::MatrixXd& J, int joint_id)
{
	// multiplies q by the jacobian of the end effector (joint_id)
	int curr_id = joint_id;
	cSpAlg::tSpVec joint_vel = cSpAlg::tSpVec::Zero();
	while (curr_id != cKinTree::gInvalidJointID)
	{
		int offset = cKinTree::GetParamOffset(joint_mat, curr_id);
		int size = cKinTree::GetParamSize(joint_mat, curr_id);
		const auto curr_q = cKinTree::GetJointParams(joint_mat, q, curr_id);
		const auto curr_J = J.block(0, offset, cSpAlg::gSpVecSize, size);
		
		joint_vel += curr_J * curr_q;
		curr_id = cKinTree::GetParent(joint_mat, curr_id);
	}
	return joint_vel;
}

void cRBDUtil::BuildJacobian(const cRBDModel& model, Eigen::MatrixXd& out_J)
{
	const Eigen::MatrixXd& joint_mat = model.GetJointMat();
	const Eigen::VectorXd& pose = model.GetPose();

	int num_dofs = model.GetNumDof();
	out_J = Eigen::MatrixXd::Zero(cSpAlg::gSpVecSize, num_dofs);

	int num_joints = cKinTree::GetNumJoints(joint_mat);
	for (int j = 0; j < num_joints; ++j)
	{
		cSpAlg::tSpTrans world_joint_trans = model.GetSpWorldJointTrans(j);

		int offset = cKinTree::GetParamOffset(joint_mat, j);
		int size = cKinTree::GetParamSize(joint_mat, j);
		const Eigen::MatrixXd S = model.GetJointSubspace(j);

		out_J.block(0, offset, cSpAlg::gSpVecSize, size) = cSpAlg::ApplyInvTransM(world_joint_trans, S);
	}
}

Eigen::MatrixXd cRBDUtil::ExtractEndEffJacobian(const Eigen::MatrixXd& joint_mat, const Eigen::MatrixXd& J, int joint_id)
{
	int curr_id = joint_id;
	Eigen::MatrixXd J_end_eff = Eigen::MatrixXd::Zero(J.rows(), J.cols());

	while (curr_id != cKinTree::gInvalidJointID)
	{
		int offset = cKinTree::GetParamOffset(joint_mat, curr_id);
		int size = cKinTree::GetParamSize(joint_mat, curr_id);
		const auto curr_J = J.block(0, offset, cSpAlg::gSpVecSize, size);

		J_end_eff.block(0, offset, cSpAlg::gSpVecSize, size) = curr_J;
		curr_id = cKinTree::GetParent(joint_mat, curr_id);
	}
	return J_end_eff;
}


void cRBDUtil::BuildCOMJacobian(const cRBDModel& model, Eigen::MatrixXd& out_J)
{
	// coord frame for jacobian has origin at the com
	const Eigen::MatrixXd& joint_mat = model.GetJointMat();
	const Eigen::MatrixXd& body_defs = model.GetBodyDefs();
	const Eigen::VectorXd& pose = model.GetPose();

	Eigen::MatrixXd J;
	BuildJacobian(model, J);
	BuildCOMJacobian(model, J, out_J);
}

void cRBDUtil::BuildCOMJacobian(const cRBDModel& model, const Eigen::MatrixXd& J, Eigen::MatrixXd& out_J)
{
	const Eigen::MatrixXd& joint_mat = model.GetJointMat();
	const Eigen::MatrixXd& body_defs = model.GetBodyDefs();

	int num_dofs = cKinTree::GetNumDof(joint_mat);
	out_J = Eigen::MatrixXd::Zero(cSpAlg::gSpVecSize, num_dofs);
	double total_mass = cKinTree::CalcTotalMass(body_defs);

	int num_joints = cKinTree::GetNumJoints(joint_mat);
	for (int j = num_joints - 1; j >= 0; --j)
	{
		if (cKinTree::IsValidBody(body_defs, j))
		{
			double mass = cKinTree::GetBodyMass(body_defs, j);
			double mass_frac = mass / total_mass;

			cSpAlg::tSpTrans world_child_trans = model.GetSpWorldJointTrans(j);
			tMatrix body_joint_mat = cKinTree::BodyJointTrans(body_defs, j);
			cSpAlg::tSpTrans body_world_trans = cSpAlg::CompTrans(cSpAlg::InvTrans(world_child_trans), cSpAlg::MatToTrans(body_joint_mat));
			tMatrix body_world_mat = cSpAlg::TransToMat(body_world_trans);

			tVector body_pos = body_world_mat * tVector(0, 0, 0, 1);
			cSpAlg::tSpTrans body_pos_trans = cSpAlg::BuildTrans(body_pos);

			int curr_id = j;
			while (curr_id != cKinTree::gInvalidJointID)
			{
				int offset = cKinTree::GetParamOffset(joint_mat, curr_id);
				int size = cKinTree::GetParamSize(joint_mat, curr_id);
				
				const Eigen::MatrixXd& J_block = J.block(0, offset, cSpAlg::gSpVecSize, size);
				auto J_com_block = out_J.block(0, offset, cSpAlg::gSpVecSize, size);
				J_com_block += mass_frac * cSpAlg::ApplyTransM(body_pos_trans, J_block);

				curr_id = cKinTree::GetParent(joint_mat, curr_id);
			}
		}
	}
}


void cRBDUtil::BuildJacobianDot(const cRBDModel& model, Eigen::MatrixXd& out_J_dot)
{
	// for comput the velocity product acceleration
	const Eigen::MatrixXd& joint_mat = model.GetJointMat();
	const Eigen::VectorXd& pose = model.GetPose();
	const Eigen::VectorXd& vel = model.GetVel();

	int num_dofs = cKinTree::GetNumDof(joint_mat);
	int num_joints = cKinTree::GetNumJoints(joint_mat);
	out_J_dot = Eigen::MatrixXd(cSpAlg::gSpVecSize, num_dofs);
	Eigen::MatrixXd Sqs(cSpAlg::gSpVecSize, num_joints);

	for (int j = 0; j < num_joints; ++j)
	{
		cSpAlg::tSpTrans world_child_trans = model.GetSpWorldJointTrans(j);
		Eigen::MatrixXd S = model.GetJointSubspace(j);
		S = cSpAlg::ApplyInvTransM(world_child_trans, S);
		Eigen::VectorXd dq = cKinTree::GetJointParams(joint_mat, vel, j);
		cSpAlg::tSpVec Sq = S * dq;

		cSpAlg::tSpVec parent_Sq = cSpAlg::tSpVec::Zero();
		int parent_id = cKinTree::GetParent(joint_mat, j);
		if (parent_id != cKinTree::gInvalidJointID)
		{
			parent_Sq = Sqs.col(parent_id);
		}
		Sqs.col(j) = parent_Sq + Sq;

		int offset = cKinTree::GetParamOffset(joint_mat, j);
		int size = cKinTree::GetParamSize(joint_mat, j);
		out_J_dot.block(0, offset, cSpAlg::gSpVecSize, size) = cSpAlg::CrossMs(parent_Sq, S);
	}
}

cSpAlg::tSpVec cRBDUtil::BuildCOMVelProdAcc(const cRBDModel& model)
{
	Eigen::MatrixXd Jd;
	BuildJacobianDot(model, Jd);
	return BuildCOMVelProdAccAux(model, Jd);
}

cSpAlg::tSpVec cRBDUtil::BuildCOMVelProdAccAux(const cRBDModel& model, const Eigen::MatrixXd& Jd)
{
	const Eigen::MatrixXd& joint_mat = model.GetJointMat();
	const Eigen::MatrixXd& body_defs = model.GetBodyDefs();
	const Eigen::VectorXd& pose = model.GetPose();
	const Eigen::VectorXd& vel = model.GetVel();
	const tVector& gravity = model.GetGravity();

	// coord frame origin at com
	int num_dofs = cKinTree::GetNumDof(joint_mat);
	int num_joints = cKinTree::GetNumJoints(joint_mat);
	
	double total_mass = cKinTree::CalcTotalMass(body_defs);
	cSpAlg::tSpVec com_vpa = cSpAlg::tSpVec::Zero();
	for (int j = 0; j < num_joints; ++j)
	{
		if (cKinTree::IsValidBody(body_defs, j))
		{
			double mass = cKinTree::GetBodyMass(body_defs, j);
			double mass_frac = mass / total_mass;

			cSpAlg::tSpVec vpa = CalcVelProdAcc(model, Jd, j);
			com_vpa += mass_frac * vpa;
		}
	}

	tVector com = CalcCoMPos(model);
	cSpAlg::tSpTrans trans = cSpAlg::BuildTrans(com);
	com_vpa = cSpAlg::ApplyTransM(trans, com_vpa);

	return com_vpa;
}

cSpAlg::tSpVec cRBDUtil::CalcVelProdAcc(const cRBDModel& model, const Eigen::MatrixXd& Jd, int joint_id)
{
	const Eigen::MatrixXd& joint_mat = model.GetJointMat();
	const Eigen::VectorXd& pose = model.GetPose();
	const Eigen::VectorXd& vel = model.GetVel();
	const tVector& gravity = model.GetGravity();

	int curr_id = joint_id;
	cSpAlg::tSpVec acc = cSpAlg::BuildSV(tVector::Zero(), -gravity);

	while (curr_id != cKinTree::gInvalidJointID)
	{
		int offset = cKinTree::GetParamOffset(joint_mat, curr_id);
		int size = cKinTree::GetParamSize(joint_mat, curr_id);
		const auto dq = cKinTree::GetJointParams(joint_mat, vel, curr_id);
		const auto curr_Jd = Jd.block(0, offset, cSpAlg::gSpVecSize, size);

		cSpAlg::tSpVec cj = BuildCj(joint_mat, dq, curr_id);
		if (cj.squaredNorm() > 0)
		{
			cSpAlg::tSpTrans world_joint_trans = model.GetSpWorldJointTrans(curr_id);
			cj = cSpAlg::ApplyInvTransM(world_joint_trans, cj);
		}

		acc += curr_Jd * dq + cj;
		curr_id = cKinTree::GetParent(joint_mat, curr_id);
	}

	return acc;
}

cSpAlg::tSpVec cRBDUtil::CalcJointWorldVel(const cRBDModel& model, int joint_id)
{
	cSpAlg::tSpVec joint_vel = CalcJointWorldVel(model.GetJointMat(), model.GetPose(), model.GetVel(), joint_id);
	return joint_vel;
}

cSpAlg::tSpVec cRBDUtil::CalcJointWorldVel(const Eigen::MatrixXd& joint_mat, const Eigen::VectorXd& state, const Eigen::VectorXd& vel, int joint_id)
{
	cSpAlg::tSpVec joint_vel = CalcWorldVel(joint_mat, state, vel, joint_id, tVector::Zero());
	return joint_vel;
}

cSpAlg::tSpVec cRBDUtil::CalcWorldVel(const cRBDModel& model, int parent_id, const tVector& attach_pt)
{
	return CalcWorldVel(model.GetJointMat(), model.GetPose(), model.GetVel(), parent_id, attach_pt);
}

cSpAlg::tSpVec cRBDUtil::CalcWorldVel(const Eigen::MatrixXd& joint_mat, const Eigen::VectorXd& pose, const Eigen::VectorXd& vel, int parent_id, const tVector& attach_pt)
{
	Eigen::MatrixXd J;
	BuildEndEffectorJacobian(joint_mat, pose, parent_id, J);
	cSpAlg::tSpVec sv = J * vel;
	return sv;
}

tVector cRBDUtil::CalcCoMPos(const cRBDModel& model)
{
	tVector com;
	tVector com_vel;
	Eigen::VectorXd vel = Eigen::VectorXd::Zero(model.GetNumDof());
	CalcCoM(model, com, com_vel);
	return com;
}

tVector cRBDUtil::CalcCoMVel(const cRBDModel& model)
{
	tVector com;
	tVector com_vel;
	CalcCoM(model, com, com_vel);
	return com_vel;
}

void cRBDUtil::CalcCoM(const cRBDModel& model, tVector& out_com, tVector& out_vel)
{
	const Eigen::MatrixXd& joint_mat = model.GetJointMat();
	const Eigen::MatrixXd& body_defs = model.GetBodyDefs();
	const Eigen::VectorXd& pose = model.GetPose();
	const Eigen::VectorXd& vel = model.GetVel();

	int num_joints = cKinTree::GetNumJoints(joint_mat);
	out_com.setZero();
	out_vel.setZero();
	double total_mass = 0;

	for (int j = 0; j < num_joints; ++j)
	{
		if (cKinTree::IsValidBody(body_defs, j))
		{
			tVector local_com = cKinTree::GetBodyLocalCoM(body_defs, j);
			cSpAlg::tSpTrans world_joint_trans = model.GetSpWorldJointTrans(j);

			tMatrix joint_world_mat = cSpAlg::TransToMat(cSpAlg::InvTrans(world_joint_trans));
			tMatrix body_joint_mat = cKinTree::BodyJointTrans(body_defs, j);
			tMatrix body_world_mat = joint_world_mat * body_joint_mat;

			tVector attach_pt = cKinTree::GetBodyAttachPt(body_defs, j);
			attach_pt[3] = 1;
			attach_pt = body_joint_mat * attach_pt;

			local_com[3] = 1;
			tVector world_com = body_world_mat * local_com;
			world_com[3] = 0;

			cSpAlg::tSpTrans com_trans = cSpAlg::BuildTrans(world_com);
			cSpAlg::tSpVec sv = CalcWorldVel(model, j, attach_pt);
			sv = cSpAlg::ApplyTransM(com_trans, sv);

			tVector com_vel = cSpAlg::GetV(sv);

			double m = cKinTree::GetBodyMass(body_defs, j);
			out_com += m * world_com;
			out_vel += m * com_vel;
			total_mass += m;
		}
	}

	assert(total_mass > 0);
	out_com /= total_mass;
	out_vel /= total_mass;
}

cSpAlg::tSpMat cRBDUtil::BuildMomentInertia(const Eigen::MatrixXd& body_defs, int part_id)
{
	// inertia tensor of shape centered at the com
	assert(cKinTree::IsValidBody(body_defs, part_id));
	cKinTree::eBodyShape shape = cKinTree::GetBodyShape(body_defs, part_id);

	cSpAlg::tSpMat I;
	switch (shape)
	{
	case cKinTree::eBodyShapeBox:
		I = BuildMomentInertiaBox(body_defs, part_id);
		break;
	case cKinTree::eBodyShapeCapsule:
		I = BuildMomentInertiaCapsule(body_defs, part_id);
		break;
	default:
		assert(false); // unsupported shape
		break;
	}

	return I;
}

cSpAlg::tSpMat cRBDUtil::BuildMomentInertiaBox(const Eigen::MatrixXd& body_defs, int part_id)
{
	const cKinTree::tBodyDef& def = body_defs.row(part_id);
	double mass = cKinTree::GetBodyMass(body_defs, part_id);
	double sx = def(cKinTree::eBodyParam0);
	double sy = def(cKinTree::eBodyParam1);
	double sz = def(cKinTree::eBodyParam2);

	double x = mass / 12.0 * (sy * sy + sz * sz);
	double y = mass / 12.0 * (sx * sx + sz * sz);
	double z = mass / 12.0 * (sx * sx + sy * sy);

	cSpAlg::tSpMat I = cSpAlg::tSpMat::Zero();
	I(0, 0) = x;
	I(1, 1) = y;
	I(2, 2) = z;
	I(3, 3) = mass;
	I(4, 4) = mass;
	I(5, 5) = mass;

	return I;
}

cSpAlg::tSpMat cRBDUtil::BuildMomentInertiaCapsule(const Eigen::MatrixXd& body_defs, int part_id)
{
	const cKinTree::tBodyDef& def = body_defs.row(part_id);
	double mass = cKinTree::GetBodyMass(body_defs, part_id);
	double r = def(cKinTree::eBodyParam0);
	double h = def(cKinTree::eBodyParam1);

	double c_vol = M_PI * r * r * h;
	double hs_vol = M_PI * 2.0 / 3.0 * r * r * r;
	double density = mass / (c_vol + 2 * hs_vol);
	double cm = c_vol * density;
	double hsm = hs_vol * density;

	double x = cm*(0.25 *r * r + (1.0 / 12.0) * cm * h * h) +
				2 * hsm *(0.4 * r * r + 0.375 * r * h + 0.25 * h * h);
	double y = (0.5 * cm + 0.8 * hsm) * r * r;
	double z = x;

	cSpAlg::tSpMat I = cSpAlg::tSpMat::Zero();
	I(0, 0) = x;
	I(1, 1) = y;
	I(2, 2) = z;
	I(3, 3) = mass;
	I(4, 4) = mass;
	I(5, 5) = mass;

	return I;
}

cSpAlg::tSpMat cRBDUtil::BuildInertiaSpatialMat(const Eigen::MatrixXd& body_defs, int part_id)
{
	cSpAlg::tSpMat Ic = BuildMomentInertia(body_defs, part_id);
	tMatrix E = tMatrix::Identity();
	tVector r = -cKinTree::GetBodyAttachPt(body_defs, part_id);
	cSpAlg::tSpTrans X = cSpAlg::BuildTrans(E, r);

	cSpAlg::tSpMat Io = cSpAlg::BuildSpatialMatF(X) * Ic * cSpAlg::BuildSpatialMatM(cSpAlg::InvTrans(X));
	return Io;
}

void cRBDUtil::CalcWorldJointTransforms(const cRBDModel& model, Eigen::MatrixXd& out_trans_arr)
{
	const Eigen::MatrixXd& joint_mat = model.GetJointMat();
	const Eigen::VectorXd& pose = model.GetPose();

	int num_joints = cKinTree::GetNumJoints(joint_mat);
	out_trans_arr.resize(num_joints * cSpAlg::gSVTransRows, cSpAlg::gSVTransCols);

	for (int j = 0; j < num_joints; ++j)
	{
		int row_idx = j * cSpAlg::gSVTransRows;
		int parent_id = cKinTree::GetParent(joint_mat, j);

		cSpAlg::tSpTrans parent_child_trans = model.GetSpParentChildTrans(j);
		cSpAlg::tSpTrans world_parent_trans = cSpAlg::BuildTrans();

		if (parent_id != cKinTree::gInvalidJointID)
		{
			world_parent_trans = cSpAlg::GetTrans(out_trans_arr, parent_id);
		}

		cSpAlg::tSpTrans world_child_trans = cSpAlg::CompTrans(parent_child_trans, world_parent_trans);
		out_trans_arr.block(row_idx, 0, cSpAlg::gSVTransRows, cSpAlg::gSVTransCols) = world_child_trans;
	}
}

cSpAlg::tSpTrans cRBDUtil::BuildParentChildTransform(const Eigen::MatrixXd& joint_mat, const Eigen::VectorXd& pose, int j)
{
	cSpAlg::tSpTrans trans = BuildChildParentTransform(joint_mat, pose, j);
	trans = cSpAlg::InvTrans(trans);
	return trans;
}

cSpAlg::tSpTrans cRBDUtil::BuildChildParentTransform(const Eigen::MatrixXd& joint_mat, const Eigen::VectorXd& pose, int j)
{
	tMatrix m = cKinTree::ChildParentTrans(joint_mat, pose, j);
	cSpAlg::tSpTrans trans = cSpAlg::MatToTrans(m);
	return trans;
}

cSpAlg::tSpTrans cRBDUtil::BuildWorldJointTransform(const Eigen::MatrixXd& joint_mat, const Eigen::VectorXd& pose, int j)
{
	cSpAlg::tSpTrans trans = BuildJointWorldTransform(joint_mat, pose, j);
	trans = cSpAlg::InvTrans(trans);
	return trans;
}

cSpAlg::tSpTrans cRBDUtil::BuildJointWorldTransform(const Eigen::MatrixXd& joint_mat, const Eigen::VectorXd& pose, int j)
{
	tMatrix m = cKinTree::JointWorldTrans(joint_mat, pose, j);
	cSpAlg::tSpTrans trans = cSpAlg::MatToTrans(m);
	return trans;
}

bool cRBDUtil::IsConstJointSubspace(const Eigen::MatrixXd& joint_mat, int j)
{
	bool is_const = true;
	bool is_root = cKinTree::IsRoot(joint_mat, j);

	cKinTree::eJointType j_type = cKinTree::GetJointType(joint_mat, j);
	switch (j_type)
	{
	case cKinTree::eJointTypePlanar:
		is_const = !is_root;
		break;
	default:
		break;
	}
	return is_const;
}

Eigen::MatrixXd cRBDUtil::BuildJointSubspace(const Eigen::MatrixXd& joint_mat, const Eigen::VectorXd& pose, int j)
{
	cKinTree::eJointType j_type = cKinTree::GetJointType(joint_mat, j);
	Eigen::MatrixXd S;
	switch (j_type)
	{
	case cKinTree::eJointTypeRevolute:
		S = BuildJointSubspaceRevolute(joint_mat, pose, j);
		break;
	case cKinTree::eJointTypePrismatic:
		S = BuildJointSubspacePrismatic(joint_mat, pose, j);
		break;
	case cKinTree::eJointTypePlanar:
		S = BuildJointSubspacePlanar(joint_mat, pose, j);
		break;
	case cKinTree::eJointTypeFixed:
		S = BuildJointSubspaceFixed(joint_mat, pose, j);
		break;
	default:
		assert(false); // unsupported joint type;
		break;
	}
	return S;
}

Eigen::MatrixXd cRBDUtil::BuildJointSubspaceRevolute(const Eigen::MatrixXd& joint_mat, const Eigen::VectorXd& pose, int j)
{
	int dim = 1;
	Eigen::MatrixXd S = Eigen::MatrixXd::Zero(cSpAlg::gSpVecSize, dim);

	S(2, 0) = 1;
	return S;
}

Eigen::MatrixXd cRBDUtil::BuildJointSubspacePrismatic(const Eigen::MatrixXd& joint_mat, const Eigen::VectorXd& pose, int j)
{
	// assert(false); // TODO: this likely doesn't work
	int dim = 1;
	Eigen::MatrixXd S = Eigen::MatrixXd::Zero(cSpAlg::gSpVecSize, dim);

	S(3, 0) = 1;
	return S;
}

Eigen::MatrixXd cRBDUtil::BuildJointSubspacePlanar(const Eigen::MatrixXd& joint_mat, const Eigen::VectorXd& pose, int j)
{
	int dim = 3;
	Eigen::MatrixXd S = Eigen::MatrixXd::Zero(cSpAlg::gSpVecSize, dim);

	bool is_root = cKinTree::IsRoot(joint_mat, j);

	// in contrast to featherstone's convention, translation comes before rotation....
	if (is_root)
	{
		double theta = cKinTree::GetJointTheta(joint_mat, pose, j);
		tMatrix E = cMathUtil::RotateMat(tVector(0, 0, 1, 0), -theta);
		
		S.block(3, 0, 2, 2) = E.block(0, 0, 2, 2);
		S(2, 2) = 1;
	}
	else
	{
		S(3, 0) = 1;
		S(4, 1) = 1;
		S(2, 2) = 1;
	}

	return S;
}

Eigen::MatrixXd cRBDUtil::BuildJointSubspaceFixed(const Eigen::MatrixXd& joint_mat, const Eigen::VectorXd& pose, int j)
{
	bool is_root = cKinTree::IsRoot(joint_mat, j);
	int dim = (is_root) ? 3 : 1;
	Eigen::MatrixXd S = Eigen::MatrixXd::Zero(cSpAlg::gSpVecSize, dim);
	return S;
}

cSpAlg::tSpVec cRBDUtil::BuildCj(const Eigen::MatrixXd& joint_mat, const Eigen::VectorXd& q_dot, int j)
{
	cKinTree::eJointType j_type = cKinTree::GetJointType(joint_mat, j);
	cSpAlg::tSpVec cj;
	switch (j_type)
	{
	case cKinTree::eJointTypeRevolute:
		cj = BuildCjRevolute(joint_mat, q_dot, j);
		break;
	case cKinTree::eJointTypePrismatic:
		cj = BuildCjPrismatic(joint_mat, q_dot, j);
		break;
	case cKinTree::eJointTypePlanar:
		cj = BuildCjPlanar(joint_mat, q_dot, j);
		break;
	case cKinTree::eJointTypeFixed:
		cj = BuildCjFixed(joint_mat, q_dot, j);
		break;
	default:
		assert(false); // unsupported joint type;
		break;
	}
	return cj;
}

cSpAlg::tSpVec cRBDUtil::BuildCjRevolute(const Eigen::MatrixXd& joint_mat, const Eigen::VectorXd& q_dot, int j)
{
	return cSpAlg::tSpVec::Zero();
}

cSpAlg::tSpVec cRBDUtil::BuildCjPrismatic(const Eigen::MatrixXd& joint_mat, const Eigen::VectorXd& q_dot, int j)
{
	return cSpAlg::tSpVec::Zero();
}

cSpAlg::tSpVec cRBDUtil::BuildCjPlanar(const Eigen::MatrixXd& joint_mat, const Eigen::VectorXd& q_dot, int j)
{
	bool is_root = cKinTree::IsRoot(joint_mat, j);
	cSpAlg::tSpVec cj;

	if (is_root)
	{
		tVector offset = cKinTree::GetJointOffset(joint_mat, q_dot, j);
		double theta = cKinTree::GetJointTheta(joint_mat, q_dot, j);
		double x = offset[0];
		double y = offset[1];

		double c = std::cos(theta);
		double s = std::cos(theta);

		cj << 0,
			  0,
			  0,
			  (-s * x + c * y) * theta,
			  (-c * x - s * y) * theta,
			  0;
	}
	else
	{
		cj = cSpAlg::tSpVec::Zero();
	}
	return cj;
}

cSpAlg::tSpVec cRBDUtil::BuildCjFixed(const Eigen::MatrixXd& joint_mat, const Eigen::VectorXd& q_dot, int j)
{
	return cSpAlg::tSpVec::Zero();
}


void cRBDUtil::BuildBiasForce(const cRBDModel& model, Eigen::VectorXd& out_bias_force)
{
	Eigen::VectorXd acc = Eigen::VectorXd::Zero(model.GetNumDof());
	SolveInvDyna(model, acc, out_bias_force);
}

void cRBDUtil::CalcGravityForce(const cRBDModel& model, Eigen::VectorXd& out_g_force)
{
	const Eigen::MatrixXd& joint_mat = model.GetJointMat();
	const Eigen::MatrixXd& body_defs = model.GetBodyDefs();
	const tVector& gravity = model.GetGravity();
	const Eigen::VectorXd& pose = model.GetPose();

	assert(joint_mat.rows() == body_defs.rows());
	assert(cKinTree::GetNumDof(joint_mat) == pose.rows());

	cSpAlg::tSpVec acc0 = cSpAlg::BuildSV(tVector::Zero(), gravity);

	int num_joints = cKinTree::GetNumJoints(joint_mat);
	Eigen::MatrixXd fs = Eigen::MatrixXd(num_joints, cSpAlg::gSpVecSize);

	for (int j = 0; j < num_joints; ++j)
	{
		if (cKinTree::IsValidBody(body_defs, j))
		{
			cSpAlg::tSpTrans world_child_trans = model.GetSpWorldJointTrans(j);
			cSpAlg::tSpMat I = BuildInertiaSpatialMat(body_defs, j);
			cSpAlg::tSpVec curr_acc = cSpAlg::ApplyTransM(world_child_trans, acc0);
			cSpAlg::tSpVec curr_f = I * curr_acc;
			fs.row(j) = curr_f;
		}
	}

	out_g_force = Eigen::VectorXd::Zero(pose.size());
	for (int j = num_joints - 1; j >= 0; --j)
	{
		if (cKinTree::IsValidBody(body_defs, j))
		{
			cSpAlg::tSpVec curr_f = fs.row(j);
			const auto S = model.GetJointSubspace(j);
			Eigen::VectorXd curr_tau = S.transpose() * curr_f;

			cKinTree::SetJointParams(joint_mat, j, curr_tau, out_g_force);
			if (cKinTree::HasParent(joint_mat, j))
			{
				int parent_id = cKinTree::GetParent(joint_mat, j);
				cSpAlg::tSpTrans child_parent_trans = model.GetSpChildParentTrans(j);
				fs.row(parent_id) += cSpAlg::ApplyTransF(child_parent_trans, curr_f);
			}
		}
	}
}