#pragma once

#include <array>
#include <string>

#include <Eigen/Dense>
#include <pinocchio/algorithm/frames.hpp>
#include <pinocchio/algorithm/jacobian.hpp>
#include <pinocchio/algorithm/kinematics.hpp>
#include <pinocchio/multibody/data.hpp>
#include <pinocchio/multibody/model.hpp>

namespace r1_vision {

inline Eigen::VectorXd MakePinocchioQ(
    const pinocchio::Model& model,
    const std::array<float, 12>& robot_q)
{
    Eigen::VectorXd q = pinocchio::neutral(model);

    static const std::array<const char*, 12> names = {
        "left_shoulder_pitch_joint",
        "left_shoulder_roll_joint",
        "left_shoulder_yaw_joint",
        "left_elbow_joint",
        "left_wrist_roll_joint",
        "right_shoulder_pitch_joint",
        "right_shoulder_roll_joint",
        "right_shoulder_yaw_joint",
        "right_elbow_joint",
        "right_wrist_roll_joint",
        "head_pitch_joint",
        "head_yaw_joint"
    };

    for (int i = 0; i < 12; ++i) {
        const int jid = model.getJointId(names[i]);
        if (jid == 0) {
            continue;
        }
        q[model.joints[jid].idx_q()] = robot_q[i];
    }

    return q;
}

inline Eigen::Vector3d FramePosition(
    pinocchio::Model& model,
    pinocchio::Data& data,
    const std::array<float, 12>& robot_q,
    const std::string& frame_name)
{
    const Eigen::VectorXd q = MakePinocchioQ(model, robot_q);
    const int frame_id = model.getFrameId(frame_name);
    if (frame_id >= static_cast<int>(model.frames.size())) {
        throw std::runtime_error("Pinocchio frame not found: " + frame_name);
    }

    pinocchio::forwardKinematics(model, data, q);
    pinocchio::updateFramePlacements(model, data);
    return data.oMf[frame_id].translation();
}

inline Eigen::MatrixXd FrameTranslationJacobian(
    pinocchio::Model& model,
    pinocchio::Data& data,
    const std::array<float, 12>& robot_q,
    const std::string& frame_name)
{
    const Eigen::VectorXd q = MakePinocchioQ(model, robot_q);
    const int frame_id = model.getFrameId(frame_name);
    if (frame_id >= static_cast<int>(model.frames.size())) {
        throw std::runtime_error("Pinocchio frame not found: " + frame_name);
    }

    pinocchio::forwardKinematics(model, data, q);
    pinocchio::updateFramePlacements(model, data);

    pinocchio::Data::Matrix6x J6(6, model.nv);
    pinocchio::computeFrameJacobian(
        model,
        data,
        q,
        frame_id,
        pinocchio::LOCAL_WORLD_ALIGNED,
        J6);

    return J6.topRows(3);
}

} // namespace r1_vision