#include <chrono>
#include <numeric>

#include <omp.h>
#include <opencv2/opencv.hpp>

#include "camera.hpp"
#include "triangulator.hpp"

// =================================================================================================
// =================================================================================================

[[maybe_unused]] static void print_2d_mat(const cv::Mat &mat)
{
    // Ensure the matrix is 2D
    if (mat.dims != 2)
    {
        std::cerr << "Error: The matrix is not 2D." << std::endl;
        return;
    }

    // Retrieve matrix dimensions
    int rows = mat.rows;
    int cols = mat.cols;

    // Print the matrix in a NumPy-like style
    std::cout << "cv::Mat('shape': (" << rows << ", " << cols << ")";
    std::cout << ", 'data': [" << std::endl;

    for (int i = 0; i < rows; ++i)
    {
        std::cout << "  [";
        for (int j = 0; j < cols; ++j)
        {
            std::cout << std::fixed << std::setprecision(3) << mat.at<float>(i, j);
            if (j < cols - 1)
            {
                std::cout << ", ";
            }
        }
        std::cout << "]";
        if (i < rows - 1)
        {
            std::cout << "," << std::endl;
        }
    }
    std::cout << "])" << std::endl;
}

// =================================================================================================

[[maybe_unused]] static void print_allpairs(
    const std::vector<std::pair<std::tuple<int, int, int, int>, std::pair<int, int>>> &all_pairs)
{
    std::cout << "All pairs: [" << std::endl;
    for (const auto &pair : all_pairs)
    {
        const auto &tuple_part = pair.first;
        const auto &pair_part = pair.second;

        // Print the tuple part
        std::cout << "("
                  << std::get<0>(tuple_part) << ", "
                  << std::get<1>(tuple_part) << ", "
                  << std::get<2>(tuple_part) << ", "
                  << std::get<3>(tuple_part) << ")";

        // Print the pair part
        std::cout << ", ("
                  << pair_part.first << ", "
                  << pair_part.second << ")"
                  << std::endl;
    }
    std::cout << "]" << std::endl;
}

// =================================================================================================
// =================================================================================================

CameraInternal::CameraInternal(const Camera &cam)
{
    this->cam = cam;

    // Convert camera matrix to cv::Mat for OpenCV
    K = cv::Mat(3, 3, CV_32FC1, const_cast<float *>(&cam.K[0][0])).clone();
    DC = cv::Mat(cam.DC.size(), 1, CV_32FC1, const_cast<float *>(cam.DC.data())).clone();
    R = cv::Mat(3, 3, CV_32FC1, const_cast<float *>(&cam.R[0][0])).clone();
    T = cv::Mat(3, 1, CV_32FC1, const_cast<float *>(&cam.T[0][0])).clone();
}

// =================================================================================================

void CameraInternal::update_projection_matrix()
{
    // Calculate opencv-style projection matrix
    cv::Mat Tr, RT;
    Tr = R * (T * -1);
    cv::hconcat(R, Tr, RT);
    P = K * RT;
}

// =================================================================================================
// =================================================================================================

TriangulatorInternal::TriangulatorInternal(float min_score, size_t min_group_size)
{
    this->min_score = min_score;
    this->min_group_size = min_group_size;
}

// =================================================================================================

