#pragma once

#include "robot.h"
using namespace cmvr::dyn;


template<int DOF>
Robot<DOF>::Robot(const RobotConfiguration &robot_configuration) {
    Build(robot_configuration);
}

template<int DOF>
std::shared_ptr<Link> Robot<DOF>::GetBase() {
    return links_[0].links[0].link;
}

template<int DOF>
std::vector<std::string> Robot<DOF>::GetLinkNames() const {
    std::vector<std::string> names;
    for (const auto &[n, idx]: link_idx_) {
        names.push_back(n);
    }
    return names;
}

template<int DOF>
std::vector<std::string> Robot<DOF>::GetJointNames() const {
    std::vector<std::string> names;
    for (const auto &[n, idx]: joint_idx_) {
        names.push_back(n);
    }
    return names;
}

template<int DOF>
std::shared_ptr<Link> Robot<DOF>::GetLink(const std::string &name) const {
    if (link_idx_.find(name) == link_idx_.end()) {
        return nullptr;
    }
    const auto &idx = link_idx_.at(name);
    return links_[idx.link_idx].links[idx.sub_link_idx].link;
}

template<int DOF>
std::shared_ptr<Link> Robot<DOF>::GetLink(std::shared_ptr<State<DOF> > state, int index) const {
    if (index >= static_cast<int>(state->utr_link_map.size())) {
        return nullptr;
    }
    const auto &idx = state->utr_link_map[index];
    return links_[idx.link_idx].links[idx.sub_link_idx].link;
}

template<int DOF>
typename Robot<DOF>::LinkIdx_ Robot<DOF>::GetLinkIdx(const std::string& name) const {
    auto it = link_idx_.find(name);
    if (it == link_idx_.end()) {
        throw std::runtime_error("Link name not found: " + name);
    }
    return it->second;
}

template<int DOF>
template<typename LinkContainer, typename JointContainer>
std::shared_ptr<State<DOF>> Robot<DOF>::MakeState(const LinkContainer &link_names, const JointContainer &joint_names) {
    static_assert(std::is_same_v<typename LinkContainer::value_type, std::string> ||
                  std::is_same_v<typename LinkContainer::value_type, std::string_view>,
                  "LinkContainer value_type must be std::string");
    static_assert(std::is_same_v<typename JointContainer::value_type, std::string> ||
                  std::is_same_v<typename JointContainer::value_type, std::string_view>,
                  "JointContainer value_type must be std::string");

    std::vector<bool> flag;
    flag.resize(n_joints_);
    std::fill(flag.begin(), flag.end(), false);

    auto state = std::shared_ptr<State<DOF>>(new State < DOF > (n_joints_));
    for (int i = 0; i < static_cast<int>(joint_names.size()); i++) {
        const auto &name = joint_names[i];

        auto it = std::find_if(joints_.begin(), joints_.end(),
                               [name](const auto &j) { return j.joint->name_ == name; });
        if (it == joints_.end()) {
            throw std::runtime_error("Cannot find the joint with name");
        }
        state->utr_joint_map[i] = it - joints_.begin();
        state->joint_names[i] = it->joint->name_;
        flag[state->utr_joint_map[i]] = true;
    }
    for (unsigned int i = joint_names.size(); i < n_joints_; i++) {
        auto it = std::find(flag.begin(), flag.end(), false);
        state->utr_joint_map[i] = it - flag.begin();
        state->joint_names[i] = (joints_.begin() + (it - flag.begin()))->joint->name_;
        *it = true;
    }
    for (int i = 0; i < state->utr_joint_map.size(); i++) {
        state->rtu_joint_map[state->utr_joint_map[i]] = i;
    }

    std::unordered_map<std::string, bool> name_flag;
    state->link_names.clear();
    state->utr_link_map.clear();
    for (const auto &name: link_names) {
        auto it = link_idx_.find(name);
        if (it == link_idx_.end()) {
            throw std::runtime_error("The link with the given name does not exist");
        }
        state->utr_link_map.push_back(it->second);
        state->link_names.push_back(name);
        name_flag[name] = true;
    }
    for (const auto &[k, v]: link_idx_) {
        if (name_flag.find(k) == name_flag.end()) {
            state->utr_link_map.push_back(v);
            state->link_names.push_back(k);
        }
    }
    for (auto i = 0; i < state->link_names.size(); i++) {
        const auto &n = state->link_names[i];
        auto it = link_idx_.find(n);
        if (it == link_idx_.end()) {
            throw std::runtime_error("Fatal error; link names in state should be found in link idx");
        }
        if (it->second.link_idx == 0 && it->second.sub_link_idx == 0) {
            state->base_link_user_idx = i;
        }
    }

    return state;
}

