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

#ifndef CMVR_ES_CARTESIAN_CONTROLLER_H
#define CMVR_ES_CARTESIAN_CONTROLLER_H

#pragma once

#include <Eigen/Core>
#include <Eigen/Dense>
#include <vector>
#include <string>
#include <memory>
#include <cmath>
#include <fcl/fcl.h>
#include <optional>
#include <unordered_map>

#include "utils/math/se3.h"
#include "utils/math/qp_solver.h"
#include "utils/dynamics/robot.h"

namespace cmvr::ctrl{
    using cmvr::math::SE3;

    // -----------------------------------------------------------------------------
    // Task & collision descriptions
    // -----------------------------------------------------------------------------
    struct PoseTarget {
        std::string     link_name;    //!< link to be controlled
        Eigen::Matrix4d T_target;     //!< desired pose w.r.t *base_link*
        double          w_posrot = 0.5; // [0→pure pos, 1→pure rot]
        double          weight   = 1.0; // task weight (0‑1 soft, 1 hard)
    };

    struct CollisionPair { std::string link_A, link_B; };

    // -----------------------------------------------------------------------------
    // CartesianController declaration (implementation in .cpp)
    // -----------------------------------------------------------------------------
    template<int DOF>
    class QPSolver {
    public:
        using RobotT = cmvr::dyn::Robot<DOF>;
        using VecD   = Eigen::Vector<double, DOF>;
        using Mat6D  = Eigen::Matrix<double, 6, DOF>;

        enum class Mode { Position, Velocity, Torque };

        QPSolver(std::shared_ptr<RobotT> robot, double dsafe = 0.05, double lambda = 1e-2);
        /* ------------------------- configuration ----------------------------- */
        void setCollisionPairs(const std::vector<CollisionPair>& pairs);

        // Joint limits (size = DOF).  Any value = std::nullopt → no limit.
        void setPositionLimits(const std::vector<std::optional<double>>& q_min, const std::vector<std::optional<double>>& q_max);
        void setVelocityLimits(const std::vector<std::optional<double>>& qd_max);
        void setAccelerationLimits(const std::vector<std::optional<double>>& qdd_max);

        /* --------------------------- main API -------------------------------- */
        /**
         * Compute joint command given Cartesian objectives.
         * @param state   current (mutable) robot state
         * @param base    reference frame (string)
         * @param targets list of PoseTarget
         * @param dt      controller step [s]
         * @param mode    desired output type
         * @param out_cmd filled with q / q̇ / τ
         * @return        true if IK converged (position) or soft OK (others)
         */
        bool compute(
                std::shared_ptr<cmvr::dyn::State<DOF>>& state,
                const std::string& base_link,
                const std::vector<PoseTarget>& targets,
                double dt,
                Mode mode,
                VecD& out_cmd,
                int max_iters = 60,
                double tol = 1e-3,
                double alpha_h = 10.0);

    private:
        /* -------- internal helpers (implemented in .cpp) --------------------- */
        bool solveIK(std::shared_ptr<cmvr::dyn::State<DOF>>& state, const std::string& base, const std::vector<PoseTarget>& targets,
                     const VecD& q_init, VecD& q_out, int max_iters, double tol, double alpha_h) const;

        bool distanceConstraint(const CollisionPair& cp, std::shared_ptr<cmvr::dyn::State<DOF>>& state,
                                const std::vector<std::shared_ptr<fcl::CollisionObjectd>>& objs,
                                double& h_out, Eigen::RowVectorXd& J_out) const;

        inline void clampVec(VecD& v, const std::vector<std::optional<double>>& lo, const std::vector<std::optional<double>>& hi) const;

        /* -------------------------- data ------------------------------------ */
        std::shared_ptr<RobotT> robot_;
        std::vector<CollisionPair> pairs_;
        double d_safe_  = 0.05;
        double lambda_  = 1e-2;
        mutable math::QPSolver solver_;

        // joint limits (optional)
        std::vector<std::optional<double>> q_min_  = std::vector<std::optional<double>>(DOF);
        std::vector<std::optional<double>> q_max_  = std::vector<std::optional<double>>(DOF);
        std::vector<std::optional<double>> qd_max_ = std::vector<std::optional<double>>(DOF);
        std::vector<std::optional<double>> qdd_max_= std::vector<std::optional<double>>(DOF);

        // FCL reusable buffers
        mutable fcl::DistanceRequestd dist_req_{};
        mutable fcl::DistanceResultd dist_res_{};
    };
}

// ------------------------------ explicit instantiation -----------------------
extern template class cmvr::ctrl::QPSolver<7>;
extern template class cmvr::ctrl::QPSolver<14>;
extern template class cmvr::ctrl::QPSolver<20>;


#endif //CMVR_ES_CARTESIAN_CONTROLLER_H