std::vector<std::vector<std::array<float, 4>>> TriangulatorInternal::triangulate_poses(
    const std::vector<std::vector<std::vector<std::array<float, 3>>>> &poses_2d,
    const std::vector<Camera> &cameras,
    const std::array<std::array<float, 3>, 2> &roomparams,
    const std::vector<std::string> &joint_names)
{
    auto ttime = std::chrono::high_resolution_clock::now();
    auto stime = std::chrono::high_resolution_clock::now();
    std::chrono::duration<float> elapsed;

    // Check inputs
    if (poses_2d.empty())
    {
        throw std::invalid_argument("No 2D poses provided.");
    }
    if (poses_2d.size() != cameras.size())
    {
        throw std::invalid_argument("Number of cameras and 2D poses must be the same.");
    }
    if (joint_names.size() != poses_2d[0][0].size())
    {
        throw std::invalid_argument("Number of joint names and 2D poses must be the same.");
    }
    for (const auto &joint : core_joints)
    {
        if (std::find(joint_names.begin(), joint_names.end(), joint) == joint_names.end())
        {
            throw std::invalid_argument("Core joint '" + joint + "' not found in joint names.");
        }
    }

    // Convert inputs to internal formats
    std::vector<std::vector<cv::Mat>> poses_2d_mats;
    poses_2d_mats.reserve(cameras.size());
    for (size_t i = 0; i < cameras.size(); ++i)
    {
        size_t num_persons = poses_2d[i].size();
        size_t num_joints = poses_2d[i][0].size();
        std::vector<cv::Mat> camera_poses;
        camera_poses.reserve(num_persons);

        for (size_t j = 0; j < num_persons; ++j)
        {
            std::vector<int> dims = {(int)num_joints, 3};
            cv::Mat pose_mat = cv::Mat(dims, CV_32F);

            // Use pointer to copy data efficiently
            for (size_t k = 0; k < num_joints; ++k)
            {
                float *mat_ptr = pose_mat.ptr<float>(k);
                const auto &joint = poses_2d[i][j][k];
                mat_ptr[0] = joint[0];
                mat_ptr[1] = joint[1];
                mat_ptr[2] = joint[2];
            }
            camera_poses.push_back(std::move(pose_mat));
        }
        poses_2d_mats.push_back(std::move(camera_poses));
    }
    std::vector<CameraInternal> internal_cameras;
    for (size_t i = 0; i < cameras.size(); ++i)
    {
        internal_cameras.emplace_back(cameras[i]);
    }
    std::vector<size_t> core_joint_idx;
    for (const auto &joint : core_joints)
    {
        auto it = std::find(joint_names.begin(), joint_names.end(), joint);
        core_joint_idx.push_back(std::distance(joint_names.begin(), it));
    }
    std::vector<std::array<size_t, 2>> core_limbs_idx;
    for (const auto &limb : core_limbs)
    {
        auto it1 = std::find(core_joints.begin(), core_joints.end(), limb.first);
        auto it2 = std::find(core_joints.begin(), core_joints.end(), limb.second);
        core_limbs_idx.push_back({(size_t)std::distance(core_joints.begin(), it1),
                                  (size_t)std::distance(core_joints.begin(), it2)});
    }

    elapsed = std::chrono::high_resolution_clock::now() - stime;
    init_time += elapsed.count();
    stime = std::chrono::high_resolution_clock::now();

    // Undistort 2D poses
    #pragma omp parallel for
    for (size_t i = 0; i < cameras.size(); ++i)
    {
        undistort_poses(poses_2d_mats[i], internal_cameras[i]);
        internal_cameras[i].update_projection_matrix();
    }

    elapsed = std::chrono::high_resolution_clock::now() - stime;
    undistort_time += elapsed.count();
    stime = std::chrono::high_resolution_clock::now();

    // Project last 3D poses to 2D
    std::vector<std::tuple<std::vector<cv::Mat>, std::vector<cv::Mat>>> last_poses_2d;
    if (!last_poses_3d.empty())
    {
        // Select core joints
        std::vector<cv::Mat> last_core_poses;
        last_core_poses.resize(last_poses_3d.size());
        #pragma omp parallel for
        for (size_t i = 0; i < last_poses_3d.size(); ++i)
        {
            cv::Mat &pose = last_poses_3d[i];
            std::vector<int> dims = {(int)core_joint_idx.size(), 4};
            cv::Mat last_poses_core(dims, pose.type());

            for (size_t j = 0; j < core_joint_idx.size(); ++j)
            {
                pose.row(core_joint_idx[j]).copyTo(last_poses_core.row(j));
            }
            last_core_poses[i] = last_poses_core;
        }

        // Project
        last_poses_2d.resize(cameras.size());
        #pragma omp parallel for
        for (size_t i = 0; i < cameras.size(); ++i)
        {
            auto [poses2d, dists] = project_poses(last_core_poses, internal_cameras[i], true);
            last_poses_2d[i] = std::make_tuple(poses2d, dists);
        }
    }

    elapsed = std::chrono::high_resolution_clock::now() - stime;
    project_time += elapsed.count();
    stime = std::chrono::high_resolution_clock::now();

    // Check matches to old poses
    float threshold = min_score - 0.2;
    std::map<size_t, std::map<size_t, std::vector<size_t>>> scored_pasts;
    if (!last_poses_3d.empty())
    {
        // Calculate index pairs and initialize vectors
        std::vector<std::array<size_t, 3>> indices_ijk;
        for (size_t i = 0; i < cameras.size(); ++i)
        {
            size_t num_last_persons = std::get<0>(last_poses_2d[i]).size();
            scored_pasts[i] = std::map<size_t, std::vector<size_t>>();

            for (size_t j = 0; j < num_last_persons; ++j)
            {
                size_t num_new_persons = poses_2d_mats[i].size();
                scored_pasts[i][j] = std::vector<size_t>();

                for (size_t k = 0; k < num_new_persons; ++k)
                {
                    indices_ijk.push_back({i, j, k});
                }
            }
        }
        std::vector<std::array<size_t, 2>> indices_ik;
        for (size_t i = 0; i < cameras.size(); ++i)
        {
            size_t num_new_persons = poses_2d_mats[i].size();
            for (size_t k = 0; k < num_new_persons; ++k)
            {
                indices_ik.push_back({i, k});
            }
        }

        // Precalculate core poses
        std::vector<cv::Mat> poses_2d_mats_core_list;
        poses_2d_mats_core_list.resize(indices_ik.size());
        std::vector<std::vector<size_t>> mats_core_map;
        mats_core_map.resize(cameras.size());
        for (size_t i = 0; i < cameras.size(); ++i)
        {
            size_t num_new_persons = poses_2d_mats[i].size();
            for (size_t k = 0; k < num_new_persons; ++k)
            {
                mats_core_map[i].push_back(0);
            }
        }
        #pragma omp parallel for
        for (size_t e = 0; e < indices_ik.size(); ++e)
        {
            const auto [i, k] = indices_ik[e];

            const cv::Mat &pose = poses_2d_mats[i][k];
            std::vector<int> dims = {(int)core_joint_idx.size(), 3};
            cv::Mat pose_core(dims, pose.type());

            for (size_t j = 0; j < core_joint_idx.size(); ++j)
            {
                pose.row(core_joint_idx[j]).copyTo(pose_core.row(j));
            }

            poses_2d_mats_core_list[e] = pose_core;
            mats_core_map[i][k] = e;
        }

        // Calculate matching score
        #pragma omp parallel for
        for (size_t e = 0; e < indices_ijk.size(); ++e)
        {
            const auto [i, j, k] = indices_ijk[e];

            const cv::Mat &last_pose = std::get<0>(last_poses_2d[i])[j];
            const cv::Mat &last_dist = std::get<1>(last_poses_2d[i])[j];
            const cv::Mat &new_pose = poses_2d_mats_core_list[mats_core_map[i][k]];

            float score = calc_pose_score(new_pose, last_pose, last_dist, internal_cameras[i]);
            if (score > threshold)
            {
                #pragma omp critical
                {
                    scored_pasts[i][j].push_back(k);
                }
            }
        }
    }

    elapsed = std::chrono::high_resolution_clock::now() - stime;
    match_time += elapsed.count();
    stime = std::chrono::high_resolution_clock::now();

    // Create pairs of persons
    // Checks if the person was already matched to the last frame and if so only creates pairs
    // with those, else it creates all possible pairs
    std::vector<int> num_persons_sum;
    for (size_t i = 0; i < cameras.size(); ++i)
    {
        int nsum = poses_2d[i].size();
        if (i > 0)
        {
            nsum += num_persons_sum[i - 1];
        }
        num_persons_sum.push_back(nsum);
    }
    std::vector<std::pair<std::tuple<int, int, int, int>, std::pair<int, int>>> all_pairs;
    std::vector<std::array<size_t, 4>> indices;
    for (size_t i = 0; i < cameras.size(); ++i)
    {
        for (size_t j = i + 1; j < cameras.size(); ++j)
        {
            for (size_t k = 0; k < poses_2d[i].size(); ++k)
            {
                for (size_t l = 0; l < poses_2d[j].size(); ++l)
                {
                    indices.push_back({i, j, k, l});
                }
            }
        }
    }
    #pragma omp parallel for ordered schedule(dynamic)
    for (size_t e = 0; e < indices.size(); ++e)
    {
        const auto [i, j, k, l] = indices[e];

        int pid1 = num_persons_sum[i] + k;
        int pid2 = num_persons_sum[k] + l;
        bool match = false;

        if (!last_poses_3d.empty())
        {
            for (size_t m = 0; m < last_poses_3d.size(); ++m)
            {
                auto &smi = scored_pasts[i][m];
                auto &smj = scored_pasts[j][m];
                bool in_smi = std::find(smi.begin(), smi.end(), k) != smi.end();
                bool in_smj = std::find(smj.begin(), smj.end(), l) != smj.end();

                if (in_smi && in_smj)
                {
                    match = true;
                    auto item = std::make_pair(
                        std::make_tuple(i, j, k, l), std::make_pair(pid1, pid2));

                    #pragma omp ordered
                    all_pairs.push_back(item);
                }
                else if (in_smi || in_smj)
                {
                    match = true;
                }
            }
        }

        if (!match)
        {
            auto item = std::make_pair(
                std::make_tuple(i, j, k, l), std::make_pair(pid1, pid2));

            // Needed to prevent randomized grouping/merging with slightly different results
            #pragma omp ordered
            all_pairs.push_back(item);
        }
    }

    elapsed = std::chrono::high_resolution_clock::now() - stime;
    pairs_time += elapsed.count();
    stime = std::chrono::high_resolution_clock::now();

    // Calculate pair scores
    std::vector<std::pair<cv::Mat, float>> all_scored_poses;
    all_scored_poses.resize(all_pairs.size());
    #pragma omp parallel for
    for (size_t i = 0; i < all_pairs.size(); ++i)
    {
        const auto &pids = all_pairs[i].first;

        // Extract camera parameters
        const auto &cam1 = internal_cameras[std::get<0>(pids)];
        const auto &cam2 = internal_cameras[std::get<1>(pids)];

        // Extract 2D poses
        const cv::Mat &pose1 = poses_2d_mats[std::get<0>(pids)][std::get<2>(pids)];
        const cv::Mat &pose2 = poses_2d_mats[std::get<1>(pids)][std::get<3>(pids)];

        // Select core joints
        std::vector<int> dims = {(int)core_joint_idx.size(), 3};
        cv::Mat pose1_core(dims, pose1.type());
        cv::Mat pose2_core(dims, pose2.type());
        for (size_t j = 0; j < core_joint_idx.size(); ++j)
        {
            size_t idx = core_joint_idx[j];
            pose1.row(idx).copyTo(pose1_core.row(j));
            pose2.row(idx).copyTo(pose2_core.row(j));
        }

        // Triangulate and score
        auto [pose3d, score] = triangulate_and_score(
            pose1_core, pose2_core, cam1, cam2, roomparams, core_limbs_idx);
        all_scored_poses[i] = std::make_pair(pose3d, score);
    }

    elapsed = std::chrono::high_resolution_clock::now() - stime;
    pair_scoring_time += elapsed.count();
    stime = std::chrono::high_resolution_clock::now();

    // Drop low scoring poses
    size_t num_poses = all_scored_poses.size();
    for (size_t i = num_poses; i > 0; --i)
    {
        if (all_scored_poses[i - 1].second < min_score)
        {
            all_scored_poses.erase(all_scored_poses.begin() + i - 1);
            all_pairs.erase(all_pairs.begin() + i - 1);
        }
    }

    // Group pairs that share a person
    std::vector<std::tuple<cv::Point3d, cv::Mat, std::vector<int>>> groups;
    groups = calc_grouping(all_pairs, all_scored_poses, min_score);

    // Drop groups with too few matches
    size_t num_groups = groups.size();
    for (size_t i = num_groups; i > 0; --i)
    {
        if (std::get<2>(groups[i - 1]).size() < this->min_group_size)
        {
            groups.erase(groups.begin() + i - 1);
        }
    }

    elapsed = std::chrono::high_resolution_clock::now() - stime;
    grouping_time += elapsed.count();
    stime = std::chrono::high_resolution_clock::now();

    // Calculate full 3D poses
    std::vector<cv::Mat> all_full_poses;
    all_full_poses.resize(all_pairs.size());
    #pragma omp parallel for
    for (size_t i = 0; i < all_pairs.size(); ++i)
    {
        const auto &pids = all_pairs[i].first;
        const auto &cam1 = internal_cameras[std::get<0>(pids)];
        const auto &cam2 = internal_cameras[std::get<1>(pids)];
        const auto &pose1 = poses_2d_mats[std::get<0>(pids)][std::get<2>(pids)];
        const auto &pose2 = poses_2d_mats[std::get<1>(pids)][std::get<3>(pids)];

        auto [pose3d, score] = triangulate_and_score(
            pose1, pose2, cam1, cam2, roomparams, {});
        all_full_poses[i] = (pose3d);
    }

    elapsed = std::chrono::high_resolution_clock::now() - stime;
    full_time += elapsed.count();
    stime = std::chrono::high_resolution_clock::now();

    // Merge groups
    std::vector<cv::Mat> all_merged_poses;
    all_merged_poses.resize(groups.size());
    #pragma omp parallel for
    for (size_t i = 0; i < groups.size(); ++i)
    {
        const auto &group = groups[i];
        std::vector<cv::Mat> poses;
        poses.reserve(std::get<2>(group).size());

        for (const auto &idx : std::get<2>(group))
        {
            poses.push_back(all_full_poses[idx]);
        }

        auto merged_pose = merge_group(poses, min_score);
        all_merged_poses[i] = (merged_pose);
    }

    elapsed = std::chrono::high_resolution_clock::now() - stime;
    merge_time += elapsed.count();
    stime = std::chrono::high_resolution_clock::now();

    // Run post-processing steps
    std::vector<cv::Mat> final_poses_3d = all_merged_poses;
    add_extra_joints(final_poses_3d, joint_names);
    filter_poses(final_poses_3d, roomparams, core_joint_idx, core_limbs_idx);
    add_missing_joints(final_poses_3d, joint_names);
    last_poses_3d = final_poses_3d;

    elapsed = std::chrono::high_resolution_clock::now() - stime;
    post_time += elapsed.count();
    stime = std::chrono::high_resolution_clock::now();

    // Convert to output format
    std::vector<std::vector<std::array<float, 4>>> poses_3d;
    poses_3d.reserve(final_poses_3d.size());
    for (size_t i = 0; i < final_poses_3d.size(); ++i)
    {
        const auto &mat = final_poses_3d[i];
        std::vector<std::array<float, 4>> pose;
        size_t num_joints = mat.rows;
        pose.reserve(num_joints);
        size_t num_valid = 0;

        for (size_t j = 0; j < num_joints; ++j)
        {
            const float *mat_ptr = mat.ptr<float>(j);
            std::array<float, 4> point;
            for (size_t k = 0; k < 4; ++k)
            {
                point[k] = mat_ptr[k];
            }
            pose.push_back(point);

            if (point[3] > min_score)
            {
                num_valid++;
            }
        }

        if (num_valid > 0)
        {
            poses_3d.push_back(std::move(pose));
        }
    }

    elapsed = std::chrono::high_resolution_clock::now() - stime;
    convert_time += elapsed.count();

    elapsed = std::chrono::high_resolution_clock::now() - ttime;
    total_time += elapsed.count();
    num_calls++;

    return poses_3d;
}