template<int DOF>
int Robot<DOF>::GetDOF() const {
    return n_joints_;
}

template<int DOF>
int Robot<DOF>::GetNumberOfJoints() const {
    return n_joints_;
}

template<int DOF>
void Robot<DOF>::ComputeForwardKinematics(std::shared_ptr<State<DOF>> state) {
    for (auto i = 0; i < n_joints_; i++) {
        state->E[i] = math::SE3::Exp(joints_[i].S, state->q(i));

        int parent_joint_idx = links_[joints_[i].parent_link_idx].parent_joint_idx;
        state->T[i] = GetJointT(state, parent_joint_idx) * state->E[i];

        state->S.col(i) = math::SE3::Ad(GetJointT(state, parent_joint_idx), joints_[i].S);
    }
}

template<int DOF>
void Robot<DOF>::ComputeDiffForwardKinematics(std::shared_ptr<State<DOF>> state) {
    // In Body Frame
    for (auto i = 0; i < n_joints_; i++) {
        state->V.col(i) = joints_[i].S * state->qdot(i);

        int parent_joint_idx = links_[joints_[i].parent_link_idx].parent_joint_idx;
        state->V.col(i) += math::SE3::InvAd(state->E[i], GetJointV(state, parent_joint_idx));
    }
}

template<int DOF>
void Robot<DOF>::Compute2ndDiffForwardKinematics(std::shared_ptr<State<DOF>> state) {
    for (auto i = 0; i < n_joints_; i++) {
        state->Vdot.col(i) = joints_[i].S * state->qddot(i);

        int parent_joint_idx = links_[joints_[i].parent_link_idx].parent_joint_idx;
        state->Vdot.col(i) += math::SE3::InvAd(state->E[i], GetJointVdot(state, parent_joint_idx)) +
                math::SE3::ad(state->V.col(i), joints_[i].S * state->qdot(i));
    }
}

template<int DOF>
void Robot<DOF>::ComputeInverseDynamics(std::shared_ptr<State<DOF>> state) {
    for (int i = n_joints_ - 1; i >= 0; i--) {
        Eigen::Vector<double, 6> F;
        F.setZero();

        for (const auto &child_joint_idx: links_[joints_[i].child_link_idx].child_joint_idx) {
            F += math::SE3::InvAd(state->E[child_joint_idx]).transpose() * state->F.col(child_joint_idx);
        }
        F += links_[joints_[i].child_link_idx].J * state->Vdot.col(i) -
                math::SE3::adTranspose(state->V.col(i), links_[joints_[i].child_link_idx].J * state->V.col(i));

        state->F.col(i) = F;
        state->tau(i) = joints_[i].S.dot(F);
    }
}

template<int DOF>
Eigen::Vector<double, DOF> Robot<DOF>::ComputeGravityTerm(std::shared_ptr<State<DOF>> state) {
    Eigen::Vector<double, DOF> gravity_term;
    gravity_term.resize(n_joints_);
    std::vector<Eigen::Vector<double, 6> > Vdot{}, Fs{};
    Vdot.resize(n_joints_);
    Fs.resize(n_joints_);

    // Calculate Vdot
    for (auto i = 0; i < n_joints_; i++) {
        int parent_joint_idx = links_[joints_[i].parent_link_idx].parent_joint_idx;
        if (parent_joint_idx < 0) {
            Vdot[i] = math::SE3::InvAd(state->E[i], state->Vdot0);
        } else {
            Vdot[i] = math::SE3::InvAd(state->E[i], Vdot[parent_joint_idx]);
        }
    }

    //
    for (int i = n_joints_ - 1; i >= 0; i--) {
        Eigen::Vector<double, 6> F;
        F.setZero();

        for (const auto &child_joint_idx: links_[joints_[i].child_link_idx].child_joint_idx) {
            F += math::SE3::InvAd(state->E[child_joint_idx]).transpose() * Fs[child_joint_idx];
        }
        F += links_[joints_[i].child_link_idx].J * Vdot[i];

        Fs[i] = F;
        gravity_term(state->rtu_joint_map[i]) = joints_[i].S.dot(F);
    }

    return gravity_term;
}

