//
// Created by xtkuang on 2025/7/8.
//

#ifndef ROBOT_H
#define ROBOT_H


#pragma once

#include <Eigen/Core>
#include <iostream>
#include <memory>
#include <queue>
#include <string>
#include <type_traits>
#include <utility>
#include <vector>
#include <iomanip>
#include <limits>

#include "inertial.h"
#include "joint.h"
#include "link.h"
#include "state.h"
#include "utils/math/liegroup.h"
#include "utils/math/qp_solver.h"

namespace cmvr::dyn {
/**********************************************************
 * FORWARD DECLARATION
 **********************************************************/

//
template<int DOF>
class Robot;
class Link;
class Joint;

/**********************************************************
 * MOBILE BASE
 **********************************************************/

enum class MobileBaseType {
    None, //
    Differential, //
    Mecanum //
};

struct MobileBase {
    MobileBaseType type{MobileBaseType::None};
    math::SE3::MatrixType T; // front = x-axis
    std::vector<std::string> joints;
    std::vector<double> params;
};

struct MobileBaseDifferential : public MobileBase {
    unsigned int right_wheel_idx{};
    unsigned int left_wheel_idx{};
    double wheel_base{0.};
    double wheel_radius{0.};
};

struct MobileBaseMecanum : public MobileBase {
    unsigned int fr_wheel_idx{};
    unsigned int fl_wheel_idx{};
    unsigned int rr_wheel_idx{};
    unsigned int rl_wheel_idx{};
    double L_x{0.};
    double L_y{0.};
    double wheel_radius{0.};
};

/**********************************************************
 * ROBOT
 **********************************************************/

struct RobotConfiguration {
    std::string name;
    std::shared_ptr<Link> base_link;
    std::shared_ptr<MobileBase> mobile_base;
};

template<int DOF>
class Robot {
public:
    template<typename T, int N>
    using ContainerType = typename std::conditional_t<(N > 0), std::array<T, (unsigned int) N>, std::vector<T> >;

    struct Link_ {
        struct SubLink_ {
            std::shared_ptr<Link> link;
            Inertial::MatrixType J_wrt_p;
            math::SE3::MatrixType M_wrt_p;
            math::SE3::MatrixType M_wrt_base;
        };

        std::vector<SubLink_> links;
        Inertial::MatrixType J{Inertial::MatrixType::Zero()}; // Link_ inertial wrt base link
        math::SE3::MatrixType M{math::SE3::Identity()}; // Base link to Link_ frame
        Inertial::MatrixType I{Inertial::MatrixType::Zero()}; // Link_ inertial wrt Link_ frame
        int depth{0};
        int parent_joint_idx{-1};
        std::vector<int> child_joint_idx{};

        void SetBaseLink(const std::shared_ptr<Link> &link, const math::SE3::MatrixType &T) {
            // Rest link information
            links.clear();
            J = Inertial::MatrixType::Zero();
            M = T;
            AddLink(link, math::SE3::Identity());
        }

        int AddLink(const std::shared_ptr<Link> &link, const math::SE3::MatrixType &M_wrt_p) {
            Inertial::MatrixType J_wrt_p = Inertial::Transform(M_wrt_p, link->I_);
            int idx = links.size();
            SubLink_ l;
            l.link = link;
            l.J_wrt_p = J_wrt_p;
            l.M_wrt_p = M_wrt_p;
            l.M_wrt_base = M * M_wrt_p;
            links.push_back(l);
            I += J_wrt_p;
            J += Inertial::Transform(M, J_wrt_p);
            return idx;
        }
    };

    struct Joint_ {
        std::shared_ptr<Joint> joint{nullptr};
        math::se3v::MatrixType S{math::se3v::MatrixType::Zero()}; // zero pose

        int parent_link_idx{};
        int child_link_idx{};
    };

    struct LinkIdx_ {
        int link_idx;
        int sub_link_idx;
    };

    explicit Robot(const RobotConfiguration &robot_configuration);
    
    std::shared_ptr<Link> GetBase();
    
    std::vector<std::string> GetLinkNames() const;
    
    std::vector<std::string> GetJointNames() const;
    
    std::shared_ptr<Link> GetLink(const std::string &name) const;
    
    std::shared_ptr<Link> GetLink(std::shared_ptr<State<DOF>> state, int index) const;

    LinkIdx_ GetLinkIdx(const std::string& name) const;

    template<typename LinkContainer = std::vector<std::string>, typename JointContainer = std::vector<std::string> >
    std::shared_ptr<State<DOF>> MakeState(const LinkContainer &link_names, const JointContainer &joint_names);

    int GetDOF() const;

    int GetNumberOfJoints() const;

    void ComputeForwardKinematics(std::shared_ptr<State<DOF>> state);

    void ComputeDiffForwardKinematics(std::shared_ptr<State<DOF>> state);

    void Compute2ndDiffForwardKinematics(std::shared_ptr<State<DOF>> state);

    void ComputeInverseDynamics(std::shared_ptr<State<DOF>> state);

    Eigen::Vector<double, DOF> ComputeGravityTerm(std::shared_ptr<State<DOF>> state);

    Eigen::Matrix<double, DOF, DOF> ComputeMassMatrix(std::shared_ptr<State<DOF>> state);

    Eigen::Matrix<double, 6, 6> ComputeReflectiveInertia(std::shared_ptr<State<DOF>> state, unsigned int from, unsigned int to);