// =================================================================================================

void TriangulatorInternal::reset()
{
    last_poses_3d.clear();
}

// =================================================================================================

void TriangulatorInternal::print_stats()
{
    std::cout << "Triangulator statistics:" << std::endl;
    std::cout << " Number of calls: " << num_calls << std::endl;
    std::cout << " Init time: " << init_time / num_calls << std::endl;
    std::cout << " Undistort time: " << undistort_time / num_calls << std::endl;
    std::cout << " Project time: " << project_time / num_calls << std::endl;
    std::cout << " Match time: " << match_time / num_calls << std::endl;
    std::cout << " Pairs time: " << pairs_time / num_calls << std::endl;
    std::cout << " Pair scoring time: " << pair_scoring_time / num_calls << std::endl;
    std::cout << " Grouping time: " << grouping_time / num_calls << std::endl;
    std::cout << " Full time: " << full_time / num_calls << std::endl;
    std::cout << " Merge time: " << merge_time / num_calls << std::endl;
    std::cout << " Post time: " << post_time / num_calls << std::endl;
    std::cout << " Convert time: " << convert_time / num_calls << std::endl;
    std::cout << " Total time: " << total_time / num_calls << std::endl;
}

// =================================================================================================

void TriangulatorInternal::undistort_poses(std::vector<cv::Mat> &poses, CameraInternal &icam)
{
    int width = icam.cam.width;
    int height = icam.cam.height;

    // Undistort camera matrix
    cv::Mat newK = cv::getOptimalNewCameraMatrix(
        icam.K, icam.DC, cv::Size(width, height), 1, cv::Size(width, height));

    for (size_t p = 0; p < poses.size(); ++p)
    {
        // Extract the (x, y) coordinates
        cv::Mat points = poses[p].colRange(0, 2).clone();
        points = points.reshape(2);

        // Undistort the points
        cv::undistortPoints(points, points, icam.K, icam.DC, cv::noArray(), newK);

        // Update the original poses with the undistorted points
        points = points.reshape(1);
        points.copyTo(poses[p].colRange(0, 2));

        // Mask out points that are far outside the image (points slightly outside are still valid)
        float mask_offset = (width + height) / 40.0;
        int num_joints = poses[p].rows;
        for (int j = 0; j < num_joints; ++j)
        {
            float *poses_ptr = poses[p].ptr<float>(j);
            float x = poses_ptr[0];
            float y = poses_ptr[1];
            bool in_x = x >= -mask_offset && x < width + mask_offset;
            bool in_y = y >= -mask_offset && y < height + mask_offset;
            if (!in_x || !in_y)
            {
                poses_ptr[0] = 0.0;
                poses_ptr[1] = 0.0;
                poses_ptr[2] = 0.0;
            }
        }
    }

    // Update the camera matrix
    icam.K = newK.clone();
    icam.DC = cv::Mat::zeros(5, 1, CV_32F);
}

// =================================================================================================