template<int DOF>
Eigen::Matrix<double, DOF, DOF> Robot<DOF>::ComputeMassMatrix(std::shared_ptr<State<DOF>> state) {
    Eigen::Matrix<double, DOF, DOF> M(n_joints_, n_joints_);
    M.setZero();
    std::vector<Eigen::Matrix<double, 6, DOF> > J;
    J.resize(n_joints_);

    for (auto i = 0; i < n_joints_; i++) {
        int parent_joint_idx = links_[joints_[i].parent_link_idx].parent_joint_idx;
        if (parent_joint_idx < 0) {
            J[i].setZero();
            J[i].col(state->rtu_joint_map[i]) = joints_[i].S;
        } else {
            J[i] = math::SE3::Ad(state->E[i].inverse()) * J[parent_joint_idx];
            J[i].col(state->rtu_joint_map[i]) = joints_[i].S;
        }
    }

    for (auto i = 0; i < n_joints_; i++) {
        M += J[i].transpose() * links_[joints_[i].child_link_idx].J * J[i];
    }

    return M;
}

template<int DOF>
Eigen::Matrix<double, 6, 6> Robot<DOF>::ComputeReflectiveInertia(std::shared_ptr<State<DOF>> state, unsigned int from, unsigned int to) {
    if (from >= state->utr_link_map.size() || to >= state->utr_link_map.size()) {
        throw std::runtime_error("Out of range state link");
    }
    const auto &J = ComputeBodyJacobian(state, from, to);
    Eigen::Matrix<double, DOF, DOF> m_inv = ComputeMassMatrix(state).completeOrthogonalDecomposition().
            pseudoInverse();
    return (J * m_inv * J.transpose()).completeOrthogonalDecomposition().pseudoInverse();
}

template<int DOF>
cmvr::math::SE3::MatrixType Robot<DOF>::ComputeTransformation(std::shared_ptr<State<DOF>> state, unsigned int from, unsigned int to) {
    if (from >= state->utr_link_map.size() || to >= state->utr_link_map.size()) {
        throw std::runtime_error("Out of range state link");
    }
    return GetTransformation(state, state->utr_link_map[from], state->utr_link_map[to]);
}

template<int DOF>
cmvr::math::se3v::MatrixType Robot<DOF>::ComputeBodyVelocity(std::shared_ptr<State<DOF>> state, unsigned int from, unsigned int to) {
    if (from >= state->utr_link_map.size() || to >= state->utr_link_map.size()) {
        throw std::runtime_error("Out of range state link");
    }
    math::SE3::MatrixType T_from_to = ComputeTransformation(state, from, to);
    math::se3v::MatrixType V_from = GetLinkV(state, state->utr_link_map[from]);
    math::se3v::MatrixType V_to = GetLinkV(state, state->utr_link_map[to]);
    return V_to - math::SE3::InvAd(T_from_to, V_from);
}

template<int DOF>
Eigen::Matrix<double, 6, DOF> Robot<DOF>::ComputeSpaceJacobian(std::shared_ptr<State<DOF>> state, unsigned int from, unsigned int to) {
    return GetSpaceJacobian(state, state->utr_link_map[from], state->utr_link_map[to]);
}

template<int DOF>
Eigen::Matrix<double, 6, DOF> Robot<DOF>::ComputeBodyJacobian(std::shared_ptr<State<DOF>> state, unsigned int from, unsigned int to) {
    math::SE3::MatrixType T = ComputeTransformation(state, from, to);
    return math::SE3::InvAd(T) * ComputeSpaceJacobian(state, from, to);
}

template<int DOF>
double Robot<DOF>::ComputeMass(std::shared_ptr<State<DOF>> state, unsigned int target_link) {
    const auto &idx = state->utr_link_map[target_link];
    return Inertial::GetMass(links_[idx.link_idx].links[idx.sub_link_idx].link->I_);
}

template<int DOF>
Eigen::Vector3d Robot<DOF>::ComputeCenterOfMass(std::shared_ptr<State<DOF>> state, unsigned int ref_link, unsigned int target_link) {
    const auto &idx = state->utr_link_map[target_link];
    return math::SE3::Multiply(ComputeTransformation(state, ref_link, target_link),
                               Inertial::GetCOM(links_[idx.link_idx].links[idx.sub_link_idx].link->I_));
}

template<int DOF>
Eigen::Vector3d Robot<DOF>::ComputeCenterOfMass(std::shared_ptr<State<DOF>> state, unsigned int ref_link, const std::vector<unsigned int> &target_links) {
    Eigen::Vector3d com{Eigen::Vector3d::Zero()};
    double mass = 0;
    for (auto target_link: target_links) {
        const auto &idx = state->utr_link_map[target_link];
        const auto &p = Inertial::GetCOM(links_[idx.link_idx].links[idx.sub_link_idx].link->I_);
        double m = Inertial::GetMass(links_[idx.link_idx].links[idx.sub_link_idx].link->I_);
        com += math::SE3::Multiply(ComputeTransformation(state, ref_link, target_link), p) * m;
        mass += m;
    }
    return com / mass;
}

