//
// Created by SJQ on 2024/4/26.
//
#include "myrobotics.h"

#define R_COL 3
#define R_ROW 3

#define T_COL 4
#define T_ROW 4


void r2rpy(const mat* R,float rpy[3])
{
    rpy[2] = atan2f(R->pData[2 * R_COL + 1], R->pData[2 * R_COL + 2]);  //roll
    float temp = sqrtf(R->pData[2 * R_COL + 1] * R->pData[2 * R_COL + 1] + R->pData[2 * R_COL + 2] * R->pData[2 * R_COL + 2]);
    rpy[1] = atan2f(-R->pData[2 * R_COL + 0], temp);  //pitch
    rpy[0] = atan2f(R->pData[1 * R_COL + 0], R->pData[0 * R_COL + 0]);  //yaw
}

void rpy2r(float rpy[3],mat* R)
{
    float c[3] = {arm_cos_f32(rpy[0]), arm_cos_f32(rpy[1]), arm_cos_f32(rpy[2])};
    float s[3] = {arm_sin_f32(rpy[0]), arm_sin_f32(rpy[1]), arm_sin_f32(rpy[2])};
    float R_data[9] = {
            c[0] * c[1],                       // R11
            c[0] * s[1] * s[2] - s[0] * c[2],  // R12
            c[0] * s[1] * c[2] + s[0] * s[2],  // R13
            s[0] * c[1],                       // R21
            s[0] * s[1] * s[2] + c[0] * c[2],  // R22
            s[0] * s[1] * c[2] - c[0] * s[2],  // R23
            -s[1],                             // R31
            c[1] * s[2],                       // R32
            c[1] * c[2]                        // R33
    };
    MatInit(R,R_ROW,R_COL);
    memcpy(R->pData,R_data,sizeof(float) * R_COL * R_ROW);
}

void r2quat(const mat*R,float q[4])
{
    float R_trace = R->pData[0 * R_COL + 0] + R->pData[1 * R_COL + 1] + R->pData[2 * R_COL + 2];

    float q_data[4] = {
            0.5f * sqrtf(float_constrain(R_trace, -1, 1) + 1),  // q0=cos(θ/2)
            sign(R->pData[2 * R_COL + 1] - R->pData[1 * R_COL + 2]) * 0.5f *
            sqrtf(float_constrain(R->pData[0 * R_COL + 0] - R->pData[1 * R_COL + 1] - R->pData[2 * R_COL + 2], -1, 1) +
                  1),  // q1=rx*sin(θ/2)
            sign(R->pData[0 * R_COL + 2] - R->pData[2 * R_COL + 0]) * 0.5f *
            sqrtf(float_constrain(-R->pData[0 * R_COL + 0] + R->pData[1 * R_COL + 1] - R->pData[2 * R_COL + 2], -1, 1) +
                  1),  // q2=ry*sin(θ/2)
            sign(R->pData[1 * R_COL + 0] - R->pData[0 * R_COL + 1]) * 0.5f *
            sqrtf(float_constrain(-R->pData[0 * R_COL + 0] - R->pData[1 * R_COL + 1] + R->pData[2 * R_COL + 2], -1, 1) +
                  1),  // q3=rz*sin(θ/2)
    };
    memcpy(q,q_data,sizeof(float) * 4);
}

void quat2r(float q[4],mat* R)
{
    float R_data[9] = {
            1 - 2.0f * (q[2] * q[2] + q[3] * q[3]),  // R11
            2.0f * (q[1] * q[2] - q[0] * q[3]),      // R12
            2.0f * (q[1] * q[3] + q[0] * q[2]),      // R13
            2.0f * (q[1] * q[2] + q[0] * q[3]),      // R21
            1 - 2.0f * (q[1] * q[1] + q[3] * q[3]),  // R22
            2.0f * (q[2] * q[3] - q[0] * q[1]),      // R23
            2.0f * (q[1] * q[3] - q[0] * q[2]),      // R31
            2.0f * (q[2] * q[3] + q[0] * q[1]),      // R32
            1 - 2.0f * (q[1] * q[1] + q[2] * q[2])   // R33
    };
    memcpy(R->pData,R_data,sizeof(float) * R_COL * R_ROW);
}

void quat2rpy(const float q[4],float rpy[3])
{
    float rpy_data[3] = {
            atan2f(2.0f * (q[1] * q[2] + q[0] * q[3]),
                   1 - 2.0f * (q[2] * q[2] + q[3] * q[3])),  // yaw
            asinf(2.0f * (q[0] * q[2] - q[1] * q[3])),       // pitch
            atan2f(2.0f * (q[2] * q[3] + q[0] * q[1]),
                   1 - 2.0f * (q[1] * q[1] + q[2] * q[2]))  // rol
    };
    memcpy(rpy,rpy_data,sizeof(float) * 3);
}

void rpy2quat(const float rpy[3],float q[4])
{
    float c[3] = {cosf(0.5f * rpy[0]), cosf(0.5f * rpy[1]),
                  cosf(0.5f * rpy[2])};  // cos(*/2)
    float s[3] = {sinf(0.5f * rpy[0]), sinf(0.5f * rpy[1]),
                  sinf(0.5f * rpy[2])};  // sin(*/2)
    float q_data[4] = {
            c[0] * c[1] * c[2] + s[0] * s[1] * s[2],  // q0=cos(θ/2)
            c[0] * c[1] * s[2] - s[0] * s[1] * c[2],  // q1=rx*sin(θ/2)
            c[0] * s[1] * c[2] + s[0] * c[1] * s[2],  // q2=ry*sin(θ/2)
            s[0] * c[1] * c[2] - c[0] * s[1] * s[2]   // q3=rz*sin(θ/2)
    };
    memcpy(q,q_data,sizeof(float) * 4);
}

void r2t(const mat* R,mat* T)
{
    float T_data[16] = {R->pData[0], R->pData[1], R->pData[2], 0,
                        R->pData[3], R->pData[4], R->pData[5], 0,
                        R->pData[6], R->pData[7], R->pData[8],0,
                        0,          0,          0,          1};
    MatInit(T,T_ROW,T_COL);
    memcpy(T->pData,T_data,sizeof(float) * T_COL * T_ROW);
}

void rp2t(const mat* R,const float p[3],mat* T)
{
    float T_data[16] = {R->pData[0], R->pData[1], R->pData[2], p[0],
                        R->pData[3], R->pData[4], R->pData[5], p[1],
                        R->pData[6], R->pData[7], R->pData[8],p[2],
                        0,          0,          0,          1};
    MatInit(T,T_ROW,T_COL);
    memcpy(T->pData,T_data,sizeof(float) * T_COL * T_ROW);
}