std::tuple<std::vector<cv::Mat>, std::vector<cv::Mat>> TriangulatorInternal::project_poses(
    const std::vector<cv::Mat> &bodies3D, const CameraInternal &icam, bool calc_dists)
{
    size_t num_persons = bodies3D.size();
    size_t num_joints = bodies3D[0].rows;

    std::vector<cv::Mat> bodies2D_list(num_persons);
    std::vector<cv::Mat> dists_list(num_persons);

    cv::Mat T_repeated = cv::repeat(icam.T, 1, num_joints).t();
    cv::Mat R_transposed = icam.R.t();

    for (size_t i = 0; i < num_persons; ++i)
    {
        const cv::Mat &body3D = bodies3D[i];

        // Extract coordinates
        const cv::Mat points3d = body3D.colRange(0, 3);

        // Project from world to camera coordinate system
        cv::Mat xyz = (points3d - T_repeated) * R_transposed;

        // Set points behind the camera to zero
        for (size_t j = 0; j < num_joints; ++j)
        {
            float *xyz_row_ptr = xyz.ptr<float>(j);
            float z = xyz_row_ptr[2];
            if (z <= 0)
            {
                xyz_row_ptr[0] = 0.0;
                xyz_row_ptr[1] = 0.0;
                xyz_row_ptr[2] = 0.0;
            }
        }

        // Calculate distance from camera center if required
        cv::Mat dists;
        if (calc_dists)
        {
            cv::multiply(xyz, xyz, dists);
            cv::reduce(dists, dists, 1, cv::REDUCE_SUM, CV_32F);
            cv::sqrt(dists, dists);
        }
        else
        {
            dists = cv::Mat::zeros(num_joints, 1, CV_32F);
        }

        // Project points to image plane
        cv::Mat uv;
        if (icam.cam.type == "fisheye")
        {
        }
        else
        {
            cv::Mat DCc = icam.DC.rowRange(0, 5);
            cv::projectPoints(
                xyz, cv::Mat::zeros(3, 1, CV_32F), cv::Mat::zeros(3, 1, CV_32F), icam.K, DCc, uv);
        }
        uv = uv.reshape(1, {xyz.rows, 2});

        // Add scores again
        std::vector<int> dimsB = {(int)num_joints, 3};
        cv::Mat bodies2D = cv::Mat(dimsB, CV_32F);
        for (size_t j = 0; j < num_joints; ++j)
        {
            float *bodies2D_row_ptr = bodies2D.ptr<float>(j);
            const float *uv_row_ptr = uv.ptr<float>(j);
            const float *bodies3D_row_ptr = body3D.ptr<float>(j);

            bodies2D_row_ptr[0] = uv_row_ptr[0];
            bodies2D_row_ptr[1] = uv_row_ptr[1];
            bodies2D_row_ptr[2] = bodies3D_row_ptr[3];
        }

        // Filter invalid projections
        for (size_t j = 0; j < num_joints; ++j)
        {
            float *bodies2D_row_ptr = bodies2D.ptr<float>(j);
            float x = bodies2D_row_ptr[0];
            float y = bodies2D_row_ptr[1];
            bool in_x = x >= 0.0 && x < icam.cam.width;
            bool in_y = y >= 0.0 && y < icam.cam.height;
            if (!in_x || !in_y)
            {
                bodies2D_row_ptr[0] = 0.0;
                bodies2D_row_ptr[1] = 0.0;
                bodies2D_row_ptr[2] = 0.0;
            }
        }

        // Store results
        bodies2D_list[i] = bodies2D;
        dists_list[i] = dists;
    }

    return std::make_tuple(bodies2D_list, dists_list);
}

// =================================================================================================

float TriangulatorInternal::calc_pose_score(
    const cv::Mat &pose1,
    const cv::Mat &pose2,
    const cv::Mat &dist1,
    const CameraInternal &icam)
{
    const float min_score = 0.1;

    // Create mask for valid points
    size_t num_joints = pose1.rows;
    cv::Mat mask(num_joints, 1, CV_8U);
    for (size_t i = 0; i < num_joints; ++i)
    {
        u_char *mask_ptr = mask.ptr<u_char>(i);
        const float *pose1_ptr = pose1.ptr<float>(i);
        const float *pose2_ptr = pose2.ptr<float>(i);

        mask_ptr[0] = (pose1_ptr[2] > min_score) && (pose2_ptr[2] > min_score);
    }

    // Drop if not enough valid points
    int valid_count = cv::countNonZero(mask);
    if (valid_count < 3)
    {
        return 0.0;
    }

    // Calculate scores
    float iscale = (icam.cam.width + icam.cam.height) / 2;
    cv::Mat scores = score_projection(pose1, pose2, dist1, mask, iscale);

    // Drop lowest scores
    size_t drop_k = static_cast<size_t>(pose1.rows * 0.2);
    const size_t min_k = 3;
    std::vector<float> scores_vec;
    scores_vec.reserve(valid_count);
    for (int i = 0; i < scores.rows; ++i)
    {
        const float *scores_ptr = scores.ptr<float>(i);
        const u_char *mask_ptr = mask.ptr<u_char>(i);
        if (mask_ptr[0] > 0)
        {
            scores_vec.push_back(scores_ptr[0]);
        }
    }
    size_t scores_size = scores_vec.size();
    if (scores_size >= min_k)
    {
        drop_k = std::min(drop_k, scores_size - min_k);
        std::partial_sort(scores_vec.begin(), scores_vec.begin() + drop_k, scores_vec.end());
        scores_vec.erase(scores_vec.begin(), scores_vec.begin() + drop_k);
    }

    // Calculate final score
    float score = 0.0;
    size_t n_items = scores_vec.size();
    if (n_items > 0)
    {
        float sum_scores = std::accumulate(scores_vec.begin(), scores_vec.end(), 0.0);
        score = sum_scores / static_cast<float>(n_items);
    }

    return score;
}

// =================================================================================================

cv::Mat TriangulatorInternal::score_projection(
    const cv::Mat &pose,
    const cv::Mat &repro,
    const cv::Mat &dists,
    const cv::Mat &mask,
    float iscale)
{
    const float min_score = 0.1;
    const size_t num_joints = pose.rows;

    // Calculate error
    cv::Mat error = cv::Mat::zeros(num_joints, 1, CV_32F);
    for (size_t i = 0; i < num_joints; ++i)
    {
        const u_char *mask_ptr = mask.ptr<u_char>(i);
        const float *pose_ptr = pose.ptr<float>(i);
        const float *repro_ptr = repro.ptr<float>(i);
        float *error_ptr = error.ptr<float>(i);

        if (mask_ptr)
        {
            float dx = pose_ptr[0] - repro_ptr[0];
            float dy = pose_ptr[1] - repro_ptr[1];
            float err = std::sqrt(dx * dx + dy * dy);

            // Set errors of invisible reprojections to a high value
            float score = repro_ptr[2];
            if (score < min_score)
            {
                err = iscale;
            }

            error_ptr[0] = err;
        }
    }

    // Scale error by image size
    const float inv_iscale = 1.0 / iscale;
    const float iscale_quarter = iscale / 4.0;
    for (size_t i = 0; i < num_joints; ++i)
    {
        float *error_ptr = error.ptr<float>(i);

        float err = error_ptr[0];
        err = std::max(0.0f, std::min(err, iscale_quarter)) * inv_iscale;
        error_ptr[0] = err;
    }

    // Scale error by distance to camera
    float mean_dist = 0.0;
    int count = 0;
    for (size_t i = 0; i < num_joints; ++i)
    {
        const u_char *mask_ptr = mask.ptr<u_char>(i);
        const float *dists_ptr = dists.ptr<float>(i);

        if (mask_ptr)
        {
            mean_dist += dists_ptr[0];
            count++;
        }
    }
    if (count > 0)
    {
        mean_dist /= count;
    }
    const float dscale = std::sqrt(mean_dist / 3.5);
    for (size_t i = 0; i < num_joints; ++i)
    {
        float *error_ptr = error.ptr<float>(i);

        float err = error_ptr[0];
        err *= dscale;
        error_ptr[0] = err;
    }

    // Convert error to score
    cv::Mat score = error;
    for (size_t i = 0; i < num_joints; ++i)
    {
        float *score_ptr = score.ptr<float>(i);
        score_ptr[0] = 1.0 / (1.0 + score_ptr[0] * 10.0);
    }

    return score;
}

// =================================================================================================