    math::SE3::MatrixType ComputeTransformation(std::shared_ptr<State<DOF>> state, unsigned int from, unsigned int to);

    math::se3v::MatrixType ComputeBodyVelocity(std::shared_ptr<State<DOF>> state, unsigned int from, unsigned int to);

    Eigen::Matrix<double, 6, DOF> ComputeSpaceJacobian(std::shared_ptr<State<DOF>> state, unsigned int from, unsigned int to);

    Eigen::Matrix<double, 6, DOF> ComputeBodyJacobian(std::shared_ptr<State<DOF>> state, unsigned int from, unsigned int to);

    double ComputeMass(std::shared_ptr<State<DOF>> state, unsigned int target_link);

    Eigen::Vector3d ComputeCenterOfMass(std::shared_ptr<State<DOF>> state, unsigned int ref_link, unsigned int target_link);

    Eigen::Vector3d ComputeCenterOfMass(std::shared_ptr<State<DOF>> state, unsigned int ref_link, const std::vector<unsigned int> &target_links);

    Eigen::Matrix<double, 3, DOF> ComputeCenterOfMassJacobian(std::shared_ptr<State<DOF>> state, unsigned int ref_link, unsigned int target_link);

    Inertial::MatrixType ComputeTotalInertial(std::shared_ptr<State<DOF>> state, unsigned int ref_link);

    Eigen::Vector3d ComputeCenterOfMass(std::shared_ptr<State<DOF>> state, unsigned int ref_link);

    Eigen::Matrix<double, 3, DOF> ComputeCenterOfMassJacobian(std::shared_ptr<State<DOF>> state, unsigned int ref_link);

    std::vector<CollisionResult> DetectCollisionsOrNearestLinks(std::shared_ptr<State<DOF>> state, int collision_threshold = 0);

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

    Eigen::Vector<double, DOF> GetLimitQLower(const std::shared_ptr<State<DOF>> &state);

    Eigen::Vector<double, DOF> GetLimitQUpper(const std::shared_ptr<State<DOF>> &state);

    Eigen::Vector<double, DOF> GetLimitQdotLower(const std::shared_ptr<State<DOF>> &state);

    Eigen::Vector<double, DOF> GetLimitQdotUpper(const std::shared_ptr<State<DOF>> &state);

    Eigen::Vector<double, DOF> GetLimitQddotLower(const std::shared_ptr<State<DOF>> &state);

    Eigen::Vector<double, DOF> GetLimitQddotUpper(const std::shared_ptr<State<DOF>> &state);

    Eigen::Vector<double, DOF> GetLimitTorque(const std::shared_ptr<State<DOF>> &state);

    Eigen::Vector<double, DOF> GetJointProperty(const std::shared_ptr<State<DOF>> &state, const std::function<double(std::shared_ptr<Joint>)> &getter);

    void ComputeMobilityInverseDiffKinematics(std::shared_ptr<State<DOF>> state, //
                                              const Eigen::Vector2d &linear_velocity, // (m/s)
                                              double angular_velocity // (rad/s)
    );

    void ComputeMobilityInverseDiffKinematics(std::shared_ptr<State<DOF>> state, //
                                              const math::se2v::MatrixType &body_velocity // w, x, y
    );

    math::se2v::MatrixType ComputeMobilityDiffKinematics( //
        std::shared_ptr<State<DOF>> state //
    );

    static int CountJoints(const std::shared_ptr<Link> &base_link, bool include_fixed=false);

    math::SE3::MatrixType GetLinkT(std::shared_ptr<State<DOF>> state, const LinkIdx_ &idx);

    math::se3v::MatrixType GetLinkV(std::shared_ptr<State<DOF>> state, const LinkIdx_ &idx);

    math::SE3::MatrixType GetJointT(std::shared_ptr<State<DOF>> state, int joint_idx);

    math::se3v::MatrixType GetJointV(std::shared_ptr<State<DOF>> state, int joint_idx);

    math::se3v::MatrixType GetJointVdot(std::shared_ptr<State<DOF>> state, int joint_idx);

    math::SE3::MatrixType GetTransformation(std::shared_ptr<State<DOF>> state, const LinkIdx_ &from, const LinkIdx_ &to);

    Eigen::Matrix<double, 6, DOF> GetSpaceJacobian(std::shared_ptr<State<DOF>> state, const LinkIdx_ &from, const LinkIdx_ &to);

    Eigen::Matrix<double, 6, DOF> GetBodyJacobian(std::shared_ptr<State<DOF>> state, const LinkIdx_ &from, const LinkIdx_ &to);

protected:
    Robot() = default;

    void Build(const RobotConfiguration &rc);

private:
    ContainerType<Link_, DOF + 1> links_;
    ContainerType<Joint_, DOF> joints_;

    std::unordered_map<std::string, LinkIdx_> link_idx_; // link name to (parent idx, sub link idx)
    std::unordered_map<std::string, unsigned int> joint_idx_; // joint name to joint idx

    size_t n_links_{};
    size_t n_joints_{};

    /**
     * MOBILE BASE
     */
    std::shared_ptr<MobileBase> mobile_base_;
};

RobotConfiguration LoadRobotFromURDFData(const std::string &model, const std::string &base_link_name);

RobotConfiguration LoadRobotFromURDF(const std::string &path, const std::string &base_link_name);

inline std::string to_string(GeomType type) {
    switch (type) {
        case GeomType::kCapsule:
            return "capsule";
    }
    return "unknown";
}
}

#include "robot.tpp"

#endif //ROBOT_H