#include <opencv2/opencv.hpp>
#include <vector>
#include <cmath>
#include "svp_opencv.h"

// 预定义的68个3D模型点 (需补全剩余点)
const std::vector<cv::Point3f> model_points_68 = {
    {-73.393523f, -29.801432f, -47.667532f},
    {-72.775014f, -10.949766f, -45.909403f},
    {-70.533638f, 7.929818f, -44.84258f},
    {-66.850058f, 26.07428f, -43.141114f},
    {-59.790187f, 42.56439f, -38.635298f},
    {-48.368973f, 56.48108f, -30.750622f},
    {-34.121101f, 67.246992f, -18.456453f},
    {-17.875411f, 75.056892f, -3.609035f},
    {0.098749f, 77.061286f, 0.881698f},
    {17.477031f, 74.758448f, -5.181201f},
    {32.648966f, 66.929021f, -19.176563f},
    {46.372358f, 56.311389f, -30.77057f},
    {57.34348f, 42.419126f, -37.628629f},
    {64.388482f, 25.45588f, -40.886309f},
    {68.212038f, 6.990805f, -42.281449f},
    {70.486405f, -11.666193f, -44.142567f},
    {71.375822f, -30.365191f, -47.140426f},
    {-61.119406f, -49.361602f, -14.254422f},
    {-51.287588f, -58.769795f, -7.268147f},
    {-37.8048f, -61.996155f, -0.442051f},
    {-24.022754f, -61.033399f, 6.606501f},
    {-11.635713f, -56.686759f, 11.967398f},
    {12.056636f, -57.391033f, 12.051204f},
    {25.106256f, -61.902186f, 7.315098f},
    {38.338588f, -62.777713f, 1.022953f},
    {51.191007f, -59.302347f, -5.349435f},
    {60.053851f, -50.190255f, -11.615746f},
    {0.65394f, -42.19379f, 13.380835f},
    {0.804809f, -30.993721f, 21.150853f},
    {0.992204f, -19.944596f, 29.284036f},
    {1.226783f, -8.414541f, 36.94806f},
    {-14.772472f, 2.598255f, 20.132003f},
    {-7.180239f, 4.751589f, 23.536684f},
    {0.55592f, 6.5629f, 25.944448f},
    {8.272499f, 4.661005f, 23.695741f},
    {15.214351f, 2.643046f, 20.858157f},
    {-46.04729f, -37.471411f, -7.037989f},
    {-37.674688f, -42.73051f, -3.021217f},
    {-27.883856f, -42.711517f, -1.353629f},
    {-19.648268f, -36.754742f, 0.111088f},
    {-28.272965f, -35.134493f, 0.147273f},
    {-38.082418f, -34.919043f, -1.476612f},
    {19.265868f, -37.032306f, 0.665746f},
    {27.894191f, -43.342445f, -0.24766f},
    {37.437529f, -43.110822f, -1.696435f},
    {45.170805f, -38.086515f, -4.894163f},
    {38.196454f, -35.532024f, -0.282961f},
    {28.764989f, -35.484289f, 1.172675f},
    {-28.916267f, 28.612716f, 2.24031f},
    {-17.533194f, 22.172187f, 15.934335f},
    {-6.68459f, 19.029051f, 22.611355f},
    {0.381001f, 20.721118f, 23.748437f},
    {8.375443f, 19.03546f, 22.721995f},
    {18.876618f, 22.394109f, 15.610679f},
    {28.794412f, 28.079924f, 3.217393f},
    {19.057574f, 36.298248f, 14.987997f},
    {8.956375f, 39.634575f, 22.554245f},
    {0.381549f, 40.395647f, 23.591626f},
    {-7.428895f, 39.836405f, 22.406106f},
    {-18.160634f, 36.677899f, 15.121907f},
    {-24.37749f, 28.677771f, 4.785684f},
    {-6.897633f, 25.475976f, 20.893742f},
    {0.340663f, 26.014269f, 22.220479f},
    {8.444722f, 25.326198f, 21.02552f},
    {24.474473f, 28.323008f, 5.712776f},
    {8.449166f, 30.596216f, 20.671489f},
    {0.205322f, 31.408738f, 21.90367f},
    {-7.198266f, 30.844876f, 20.328022f}
};