std::pair<cv::Mat, float> TriangulatorInternal::triangulate_and_score(
    const cv::Mat &pose1,
    const cv::Mat &pose2,
    const CameraInternal &cam1,
    const CameraInternal &cam2,
    const std::array<std::array<float, 3>, 2> &roomparams,
    const std::vector<std::array<size_t, 2>> &core_limbs_idx)
{
    const float min_score = 0.1;
    const size_t num_joints = pose1.rows;

    // Create mask for valid points
    cv::Mat mask(num_joints, 1, CV_8U);
    for (size_t i = 0; i < num_joints; ++i)
    {
        u_char *mask_ptr = mask.ptr<u_char>(i);
        const float *pose1_ptr = pose1.ptr<float>(i);
        const float *pose2_ptr = pose2.ptr<float>(i);

        mask_ptr[0] = (pose1_ptr[2] > min_score) && (pose2_ptr[2] > min_score);
    }

    // If too few joints are visible, return a low score
    int num_visible = cv::countNonZero(mask);
    if (num_visible < 3)
    {
        cv::Mat pose3d = cv::Mat::zeros(num_joints, 4, CV_32F);
        float score = 0.0;
        return std::make_pair(pose3d, score);
    }

    // Extract coordinates of visible joints
    std::vector<int> dims = {2, num_visible};
    cv::Mat points1 = cv::Mat(dims, CV_32F);
    cv::Mat points2 = cv::Mat(dims, CV_32F);
    int idx = 0;
    float *points1_ptr = points1.ptr<float>(0);
    float *points2_ptr = points2.ptr<float>(0);
    for (size_t i = 0; i < num_joints; ++i)
    {
        const u_char *mask_ptr = mask.ptr<u_char>(i);
        const float *pose1_ptr = pose1.ptr<float>(i);
        const float *pose2_ptr = pose2.ptr<float>(i);

        if (mask_ptr[0])
        {
            points1_ptr[idx + 0 * num_visible] = pose1_ptr[0];
            points1_ptr[idx + 1 * num_visible] = pose1_ptr[1];
            points2_ptr[idx + 0 * num_visible] = pose2_ptr[0];
            points2_ptr[idx + 1 * num_visible] = pose2_ptr[1];
            idx++;
        }
    }

    // Triangulate points
    cv::Mat points4d_h, points3d;
    cv::triangulatePoints(cam1.P, cam2.P, points1, points2, points4d_h);
    cv::convertPointsFromHomogeneous(points4d_h.t(), points3d);

    // Create the 3D pose matrix
    cv::Mat pose3d = cv::Mat::zeros(num_joints, 4, CV_32F);
    idx = 0;
    for (size_t i = 0; i < num_joints; ++i)
    {
        const u_char *mask_ptr = mask.ptr<u_char>(i);
        float *pose3d_ptr = pose3d.ptr<float>(i);
        const float *points3d_ptr = points3d.ptr<float>(idx);

        if (mask_ptr[0])
        {
            pose3d_ptr[0] = points3d_ptr[0];
            pose3d_ptr[1] = points3d_ptr[1];
            pose3d_ptr[2] = points3d_ptr[2];
            pose3d_ptr[3] = 1.0;
            idx++;
        }
    }

    // Check if mean of the points is inside the room bounds
    std::array<float, 3> mean = {0.0, 0.0, 0.0};
    for (size_t i = 0; i < num_joints; ++i)
    {
        const u_char *mask_ptr = mask.ptr<u_char>(i);
        const float *pose3d_ptr = pose3d.ptr<float>(i);

        if (mask_ptr[0])
        {
            mean[0] += pose3d_ptr[0];
            mean[1] += pose3d_ptr[1];
            mean[2] += pose3d_ptr[2];
        }
    }
    float inv_num_vis = 1.0 / num_visible;
    for (int j = 0; j < 3; ++j)
    {
        mean[j] *= inv_num_vis;
    }
    const std::array<float, 3> &size = roomparams[0];
    const std::array<float, 3> &center = roomparams[1];
    for (int j = 0; j < 3; ++j)
    {
        if (mean[j] > center[j] + size[j] / 2.0 ||
            mean[j] < center[j] - size[j] / 2.0)
        {
            // Very low score if outside room
            for (size_t i = 0; i < num_joints; ++i)
            {
                float *pose3d_ptr = pose3d.ptr<float>(i);
                pose3d_ptr[3] = 0.001;
            }
            return {pose3d, 0.001};
        }
    }

    // Calculate reprojections
    std::vector<cv::Mat> poses_3d = {pose3d};
    cv::Mat repro1, dists1, repro2, dists2;
    auto [repro1s, dists1s] = project_poses(poses_3d, cam1, true);
    auto [repro2s, dists2s] = project_poses(poses_3d, cam2, true);
    repro1 = repro1s[0];
    dists1 = dists1s[0];
    repro2 = repro2s[0];
    dists2 = dists2s[0];

    // Calculate scores for each view
    float iscale = (cam1.cam.width + cam1.cam.height) / 2.0;
    cv::Mat score1 = score_projection(pose1, repro1, dists1, mask, iscale);
    cv::Mat score2 = score_projection(pose2, repro2, dists2, mask, iscale);

    // Combine scores
    cv::Mat scores = (score1 + score2) * 0.5;

    // Add scores to 3D pose
    for (size_t i = 0; i < num_joints; ++i)
    {
        const u_char *mask_ptr = mask.ptr<u_char>(i);
        float *pose3d_ptr = pose3d.ptr<float>(i);
        const float *scores_ptr = scores.ptr<float>(i);

        if (mask_ptr[0])
        {
            pose3d_ptr[3] = scores_ptr[0];
        }
    }

    // Set scores outside the room to a low value
    const float wdist = 0.1;
    for (size_t i = 0; i < num_joints; ++i)
    {
        const u_char *mask_ptr = mask.ptr<u_char>(i);
        float *pose3d_ptr = pose3d.ptr<float>(i);

        if (mask_ptr[0])
        {
            for (int j = 0; j < 3; ++j)
            {
                if (pose3d_ptr[j] > center[j] + size[j] / 2.0 + wdist ||
                    pose3d_ptr[j] < center[j] - size[j] / 2.0 - wdist)
                {
                    pose3d_ptr[3] = 0.001;
                    break;
                }
            }
        }
    }

    // Filter clearly wrong limbs
    if (!core_limbs_idx.empty())
    {
        const float max_length_sq = 0.9 * 0.9;
        for (size_t i = 0; i < core_limbs_idx.size(); ++i)
        {
            size_t limb1 = core_limbs_idx[i][0];
            size_t limb2 = core_limbs_idx[i][1];
            float *pose3d_ptr1 = pose3d.ptr<float>(limb1);
            float *pose3d_ptr2 = pose3d.ptr<float>(limb2);

            if (pose3d_ptr1[3] > min_score && pose3d_ptr2[3] > min_score)
            {
                float dx = pose3d_ptr1[0] - pose3d_ptr2[0];
                float dy = pose3d_ptr1[1] - pose3d_ptr2[1];
                float dz = pose3d_ptr1[2] - pose3d_ptr2[2];
                float length_sq = dx * dx + dy * dy + dz * dz;

                if (length_sq > max_length_sq)
                {
                    pose3d_ptr2[3] = 0.001;
                }
            }
        }
    }

    // Drop lowest scores
    size_t drop_k = static_cast<size_t>(num_joints * 0.2);
    const size_t min_k = 3;
    std::vector<float> scores_vec;
    for (size_t i = 0; i < num_joints; ++i)
    {
        const float *pose3d_ptr = pose3d.ptr<float>(i);
        if (pose3d_ptr[3] > min_score)
        {
            scores_vec.push_back(pose3d_ptr[3]);
        }
    }
    size_t scores_size = scores_vec.size();
    if (scores_size >= min_k)
    {
        drop_k = std::min(drop_k, scores_size - min_k);
        std::partial_sort(scores_vec.begin(), scores_vec.begin() + drop_k, scores_vec.end());
        scores_vec.erase(scores_vec.begin(), scores_vec.begin() + drop_k);
    }

    // Calculate final score
    float score = 0.0;
    size_t n_items = scores_vec.size();
    if (n_items > 0)
    {
        float sum_scores = std::accumulate(scores_vec.begin(), scores_vec.end(), 0.0);
        score = sum_scores / static_cast<float>(n_items);
    }

    return std::make_pair(pose3d, score);
}

// =================================================================================================