template<int DOF>
Eigen::Matrix<double, 3, DOF> Robot<DOF>::ComputeCenterOfMassJacobian(std::shared_ptr<State<DOF>> state,
                                                          unsigned int ref_link,
                                                          unsigned int target_link) {
    using namespace math;
    const auto &idx = state->utr_link_map[target_link];
    return SE3::GetRotation(ComputeTransformation(state, ref_link, target_link)) *
           (SE3::InvAd(SE3::T(Inertial::GetCOM(links_[idx.link_idx].links[idx.sub_link_idx].link->I_))) *
            ComputeBodyJacobian(state, ref_link, target_link))
           .block(3, 0, 3, n_joints_);
}

template<int DOF>
Inertial::MatrixType Robot<DOF>::ComputeTotalInertial(std::shared_ptr<State<DOF>> state, unsigned int ref_link) {
    Inertial::MatrixType I{Inertial::MatrixType::Zero()};
    for (const auto &link: links_) {
        I += Inertial::Transform(GetJointT(state, link.parent_joint_idx), link.J);
    }

    return Inertial::Transform(math::SE3::Inverse(GetLinkT(state, state->utr_link_map[ref_link])), I);
}

template<int DOF>
Eigen::Vector3d Robot<DOF>::ComputeCenterOfMass(std::shared_ptr<State<DOF>> state, unsigned int ref_link) {
    Eigen::Vector3d sum{Eigen::Vector3d::Zero()};
    double mass = 0;
    for (const auto &link: links_) {
        sum += Inertial::GetMass(link.J) *
                math::SE3::Multiply(GetJointT(state, link.parent_joint_idx), Inertial::GetCOM(link.J));
        mass += Inertial::GetMass(link.J);
    }
    return math::SE3::Multiply(math::SE3::Inverse(GetLinkT(state, state->utr_link_map[ref_link])), sum / mass);
}

template<int DOF>
Eigen::Matrix<double, 3, DOF> Robot<DOF>::ComputeCenterOfMassJacobian(std::shared_ptr<State<DOF>> state,
                                                          unsigned int ref_link) {
    using namespace math;
    Eigen::Matrix<double, 3, DOF> J;
    J.resize(3, n_joints_);
    J.setZero();
    double mass = 0;
    for (auto i = 0; i < links_.size(); i++) {
        //      const auto& link = links_[i];
        double m = Inertial::GetMass(links_[i].I);
        J += m * SE3::GetRotation(GetTransformation(state, state->utr_link_map[ref_link], {i, 0})) *
                (SE3::InvAd(SE3::T(Inertial::GetCOM(links_[i].I))) *
                 GetBodyJacobian(state, state->utr_link_map[ref_link], {i, 0}))
                .block(3, 0, 3, n_joints_);
        mass += m;
    }
    return J / mass;
}

template<int DOF>
std::vector<CollisionResult> Robot<DOF>::DetectCollisionsOrNearestLinks(std::shared_ptr<State<DOF>> state, int collision_threshold) {
    std::vector<CollisionResult> dis;
    int n = state->GetLinkNames().size();
    for (int i = 0; i < n; i++) {
        auto link1 = GetLink(state, i);
        if (!link1) {
            throw std::runtime_error("Index error");
        }
        auto link1_T = ComputeTransformation(state, 0, i);
        for (int j = i + 1; j < n; j++) {
            auto link2 = GetLink(state, j);
            if (!link2) {
                throw std::runtime_error("Index error");
            }
            auto link2_T = ComputeTransformation(state, 0, j);

            for (auto link1_col: link1->GetCollisions()) {
                for (auto link2_col: link2->GetCollisions()) {
                    for (auto link1_geom: link1_col->GetGeoms()) {
                        for (auto link2_geom: link2_col->GetGeoms()) {
                            if (link1_geom->GetType() == GeomType::kCapsule && //
                                link2_geom->GetType() == GeomType::kCapsule && //
                                link1_geom->Filter(*link2_geom)) {
                                auto collision_result =
                                        link1_geom->ComputeMinimumDistance(
                                            link1_T * link1_col->GetOrigin(), *link2_geom, //
                                            link2_T * link2_col->GetOrigin());
                                if (collision_result.has_value()) {
                                    auto v = collision_result.value();
                                    v.link1 = link1->GetName();
                                    v.link2 = link2->GetName();
                                    dis.push_back(v);
                                }
                            }
                        }
                    }
                }
            }
        }
    }

    std::sort(dis.begin(), dis.end(),
              [](const CollisionResult &r1, const CollisionResult &r2) { return r1.distance < r2.distance; });

    std::vector<CollisionResult> rv;
    for (const auto &d: dis) {
        if (d.distance > 0 && static_cast<int>(rv.size()) >= collision_threshold) {
            break;
        }
        rv.push_back(d);
    }
    return rv;
}