xmedia_s32 svp_opencv_face_orientation(const xmedia_svp_keypoint *landmarks, svp_euler_angles *angles)
{
    if (landmarks == XMEDIA_NULL || angles == XMEDIA_NULL) {
        printf("svp_opencv_face_orientation param is NULL!\n");
        return XMEDIA_FAILURE;
    }

    // 转换关键点到图像坐标
    std::vector<cv::Point2f> image_points;
    for (xmedia_u8 i = 0; i < MAX_DETECT_KEYPOINT_NUM; i++) {
        image_points.emplace_back(landmarks[i].x, landmarks[i].y);
    }

    // 相机内参矩阵
    xmedia_double center_x = 640 / 2;
    xmedia_double center_y = 360 / 2;
    xmedia_double focal_length = center_x / tan(30 * CV_PI / 180);
    cv::Mat camera_matrix = (cv::Mat_<xmedia_double>(3,3) <<
        focal_length, 0, center_x,
        0, focal_length, center_y,
        0, 0, 1);

    // 畸变系数, 后续需要根据不同sensor设置？
    cv::Mat dist_coeffs = cv::Mat::zeros(4, 1, CV_64F);
    // 初始化猜测
    cv::Mat r_vec = (cv::Mat_<xmedia_double>(3,1) << 0.01891013, 0.08560084, -3.14392813);
    cv::Mat t_vec = (cv::Mat_<xmedia_double>(3,1) << -14.97821226, -10.62040383, -2053.03596872);

    bool success = cv::solvePnP(
        model_points_68,
        image_points,
        camera_matrix,
        dist_coeffs,
        r_vec,
        t_vec,
        true, // 使用外部猜测
        cv::SOLVEPNP_ITERATIVE
    );

    if (!success) {
        printf("solvePnP error\n");
        return XMEDIA_FAILURE;
    }

    // 转换旋转向量为矩阵
    cv::Mat rvec_matrix;
    cv::Rodrigues(r_vec, rvec_matrix);

    // 构建投影矩阵
    cv::Mat proj_matrix(3, 4, CV_64F);
    cv::hconcat(rvec_matrix, t_vec, proj_matrix);

    // 分解投影矩阵求欧拉角, 实现旋转矩阵到欧拉角的转换
    xmedia_double sy = std::sqrt(rvec_matrix.at<xmedia_double>(0,0) * rvec_matrix.at<xmedia_double>(0,0) +
                       rvec_matrix.at<xmedia_double>(1,0) * rvec_matrix.at<xmedia_double>(1,0));

    xmedia_double x, y, z;
    if (!(sy < 1e-6)) {
        x = std::atan2(rvec_matrix.at<xmedia_double>(2,1), rvec_matrix.at<xmedia_double>(2,2));
        y = std::atan2(-rvec_matrix.at<xmedia_double>(2,0), sy);
        z = std::atan2(rvec_matrix.at<xmedia_double>(1,0), rvec_matrix.at<xmedia_double>(0,0));
    } else {
        x = std::atan2(-rvec_matrix.at<xmedia_double>(1,2), rvec_matrix.at<xmedia_double>(1,1));
        y = std::atan2(-rvec_matrix.at<xmedia_double>(2,0), sy);
        z = 0;
    }

    // 原始弧度转角度
    x = x * 180 / CV_PI;
    y = y * 180 / CV_PI;
    z = z * 180 / CV_PI;

    // 角度范围修正
    angles->pitch = std::asin(std::sin(x * CV_PI / 180)) * 180 / CV_PI;
    angles->roll  = -std::asin(std::sin(z * CV_PI / 180)) * 180 / CV_PI;
    angles->yaw   = std::asin(std::sin(y * CV_PI / 180)) * 180 / CV_PI;

    // 确保角度在-90到90度范围内
    angles->pitch = std::max(std::min(angles->pitch, 90.0), -90.0);
    angles->roll  = std::max(std::min(angles->roll, 90.0), -90.0);
    angles->yaw   = std::max(std::min(angles->yaw, 90.0), -90.0);

    return XMEDIA_SUCCESS;
}