std::vector<std::tuple<cv::Point3d, cv::Mat, std::vector<int>>> TriangulatorInternal::calc_grouping(
    const std::vector<std::pair<std::tuple<int, int, int, int>, std::pair<int, int>>> &all_pairs,
    const std::vector<std::pair<cv::Mat, float>> &all_scored_poses,
    float min_score)
{
    float max_center_dist_sq = 0.6 * 0.6;
    float max_joint_avg_dist = 0.3;
    size_t num_pairs = all_pairs.size();

    // Calculate pose centers
    std::vector<cv::Point3d> centers;
    centers.resize(num_pairs);
    for (size_t i = 0; i < num_pairs; ++i)
    {
        const cv::Mat &pose_3d = all_scored_poses[i].first;
        size_t num_joints = pose_3d.rows;

        cv::Point3d center(0, 0, 0);
        size_t num_valid = 0;
        for (size_t j = 0; j < num_joints; ++j)
        {
            const float *pose_3d_ptr = pose_3d.ptr<float>(j);

            float score = pose_3d_ptr[3];
            if (score > min_score)
            {
                center.x += pose_3d_ptr[0];
                center.y += pose_3d_ptr[1];
                center.z += pose_3d_ptr[2];
                num_valid++;
            }
        }
        if (num_valid > 0)
        {
            float inv_num_valid = 1.0 / num_valid;
            center.x *= inv_num_valid;
            center.y *= inv_num_valid;
            center.z *= inv_num_valid;
        }
        centers[i] = center;
    }

    // Calculate Groups
    // defined as a tuple of center, pose, and all-pairs-indices of members
    std::vector<std::tuple<cv::Point3d, cv::Mat, std::vector<int>>> groups;
    for (size_t i = 0; i < num_pairs; ++i)
    {
        const cv::Mat &pose_3d = all_scored_poses[i].first;
        size_t num_joints = pose_3d.rows;

        const cv::Point3d &center = centers[i];
        float best_dist = std::numeric_limits<float>::infinity();
        int best_group = -1;

        for (size_t j = 0; j < groups.size(); ++j)
        {
            auto &group = groups[j];
            cv::Point3d &group_center = std::get<0>(group);

            // Check if the center is close enough
            float dx = group_center.x - center.x;
            float dy = group_center.y - center.y;
            float dz = group_center.z - center.z;
            float center_dist_sq = dx * dx + dy * dy + dz * dz;

            if (center_dist_sq < max_center_dist_sq)
            {
                cv::Mat &group_pose = std::get<1>(group);

                // Calculate average joint distance
                float dist_sum = 0.0;
                size_t count = 0;
                for (size_t row = 0; row < num_joints; ++row)
                {
                    const float *pose_3d_ptr = pose_3d.ptr<float>(row);
                    const float *group_pose_ptr = group_pose.ptr<float>(row);

                    float score1 = pose_3d_ptr[3];
                    float score2 = group_pose_ptr[3];

                    if (score1 > min_score && score2 > min_score)
                    {
                        float dx = pose_3d_ptr[0] - group_pose_ptr[0];
                        float dy = pose_3d_ptr[1] - group_pose_ptr[1];
                        float dz = pose_3d_ptr[2] - group_pose_ptr[2];
                        float dist_sq = dx * dx + dy * dy + dz * dz;
                        dist_sum += std::sqrt(dist_sq);
                        count++;
                    }
                }

                if (count > 0)
                {
                    // Check if the average joint distance is close enough
                    float avg_dist = dist_sum / count;
                    if (avg_dist < max_joint_avg_dist && avg_dist < best_dist)
                    {
                        best_dist = avg_dist;
                        best_group = static_cast<int>(j);
                    }
                }
            }
        }

        if (best_group == -1)
        {
            // Create a new group
            std::vector<int> new_indices{static_cast<int>(i)};
            groups.emplace_back(center, pose_3d.clone(), std::move(new_indices));
        }
        else
        {
            // Update existing group
            auto &group = groups[best_group];
            cv::Point3d &group_center = std::get<0>(group);
            cv::Mat &group_pose = std::get<1>(group);
            std::vector<int> &group_indices = std::get<2>(group);

            float n_elems = static_cast<float>(group_indices.size());
            float inv_n1 = 1.0 / (n_elems + 1.0);

            // Update group center
            group_center.x = (group_center.x * n_elems + center.x) * inv_n1;
            group_center.y = (group_center.y * n_elems + center.y) * inv_n1;
            group_center.z = (group_center.z * n_elems + center.z) * inv_n1;

            // Update group pose
            for (size_t row = 0; row < num_joints; ++row)
            {
                const float *pose_3d_ptr = pose_3d.ptr<float>(row);
                float *group_pose_ptr = group_pose.ptr<float>(row);

                group_pose_ptr[0] = (group_pose_ptr[0] * n_elems + pose_3d_ptr[0]) * inv_n1;
                group_pose_ptr[1] = (group_pose_ptr[1] * n_elems + pose_3d_ptr[1]) * inv_n1;
                group_pose_ptr[2] = (group_pose_ptr[2] * n_elems + pose_3d_ptr[2]) * inv_n1;
                group_pose_ptr[3] = (group_pose_ptr[3] * n_elems + pose_3d_ptr[3]) * inv_n1;
            }
            group_indices.push_back(static_cast<int>(i));
        }
    }

    return groups;
}

// =================================================================================================

cv::Mat TriangulatorInternal::merge_group(const std::vector<cv::Mat> &poses_3d, float min_score)
{
    int num_poses = poses_3d.size();
    int num_joints = poses_3d[0].rows;

    // Merge poses to create initial pose
    // Use only those triangulations with a high score
    cv::Mat sum_poses = cv::Mat::zeros(num_joints, 4, CV_32F);
    std::vector<int> sum_mask(num_joints, 0);
    for (int i = 0; i < num_poses; ++i)
    {
        const cv::Mat &pose = poses_3d[i];

        for (int j = 0; j < num_joints; ++j)
        {
            const float *pose_ptr = pose.ptr<float>(j);
            float *sum_ptr = sum_poses.ptr<float>(j);

            float score = pose_ptr[3];
            if (score > min_score)
            {
                sum_ptr[0] += pose_ptr[0];
                sum_ptr[1] += pose_ptr[1];
                sum_ptr[2] += pose_ptr[2];
                sum_ptr[3] += pose_ptr[3];
                sum_mask[j]++;
            }
        }
    }
    cv::Mat initial_pose_3d = cv::Mat::zeros(num_joints, 4, CV_32F);
    for (int j = 0; j < num_joints; ++j)
    {
        if (sum_mask[j] > 0)
        {
            float *initial_ptr = initial_pose_3d.ptr<float>(j);
            const float *sum_ptr = sum_poses.ptr<float>(j);

            float inv_count = 1.0 / sum_mask[j];
            initial_ptr[0] = sum_ptr[0] * inv_count;
            initial_ptr[1] = sum_ptr[1] * inv_count;
            initial_ptr[2] = sum_ptr[2] * inv_count;
            initial_ptr[3] = sum_ptr[3] * inv_count;
        }
    }

    // Use center as default if the initial pose is empty
    std::vector<u_char> jmask(num_joints, 0);
    cv::Point3d center(0, 0, 0);
    int valid_joints = 0;
    for (int j = 0; j < num_joints; ++j)
    {
        const float *initial_ptr = initial_pose_3d.ptr<float>(j);
        float score = initial_ptr[3];

        if (score > min_score)
        {
            jmask[j] = 1;
            center.x += initial_ptr[0];
            center.y += initial_ptr[1];
            center.z += initial_ptr[2];
            valid_joints++;
        }
    }
    if (valid_joints > 0)
    {
        center *= 1.0 / valid_joints;
    }
    for (int j = 0; j < num_joints; ++j)
    {
        if (!jmask[j])
        {
            float *initial_ptr = initial_pose_3d.ptr<float>(j);
            initial_ptr[0] = center.x;
            initial_ptr[1] = center.y;
            initial_ptr[2] = center.z;
        }
    }

    // Drop joints with low scores and filter outlying joints using distance threshold
    float offset = 0.1;
    float max_dist_sq = 1.2 * 1.2;
    cv::Mat mask = cv::Mat::zeros(num_poses, num_joints, CV_8U);
    cv::Mat distances = cv::Mat::zeros(num_poses, num_joints, CV_32F);
    for (int i = 0; i < num_poses; ++i)
    {
        float *distances_ptr = distances.ptr<float>(i);
        u_char *mask_ptr = mask.ptr<u_char>(i);
        const cv::Mat &pose = poses_3d[i];

        for (int j = 0; j < num_joints; ++j)
        {
            const float *initial_ptr = initial_pose_3d.ptr<float>(j);
            const float *pose_ptr = pose.ptr<float>(j);

            float dx = pose_ptr[0] - initial_ptr[0];
            float dy = pose_ptr[1] - initial_ptr[1];
            float dz = pose_ptr[2] - initial_ptr[2];
            float dist_sq = dx * dx + dy * dy + dz * dz;
            distances_ptr[j] = dist_sq;

            float score = pose_ptr[3];
            if (dist_sq <= max_dist_sq && score > (min_score - offset))
            {
                mask_ptr[j] = 1;
            }
        }
    }

    // Select the best-k proposals for each joint that are closest to the initial pose
    int keep_best = 3;
    cv::Mat best_k_mask = cv::Mat::zeros(num_poses, num_joints, CV_8U);
    for (int j = 0; j < num_joints; ++j)
    {
        std::vector<std::pair<float, int>> valid_indices;
        valid_indices.reserve(num_poses);
        for (int i = 0; i < num_poses; ++i)
        {
            if (mask.at<u_char>(i, j))
            {
                valid_indices.emplace_back(distances.at<float>(i, j), i);
            }
        }
        std::partial_sort(
            valid_indices.begin(),
            valid_indices.begin() + std::min(keep_best, static_cast<int>(valid_indices.size())),
            valid_indices.end());

        for (int k = 0; k < std::min(keep_best, static_cast<int>(valid_indices.size())); ++k)
        {
            best_k_mask.at<u_char>(valid_indices[k].second, j) = 1;
        }
    }

    // Combine masks
    mask &= best_k_mask;

    // Compute the final pose
    sum_poses = cv::Mat::zeros(num_joints, 4, CV_32F);
    sum_mask.assign(num_joints, 0);
    for (int i = 0; i < num_poses; ++i)
    {
        const u_char *mask_row_ptr = mask.ptr<u_char>(i);
        const cv::Mat &pose = poses_3d[i];

        for (int j = 0; j < num_joints; ++j)
        {
            if (mask_row_ptr[j])
            {
                float *sum_ptr = sum_poses.ptr<float>(j);
                const float *pose_ptr = pose.ptr<float>(j);

                sum_ptr[0] += pose_ptr[0];
                sum_ptr[1] += pose_ptr[1];
                sum_ptr[2] += pose_ptr[2];
                sum_ptr[3] += pose_ptr[3];

                sum_mask[j]++;
            }
        }
    }
    cv::Mat final_pose_3d = cv::Mat::zeros(num_joints, 4, CV_32F);
    for (int j = 0; j < num_joints; ++j)
    {
        if (sum_mask[j] > 0)
        {
            float *final_pose_ptr = final_pose_3d.ptr<float>(j);
            const float *sum_ptr = sum_poses.ptr<float>(j);

            float inv_count = 1.0 / sum_mask[j];
            final_pose_ptr[0] = sum_ptr[0] * inv_count;
            final_pose_ptr[1] = sum_ptr[1] * inv_count;
            final_pose_ptr[2] = sum_ptr[2] * inv_count;
            final_pose_ptr[3] = sum_ptr[3] * inv_count;
        }
    }

    return final_pose_3d;
}