// TODO: ComputeBodyJacobianDot
// TODO: ComputeBodyAcceleration
// TODO: InverseDiffDynamics

template<int DOF>
Eigen::Vector<double, DOF> Robot<DOF>::GetLimitQLower(const std::shared_ptr<State<DOF>> &state) {
    return GetJointProperty(state, [](auto j) { return j->GetLimitQLower(); });
}

template<int DOF>
Eigen::Vector<double, DOF> Robot<DOF>::GetLimitQUpper(const std::shared_ptr<State<DOF>> &state) {
    return GetJointProperty(state, [](auto j) { return j->GetLimitQUpper(); });
}

template<int DOF>
Eigen::Vector<double, DOF> Robot<DOF>::GetLimitQdotLower(const std::shared_ptr<State<DOF>> &state) {
    return GetJointProperty(state, [](auto j) { return j->GetLimitQdotLower(); });
}

template<int DOF>
Eigen::Vector<double, DOF> Robot<DOF>::GetLimitQdotUpper(const std::shared_ptr<State<DOF>> &state) {
    return GetJointProperty(state, [](auto j) { return j->GetLimitQdotUpper(); });
}

template<int DOF>
Eigen::Vector<double, DOF> Robot<DOF>::GetLimitQddotLower(const std::shared_ptr<State<DOF>> &state) {
    return GetJointProperty(state, [](auto j) { return j->GetLimitQddotLower(); });
}

template<int DOF>
Eigen::Vector<double, DOF> Robot<DOF>::GetLimitQddotUpper(const std::shared_ptr<State<DOF>> &state) {
    return GetJointProperty(state, [](auto j) { return j->GetLimitQddotUpper(); });
}

template<int DOF>
Eigen::Vector<double, DOF> Robot<DOF>::GetLimitTorque(const std::shared_ptr<State<DOF>> &state) {
    return GetJointProperty(state, [](auto j) { return j->GetLimitTorque(); });
}

template<int DOF>
Eigen::Vector<double, DOF> Robot<DOF>::GetJointProperty(const std::shared_ptr<State<DOF>> &state,
    const std::function<double(std::shared_ptr<Joint>)> &getter) {
    Eigen::Vector<double, DOF> prop;
    prop.resize(n_joints_);
    for (auto i = 0; i < n_joints_; i++) {
        prop(state->rtu_joint_map[i]) = getter(joints_[i].joint);
    }
    return prop;
}

template<int DOF>
void Robot<DOF>::ComputeMobilityInverseDiffKinematics(std::shared_ptr<State<DOF>> state, //
    const Eigen::Vector2d &linear_velocity, // (m/s)
    double angular_velocity // (rad/s)
) {
    math::se2v::MatrixType S;
    S(0) = angular_velocity;
    S.tail<2>() = linear_velocity;
    ComputeMobilityInverseDiffKinematics(state, S);
}

template<int DOF>
void Robot<DOF>::ComputeMobilityInverseDiffKinematics(std::shared_ptr<State<DOF>> state, //
    const math::se2v::MatrixType &body_velocity // w, x, y
) {
    math::se3v::MatrixType S{math::se3v::MatrixType::Zero()};
    S.block<3, 1>(2, 0) = body_velocity;

    S = math::SE3::InvAd(mobile_base_->T, S);

    switch (mobile_base_->type) {
        case MobileBaseType::None: {
            break;
        }
        case MobileBaseType::Differential: {
            auto mb = std::static_pointer_cast<MobileBaseDifferential>(mobile_base_);
            double v_right = S(3) + mb->wheel_base / 2 * S(2);
            double v_left = S(3) - mb->wheel_base / 2 * S(2);
            state->qdot(mb->right_wheel_idx) = -v_right / mb->wheel_radius;
            state->qdot(mb->left_wheel_idx) = -v_left / mb->wheel_radius;
            break;
        }
        case MobileBaseType::Mecanum: {
            double v_x = S(3);
            double v_y = S(4);
            double w = S(2);
            // std::cout<<"V_x: "<< v_x<<std::endl;
            // std::cout<<"V_y: "<< v_y<<std::endl;
            // std::cout<<"w: "<< w<<std::endl;

            auto mb = std::static_pointer_cast<MobileBaseMecanum>(mobile_base_);
            double w_1 = 1 / mb->wheel_radius * (v_x - v_y - (mb->L_x + mb->L_y) * w);
            double w_2 = 1 / mb->wheel_radius * (v_x + v_y + (mb->L_x + mb->L_y) * w);
            double w_3 = 1 / mb->wheel_radius * (v_x + v_y - (mb->L_x + mb->L_y) * w);
            double w_4 = 1 / mb->wheel_radius * (v_x - v_y + (mb->L_x + mb->L_y) * w);
            state->qdot(mb->fl_wheel_idx) = w_1;
            state->qdot(mb->fr_wheel_idx) = w_2;
            state->qdot(mb->rl_wheel_idx) = w_3;
            state->qdot(mb->rr_wheel_idx) = w_4;

            // std::cout<<"mb->fl_wheel_idx: "<<mb->fl_wheel_idx<<std::endl;
            // std::cout<<"mb->fr_wheel_idx: "<<mb->fr_wheel_idx<<std::endl;
            // std::cout<<"mb->rl_wheel_idx: "<<mb->rl_wheel_idx<<std::endl;
            // std::cout<<"mb->rr_wheel_idx: "<<mb->rr_wheel_idx<<std::endl;

            break;
        }
    }
}