// =================================================================================================

void TriangulatorInternal::add_extra_joints(
    std::vector<cv::Mat> &poses, const std::vector<std::string> &joint_names)
{
    // Find indices of "head", "ear_left", and "ear_right"
    auto it_head = std::find(joint_names.begin(), joint_names.end(), "head");
    auto it_ear_left = std::find(joint_names.begin(), joint_names.end(), "ear_left");
    auto it_ear_right = std::find(joint_names.begin(), joint_names.end(), "ear_right");
    int idx_h = std::distance(joint_names.begin(), it_head);
    int idx_el = std::distance(joint_names.begin(), it_ear_left);
    int idx_er = std::distance(joint_names.begin(), it_ear_right);

    // If the confidence score of "head" is zero compute it as the average of the ears
    for (size_t i = 0; i < poses.size(); ++i)
    {
        cv::Mat &pose = poses[i];
        float *pose_ptr_h = pose.ptr<float>(idx_h);

        if (pose_ptr_h[3] == 0.0f)
        {
            float *pose_ptr_el = pose.ptr<float>(idx_el);
            float *pose_ptr_er = pose.ptr<float>(idx_er);

            if (pose_ptr_el[3] > 0.0f && pose_ptr_er[3] > 0.0f)
            {
                pose_ptr_h[0] = (pose_ptr_el[0] + pose_ptr_er[0]) * 0.5f;
                pose_ptr_h[1] = (pose_ptr_el[1] + pose_ptr_er[1]) * 0.5f;
                pose_ptr_h[2] = (pose_ptr_el[2] + pose_ptr_er[2]) * 0.5f;
                pose_ptr_h[3] = std::min(pose_ptr_el[3], pose_ptr_er[3]);
            }
        }
    }
}

// =================================================================================================

void TriangulatorInternal::filter_poses(
    std::vector<cv::Mat> &poses,
    const std::array<std::array<float, 3>, 2> &roomparams,
    const std::vector<size_t> &core_joint_idx,
    const std::vector<std::array<size_t, 2>> &core_limbs_idx)
{
    const float min_score = 0.1;
    std::vector<int> drop_indices;

    for (size_t i = 0; i < poses.size(); ++i)
    {
        cv::Mat &pose = poses[i];
        const size_t num_joints = pose.rows;

        // Collect valid joint indices
        std::vector<size_t> valid_joint_idx;
        for (size_t j = 0; j < num_joints; ++j)
        {
            float *pose_ptr = pose.ptr<float>(j);
            if (pose_ptr[3] > min_score)
            {
                valid_joint_idx.push_back(j);
            }
        }

        // Drop poses with too few joints
        if (valid_joint_idx.size() < 5)
        {
            drop_indices.push_back(i);
            continue;
        }

        // Compute min, max and mean coordinates
        std::array<float, 3> mean = {0.0, 0.0, 0.0};
        std::array<float, 3> mins = {
            std::numeric_limits<float>::max(),
            std::numeric_limits<float>::max(),
            std::numeric_limits<float>::max()};
        std::array<float, 3> maxs = {
            std::numeric_limits<float>::lowest(),
            std::numeric_limits<float>::lowest(),
            std::numeric_limits<float>::lowest()};
        for (size_t j = 0; j < valid_joint_idx.size(); ++j)
        {
            float *joint_ptr = pose.ptr<float>(valid_joint_idx[j]);
            for (size_t k = 0; k < 3; ++k)
            {
                mins[k] = std::min(mins[k], joint_ptr[k]);
                maxs[k] = std::max(maxs[k], joint_ptr[k]);
                mean[k] += joint_ptr[k];
            }
        }
        for (size_t k = 0; k < 3; ++k)
        {
            mean[k] /= static_cast<float>(valid_joint_idx.size());
        }

        // Drop poses that are too large or too small or too flat
        float min_flatness = 0.2f;
        float max_size = 2.3f;
        float min_size = 0.3f;
        std::array<float, 3> diff;
        for (int j = 0; j < 3; ++j)
        {
            diff[j] = maxs[j] - mins[j];
        }
        if (diff[0] > max_size || diff[1] > max_size || diff[2] > max_size)
        {
            drop_indices.push_back(i);
            continue;
        }
        if (diff[0] < min_size && diff[1] < min_size && diff[2] < min_size)
        {
            drop_indices.push_back(i);
            continue;
        }
        if ((diff[0] < min_flatness && diff[1] < min_flatness) ||
            (diff[1] < min_flatness && diff[2] < min_flatness) ||
            (diff[2] < min_flatness && diff[0] < min_flatness))
        {
            drop_indices.push_back(i);
            continue;
        }

        // Drop persons outside room
        const float wdist = 0.1f;
        const auto &room_size = roomparams[0];
        const auto &room_center = roomparams[1];
        const std::array<float, 3> room_half_size = {
            room_size[0] / 2.0f,
            room_size[1] / 2.0f,
            room_size[2] / 2.0f};
        bool outside = false;
        for (int j = 0; j < 3; ++j)
        {
            // Check if the mean position is outside the room boundaries
            if (mean[j] > room_half_size[j] + room_center[j] ||
                mean[j] < -room_half_size[j] + room_center[j])
            {
                outside = true;
                break;
            }
        }
        for (int j = 0; j < 3; ++j)
        {
            // Check if any limb is too far outside the room
            if (maxs[j] > room_half_size[j] + room_center[j] + wdist ||
                mins[j] < -room_half_size[j] + room_center[j] - wdist)
            {
                outside = true;
                break;
            }
        }
        if (outside)
        {
            drop_indices.push_back(i);
            continue;
        }

        // Calculate total limb length and average limb length
        const float max_avg_length = 0.5f;
        const float min_avg_length = 0.1f;
        float total_length = 0.0f;
        int total_limbs = 0;
        for (size_t j = 0; j < core_limbs_idx.size(); ++j)
        {
            size_t start_idx = core_joint_idx[core_limbs_idx[j][0]];
            size_t end_idx = core_joint_idx[core_limbs_idx[j][1]];

            float *joint_start_ptr = pose.ptr<float>(start_idx);
            float *joint_end_ptr = pose.ptr<float>(end_idx);

            if (joint_start_ptr[3] < min_score || joint_end_ptr[3] < min_score)
            {
                continue;
            }

            float dx = joint_end_ptr[0] - joint_start_ptr[0];
            float dy = joint_end_ptr[1] - joint_start_ptr[1];
            float dz = joint_end_ptr[2] - joint_start_ptr[2];
            float length = std::sqrt(dx * dx + dy * dy + dz * dz);

            total_length += length;
            total_limbs++;
        }
        if (total_limbs == 0)
        {
            drop_indices.push_back(i);
            continue;
        }
        float average_length = total_length / static_cast<float>(total_limbs);
        if (average_length < min_avg_length)
        {
            drop_indices.push_back(i);
            continue;
        }
        if (total_limbs >= 5 && average_length > max_avg_length)
        {
            drop_indices.push_back(i);
            continue;
        }
    }

    // Set confidences of invalid poses to a low value
    for (size_t i = 0; i < drop_indices.size(); ++i)
    {
        cv::Mat &pose = poses[drop_indices[i]];
        for (int j = 0; j < pose.rows; ++j)
        {
            float *joint_ptr = pose.ptr<float>(j);
            joint_ptr[3] = 0.001f;
        }
    }
}

// =================================================================================================

void TriangulatorInternal::add_missing_joints(
    std::vector<cv::Mat> &poses, const std::vector<std::string> &joint_names)
{
    // Map joint names to their indices for quick lookup
    std::unordered_map<std::string, size_t> joint_name_to_idx;
    for (size_t idx = 0; idx < joint_names.size(); ++idx)
    {
        joint_name_to_idx[joint_names[idx]] = idx;
    }

    // Adjacent mapping
    std::unordered_map<std::string, std::vector<std::string>> adjacents = {
        {"hip_right", {"hip_middle", "hip_left"}},
        {"hip_left", {"hip_middle", "hip_right"}},
        {"knee_right", {"hip_right", "knee_left"}},
        {"knee_left", {"hip_left", "knee_right"}},
        {"ankle_right", {"knee_right", "ankle_left"}},
        {"ankle_left", {"knee_left", "ankle_right"}},
        {"shoulder_right", {"shoulder_middle", "shoulder_left"}},
        {"shoulder_left", {"shoulder_middle", "shoulder_right"}},
        {"elbow_right", {"shoulder_right", "hip_right"}},
        {"elbow_left", {"shoulder_left", "hip_left"}},
        {"wrist_right", {"elbow_right"}},
        {"wrist_left", {"elbow_left"}},
        {"nose", {"shoulder_middle", "shoulder_right", "shoulder_left"}},
        {"head", {"shoulder_middle", "shoulder_right", "shoulder_left"}},
        {"foot_*_left_*", {"ankle_left"}},
        {"foot_*_right_*", {"ankle_right"}},
        {"face_*", {"nose"}},
        {"hand_*_left_*", {"wrist_left"}},
        {"hand_*_right_*", {"wrist_right"}}};

    for (size_t i = 0; i < poses.size(); ++i)
    {
        cv::Mat &pose = poses[i];
        const size_t num_joints = pose.rows;

        // Collect valid joint indices
        std::vector<size_t> valid_joint_idx;
        for (size_t j = 0; j < num_joints; ++j)
        {
            float *pose_ptr = pose.ptr<float>(j);
            if (pose_ptr[3] > min_score)
            {
                valid_joint_idx.push_back(j);
            }
        }

        if (valid_joint_idx.empty())
        {
            continue;
        }

        // Compute body center as the mean of valid joints
        std::array<float, 3> body_center = {0.0f, 0.0f, 0.0f};
        for (size_t idx : valid_joint_idx)
        {
            float *joint_ptr = pose.ptr<float>(idx);
            body_center[0] += joint_ptr[0];
            body_center[1] += joint_ptr[1];
            body_center[2] += joint_ptr[2];
        }
        for (int j = 0; j < 3; ++j)
        {
            body_center[j] /= static_cast<float>(valid_joint_idx.size());
        }

        // Iterate over each joint
        for (size_t j = 0; j < joint_names.size(); ++j)
        {
            std::string adname = "";
            const std::string &jname = joint_names[j];

            // Determine adjacent joint based on name patterns
            if (jname.substr(0, 5) == "foot_" && jname.find("_left") != std::string::npos)
            {
                adname = "foot_*_left_*";
            }
            else if (jname.substr(0, 5) == "foot_" && jname.find("_right") != std::string::npos)
            {
                adname = "foot_*_right_*";
            }
            else if (jname.substr(0, 5) == "face_")
            {
                adname = "face_*";
            }
            else if (jname.substr(0, 5) == "hand_" && jname.find("_left") != std::string::npos)
            {
                adname = "hand_*_left_*";
            }
            else if (jname.substr(0, 5) == "hand_" && jname.find("_right") != std::string::npos)
            {
                adname = "hand_*_right_*";
            }
            else if (adjacents.find(jname) != adjacents.end())
            {
                adname = jname;
            }

            if (adname == "")
            {
                // No adjacent joints defined for this joint
                continue;
            }

            float *joint_ptr = pose.ptr<float>(j);
            if (joint_ptr[3] == 0.0f)
            {
                if (adjacents.find(adname) != adjacents.end())
                {
                    // Joint is missing; attempt to estimate its position based on adjacent joints
                    std::array<float, 3> adjacent_position = {0.0f, 0.0f, 0.0f};
                    size_t adjacent_count = 0;

                    const std::vector<std::string> &adjacent_joint_names = adjacents[adname];
                    for (const std::string &adj_name : adjacent_joint_names)
                    {
                        auto it = joint_name_to_idx.find(adj_name);
                        if (it != joint_name_to_idx.end())
                        {
                            size_t adj_idx = it->second;
                            float *adj_joint_ptr = pose.ptr<float>(adj_idx);
                            if (adj_joint_ptr[3] > 0.1f)
                            {
                                adjacent_position[0] += adj_joint_ptr[0];
                                adjacent_position[1] += adj_joint_ptr[1];
                                adjacent_position[2] += adj_joint_ptr[2];
                                adjacent_count++;
                            }
                        }
                    }

                    if (adjacent_count > 0)
                    {
                        for (size_t k = 0; k < 3; ++k)
                        {
                            adjacent_position[k] /= static_cast<float>(adjacent_count);
                        }

                        // Update the missing joint's position
                        joint_ptr[0] = adjacent_position[0];
                        joint_ptr[1] = adjacent_position[1];
                        joint_ptr[2] = adjacent_position[2];
                    }
                    else
                    {
                        // No valid adjacent joints, use body center
                        joint_ptr[0] = body_center[0];
                        joint_ptr[1] = body_center[1];
                        joint_ptr[2] = body_center[2];
                    }
                }
                else
                {
                    // Adjacent joints not defined, use body center
                    joint_ptr[0] = body_center[0];
                    joint_ptr[1] = body_center[1];
                    joint_ptr[2] = body_center[2];
                }

                // Set a low confidence score
                joint_ptr[3] = 0.1f;
            }
        }
    }
}