template<int DOF>
cmvr::math::se2v::MatrixType Robot<DOF>::ComputeMobilityDiffKinematics( //
    std::shared_ptr<State<DOF>> state //
) {
    math::se3v::MatrixType S{math::se3v::MatrixType::Zero()};

    switch (mobile_base_->type) {
        case MobileBaseType::None: {
            break;
        }
        case MobileBaseType::Differential: {
            auto mb = std::static_pointer_cast<MobileBaseDifferential>(mobile_base_);
            double w_r = -state->qdot(mb->right_wheel_idx);
            double w_l = -state->qdot(mb->left_wheel_idx);
            S.block<3, 1>(2, 0) = math::se2v::MatrixType{
                (w_r - w_l) * mb->wheel_radius / mb->wheel_base,
                (w_r + w_l) * mb->wheel_radius / 2, 0
            };
            break;
        }
        case MobileBaseType::Mecanum: {
            auto mb = std::static_pointer_cast<MobileBaseMecanum>(mobile_base_);
            double w_1 = state->qdot(mb->fl_wheel_idx);
            double w_2 = state->qdot(mb->fr_wheel_idx);
            double w_3 = state->qdot(mb->rl_wheel_idx);
            double w_4 = state->qdot(mb->rr_wheel_idx);
            double v_x = mb->wheel_radius / 4 * (w_1 + w_2 + w_3 + w_4);
            double v_y = mb->wheel_radius / 4 * (-w_1 + w_2 + w_3 - w_4);
            double w = mb->wheel_radius / (4 * (mb->L_x + mb->L_y)) * (-w_1 + w_2 - w_3 + w_4);
            S.block<3, 1>(2, 0) = math::se2v::MatrixType{w, v_x, v_y};
            break;
        }
    }
    S = math::SE3::Ad(mobile_base_->T, S);
    return S.block<3, 1>(2, 0);
}

template<int DOF>
int Robot<DOF>::CountJoints(const std::shared_ptr<Link> &base_link, bool include_fixed) {
    int n_joints = 0;
    std::queue<std::shared_ptr<Link> > que;
    que.push(base_link);
    while (!que.empty()) {
        std::shared_ptr<Link> link = que.front();
        que.pop();

        for (const auto &joint: link->GetChildJointList()) {
            que.push(joint->GetChildLink());

            if (joint->IsFixed() && !include_fixed);
            else {
                n_joints++;
            }
        }
    }
    return n_joints;
}

template<int DOF>
void Robot<DOF>::Build(const RobotConfiguration &rc) {
    const auto &base = rc.base_link;
    int n_joints = CountJoints(base);
    if constexpr (DOF < 0) {
        n_links_ = n_joints + 1;
        n_joints_ = n_joints;

        links_.resize(n_links_);
        joints_.resize(n_joints_);
    } else {
        if (n_joints != DOF) {
            throw std::runtime_error("DOF does not match the number of joints");
        }
        n_links_ = DOF + 1;
        n_joints_ = DOF;
    }

    struct QueueItem {
        int parent_joint_idx{};
        int parent_link_idx{};
        std::shared_ptr<Link> link;
        int depth{};
        bool merge_parent_link{};
        math::SE3::MatrixType M_wrt_base;
        math::SE3::MatrixType M_wrt_p;
    };

    std::queue<QueueItem> que;
    int link_idx = 0;
    int joint_idx = 0;
    {
        {
            QueueItem item;
            item.parent_joint_idx = -1;
            item.parent_link_idx = link_idx++;
            item.link = base;
            item.depth = 0;
            item.merge_parent_link = false;
            item.M_wrt_base = math::SE3::Identity();
            item.M_wrt_p = math::SE3::Identity();
            que.push(item);
        }

        while (!que.empty()) {
            auto e = que.front();
            que.pop();

            if (!e.merge_parent_link) {
                links_[e.parent_link_idx].SetBaseLink(e.link, e.M_wrt_base);
                links_[e.parent_link_idx].depth = e.depth;
                links_[e.parent_link_idx].parent_joint_idx = e.parent_joint_idx;
                link_idx_[e.link->name_] = {e.parent_link_idx, 0};
            } else {
                int sub_link_idx = links_[e.parent_link_idx].AddLink(e.link, e.M_wrt_p);
                link_idx_[e.link->name_] = {e.parent_link_idx, sub_link_idx};
            }

            for (const auto &joint: e.link->child_joints_) {
                if (!joint->fixed_) {
                    math::SE3::MatrixType M_wrt_base = e.M_wrt_base * joint->T_pj_;
                    {
                        Joint_ j;
                        j.joint = joint;
                        j.S = math::SE3::Ad(M_wrt_base, joint->S_);
                        j.parent_link_idx = e.parent_link_idx;
                        j.child_link_idx = link_idx;
                        joints_[joint_idx] = j;
                    }
                    joint_idx_[joint->name_] = joint_idx;
                    links_[e.parent_link_idx].child_joint_idx.push_back(joint_idx);
                    {
                        QueueItem item;
                        item.parent_joint_idx = joint_idx++;
                        item.parent_link_idx = link_idx++;
                        item.link = joint->child_link_;
                        item.depth = e.depth + 1;
                        item.merge_parent_link = false;
                        item.M_wrt_base = M_wrt_base;
                        item.M_wrt_p = math::SE3::Identity();
                        que.push(item);
                    }
                } else {
                    QueueItem item;
                    item.parent_joint_idx = e.parent_joint_idx;
                    item.parent_link_idx = e.parent_link_idx;
                    item.link = joint->child_link_;
                    item.depth = e.depth;
                    item.merge_parent_link = true;
                    item.M_wrt_base = e.M_wrt_base * joint->T_pj_ * joint->T_jc_;
                    item.M_wrt_p = e.M_wrt_p * joint->T_pj_ * joint->T_jc_;
                    que.push(item);
                }
            }
        }
    }

    if (rc.mobile_base) {
        switch (rc.mobile_base->type) {
            case MobileBaseType::None: {
                auto mb = std::make_shared<MobileBase>();
                *std::static_pointer_cast<MobileBase>(mb) = *rc.mobile_base;
                mobile_base_ = mb;
                break;
            }
            case MobileBaseType::Differential: {
                auto mb = std::make_shared<MobileBaseDifferential>();
                *std::static_pointer_cast<MobileBase>(mb) = *rc.mobile_base;
                if (rc.mobile_base->joints.size() != 2) {
                    throw std::runtime_error("Differential type mobile should have two joints.");
                }
                if (joint_idx_.find(rc.mobile_base->joints[0]) == joint_idx_.end()) {
                    throw std::runtime_error("Right wheel has invalid parameter.");
                }
                if (joint_idx_.find(rc.mobile_base->joints[1]) == joint_idx_.end()) {
                    throw std::runtime_error("Left wheel has invalid parameter.");
                }
                if (rc.mobile_base->params.size() != 2) {
                    throw std::runtime_error("Differential type mobile should have two parameters.");
                }
                mb->right_wheel_idx = joint_idx_[rc.mobile_base->joints[0]];
                mb->left_wheel_idx = joint_idx_[rc.mobile_base->joints[1]];
                mb->wheel_base = rc.mobile_base->params[0];
                mb->wheel_radius = rc.mobile_base->params[1];
                mobile_base_ = mb;
                break;
            }
            case MobileBaseType::Mecanum:
                auto mb = std::make_shared<MobileBaseMecanum>();
                *std::static_pointer_cast<MobileBase>(mb) = *rc.mobile_base;
                if (rc.mobile_base->joints.size() != 4) {
                    throw std::runtime_error(
                        "Mecanum type mobile should have four joints. (front-right, front-left, rear-right, rear-left)");
                }
                if (joint_idx_.find(rc.mobile_base->joints[0]) == joint_idx_.end()) {
                    throw std::runtime_error("Front-right wheel has invalid parameter.");
                }
                if (joint_idx_.find(rc.mobile_base->joints[1]) == joint_idx_.end()) {
                    throw std::runtime_error("Front-left wheel has invalid parameter.");
                }
                if (joint_idx_.find(rc.mobile_base->joints[2]) == joint_idx_.end()) {
                    throw std::runtime_error("Rear-right wheel has invalid parameter.");
                }
                if (joint_idx_.find(rc.mobile_base->joints[3]) == joint_idx_.end()) {
                    throw std::runtime_error("Rear-left wheel has invalid parameter.");
                }
                if (rc.mobile_base->params.size() != 3) {
                    throw std::runtime_error(
                        "Mecanum type mobile should have three parameters. (Lx, Ly, wheel radius)");
                }
                mb->fr_wheel_idx = joint_idx_[rc.mobile_base->joints[0]];
                mb->fl_wheel_idx = joint_idx_[rc.mobile_base->joints[1]];
                mb->rr_wheel_idx = joint_idx_[rc.mobile_base->joints[2]];
                mb->rl_wheel_idx = joint_idx_[rc.mobile_base->joints[3]];
                mb->L_x = rc.mobile_base->params[0];
                mb->L_y = rc.mobile_base->params[1];
                mb->wheel_radius = rc.mobile_base->params[2];
                mobile_base_ = mb;
                break;
        }
    }
}

template<int DOF>
cmvr::math::SE3::MatrixType Robot<DOF>::GetLinkT(std::shared_ptr<State<DOF>> state, const LinkIdx_ &idx) {
    return GetJointT(state, links_[idx.link_idx].parent_joint_idx) *
           links_[idx.link_idx].links[idx.sub_link_idx].M_wrt_base;
}

template<int DOF>
cmvr::math::se3v::MatrixType Robot<DOF>::GetLinkV(std::shared_ptr<State<DOF>> state, const LinkIdx_ &idx) {
    return math::SE3::InvAd(links_[idx.link_idx].links[idx.sub_link_idx].M_wrt_base,
                            GetJointV(state, links_[idx.link_idx].parent_joint_idx));
}

template<int DOF>
cmvr::math::SE3::MatrixType Robot<DOF>::GetJointT(std::shared_ptr<State<DOF>> state, int joint_idx) {
    if (joint_idx < 0)
        return math::SE3::Identity();
    return state->T[joint_idx];
}

template<int DOF>
cmvr::math::se3v::MatrixType Robot<DOF>::GetJointV(std::shared_ptr<State<DOF>> state, int joint_idx) {
    if (joint_idx < 0)
        return state->V0;
    return state->V.col(joint_idx);
}

template<int DOF>
cmvr::math::se3v::MatrixType Robot<DOF>::GetJointVdot(std::shared_ptr<State<DOF>> state, int joint_idx) {
    if (joint_idx < 0)
        return state->Vdot0;
    return state->Vdot.col(joint_idx);
}

template<int DOF>
cmvr::math::SE3::MatrixType Robot<DOF>::GetTransformation(std::shared_ptr<State<DOF>> state, const LinkIdx_ &from, const LinkIdx_ &to) {
    return math::SE3::Inverse(GetLinkT(state, from)) * GetLinkT(state, to);
}

template<int DOF>
Eigen::Matrix<double, 6, DOF> Robot<DOF>::GetSpaceJacobian(std::shared_ptr<State<DOF>> state, const LinkIdx_ &from, const LinkIdx_ &to) {
    Eigen::Matrix<double, 6, DOF> J;
    J.resize(6, n_joints_);
    J.setZero();

    // TODO: Need to optimize

    int from_link_idx = from.link_idx;
    int to_link_idx = to.link_idx;
    while (from_link_idx != to_link_idx) {
        if (links_[from_link_idx].depth > links_[to_link_idx].depth) {
            int pj = links_[from_link_idx].parent_joint_idx;
            J.col(state->rtu_joint_map[pj]) = -state->S.col(pj);
            from_link_idx = joints_[pj].parent_link_idx;
        } else {
            int pj = links_[to_link_idx].parent_joint_idx;
            J.col(state->rtu_joint_map[pj]) = state->S.col(pj);
            to_link_idx = joints_[pj].parent_link_idx;
        }
    }

    return math::SE3::InvAd(GetLinkT(state, from)) * J;
}

template<int DOF>
Eigen::Matrix<double, 6, DOF> Robot<DOF>::GetBodyJacobian(std::shared_ptr<State<DOF>> state, const LinkIdx_ &from,
    const LinkIdx_ &to) {
    math::SE3::MatrixType T = GetTransformation(state, from, to);
    return math::SE3::InvAd(T) * GetSpaceJacobian(state, from, to);
}