#pragma once

#include <filesystem>
#include <memory>
#include <string>
#include <vector>

#include <onnxruntime_cxx_api.h>
#include <onnxruntime_c_api.h>
#include <opencv2/opencv.hpp>
#include <tensorrt_provider_options.h>

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

namespace utils_2d_pose
{

    class BaseModel
    {
    public:
        explicit BaseModel(const std::string &model_path, int warmup_iterations);
        virtual ~BaseModel() = default;

        std::vector<std::vector<std::vector<float>>> call_by_image(
            const std::vector<cv::Mat> &images);

    protected:
        static Ort::Env &get_env()
        {
            // Static method to create or retrieve a single global Ort::Env
            static Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "tensorrt_model_env");
            // static Ort::Env env(ORT_LOGGING_LEVEL_VERBOSE, "tensorrt_model_env");
            return env;
        }

        // Initialize ONNX Runtime session with TensorRT provider
        void init_onnx_runtime(const std::string &model_path);

        // Generate random input data, run forward pass
        void warmup(int epoch);

        // Actually run the session with given input tensors
        std::vector<Ort::Value> call_model(const std::vector<Ort::Value> &inputs);

        // Utility to check if a string ends with a given suffix
        bool ends_with(const std::string &str, const std::string &suffix)
        {
            if (str.size() < suffix.size())
                return false;
            return (str.compare(str.size() - suffix.size(), suffix.size(), suffix) == 0);
        }

    protected:
        Ort::SessionOptions session_options;
        std::unique_ptr<Ort::Session> session;
        Ort::AllocatorWithDefaultOptions allocator;

        std::vector<std::string> input_names;
        std::vector<std::vector<int64_t>> input_shapes;
        std::vector<ONNXTensorElementDataType> input_types;
        std::vector<std::string> output_names;
        std::vector<const char *> input_name_ptrs;
        std::vector<const char *> output_name_ptrs;
    };

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

    BaseModel::BaseModel(const std::string &model_path, int warmup_iterations)
    {
        if (!std::filesystem::exists(model_path))
        {
            throw std::runtime_error("File not found: " + model_path);
        }

        if (!ends_with(model_path, ".onnx"))
        {
            throw std::runtime_error("Only .onnx models are supported: " + model_path);
        }

        std::cout << "Loading model: " << model_path << std::endl;
        init_onnx_runtime(model_path);

        if (warmup_iterations > 0)
        {
            std::cout << "Running warmup ...\n";
            warmup((int)warmup_iterations / 2);
            warmup((int)warmup_iterations / 2);
        }
    }

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

    void BaseModel::init_onnx_runtime(const std::string &model_path)
    {
        // 0) Create env
        Ort::Env &env = get_env();

        session_options = Ort::SessionOptions();
        session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_EXTENDED);

        // 1) Create TensorRT provider options
        OrtTensorRTProviderOptionsV2 *tensorrt_options = nullptr;
        Ort::ThrowOnError(Ort::GetApi().CreateTensorRTProviderOptions(&tensorrt_options));

        // 2) Configure runtime options
        tensorrt_options->trt_max_workspace_size = 2147483647;
        tensorrt_options->trt_engine_cache_enable = 1;
        tensorrt_options->trt_engine_cache_path = "/RapidPoseTriangulation/data/trt_cache/";
        if (model_path.find("_fp16") != std::string::npos)
        {
            tensorrt_options->trt_fp16_enable = 1;
        }
        else
        {
            tensorrt_options->trt_fp16_enable = 0;
        }
        if (model_path.find("_Bx") != std::string::npos)
        {
            tensorrt_options->trt_profile_min_shapes = "image_input:1x384x288x3";
            tensorrt_options->trt_profile_max_shapes = "image_input:10x384x288x3";
            tensorrt_options->trt_profile_opt_shapes = "image_input:3x384x288x3";
        }

        // 3) Append to session options
        Ort::ThrowOnError(
            Ort::GetApi().SessionOptionsAppendExecutionProvider_TensorRT_V2(
                static_cast<OrtSessionOptions *>(session_options),
                tensorrt_options));

        // 4) Create the session
        session = std::make_unique<Ort::Session>(env, model_path.c_str(), session_options);

        // 5) Gather input info
        size_t num_inputs = session->GetInputCount();
        for (size_t i = 0; i < num_inputs; i++)
        {
            Ort::AllocatedStringPtr iname = session->GetInputNameAllocated(i, allocator);
            input_names.push_back(iname.get());

            Ort::TypeInfo type_info = session->GetInputTypeInfo(i);
            auto tensor_info = type_info.GetTensorTypeAndShapeInfo();

            input_shapes.push_back(tensor_info.GetShape());
            input_types.push_back(tensor_info.GetElementType());
        }

        // Gather output info to avoid "At least one output should be requested."
        size_t num_outputs = session->GetOutputCount();
        for (size_t i = 0; i < num_outputs; i++)
        {
            Ort::AllocatedStringPtr oname = session->GetOutputNameAllocated(i, allocator);
            output_names.push_back(oname.get());
        }

        // Build array of input name C-strings
        input_name_ptrs.resize(input_names.size());
        for (size_t i = 0; i < input_names.size(); i++)
        {
            input_name_ptrs[i] = input_names[i].c_str();
        }

        // Build array of output name C-strings
        output_name_ptrs.resize(output_names.size());
        for (size_t i = 0; i < output_names.size(); i++)
        {
            output_name_ptrs[i] = output_names[i].c_str();
        }
    }

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

    void BaseModel::warmup(int epoch)
    {
        std::vector<int> use_batch_sizes = {4, 10, 5, 1, 9, 3, 7, 2, 8, 4, 6, 1, 3, 7, 2};

        std::cout << "Warmup: ";
        for (int e = 0; e < epoch; e++)
        {
            // Extract the 4D shape.
            auto shape = input_shapes[0];
            int batch_size = static_cast<int>(shape[0]);
            int height = static_cast<int>(shape[1]);
            int width = static_cast<int>(shape[2]);
            int channels = static_cast<int>(shape[3]);

            // We'll only handle batch=1 in this example
            if (batch_size != 1)
            {
                batch_size = use_batch_sizes[e % use_batch_sizes.size()];
            }

            std::vector<cv::Mat> images;
            for (int i = 0; i < batch_size; i++)
            {
                // Create a multi-channel Mat of shape (Height, Width, Channels)
                cv::Mat img(height, width, CV_8UC(channels));
                if (channels == 1)
                {
                    cv::randu(img, 0, 256);
                }
                else if (channels == 3)
                {
                    cv::randu(img, cv::Scalar(0, 0, 0), cv::Scalar(256, 256, 256));
                }
                else
                {
                    throw std::runtime_error("Unsupported number of channels for warmup.");
                }

                images.push_back(std::move(img));
            }

            // Call the model
            auto outputs = call_by_image(images);
            std::cout << "#";
        }

        std::cout << std::endl;
    }

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

    std::vector<std::vector<std::vector<float>>> BaseModel::call_by_image(
        const std::vector<cv::Mat> &images)
    {
        size_t batch_size = images.size();
        int height = images[0].rows;
        int width = images[0].cols;
        int channels = images[0].channels();

        // Flatten into single vector
        size_t single_img_size = height * width * channels;
        size_t total_elems = batch_size * single_img_size;
        std::vector<uint8_t> batch_data(total_elems);
        for (size_t b = 0; b < batch_size; b++)
        {
            auto &img = images[b];
            if (img.rows != height || img.cols != width || img.channels() != channels)
            {
                throw std::runtime_error("All images must have the same dimensions!");
            }
            std::memcpy(batch_data.data() + b * single_img_size, img.data, single_img_size);
        }

        // Create an onnx tensor
        auto mem_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
        std::vector<int64_t> shape = {
            static_cast<int64_t>(batch_size),
            static_cast<int64_t>(height),
            static_cast<int64_t>(width),
            static_cast<int64_t>(channels),
        };
        Ort::Value input_tensor = Ort::Value::CreateTensor<uint8_t>(
            mem_info,
            batch_data.data(),
            total_elems,
            shape.data(),
            shape.size());

        // Call model
        std::vector<Ort::Value> input_tensors;
        input_tensors.push_back(std::move(input_tensor));
        auto outputs = call_model(input_tensors);

        // Get pointer to ouput tensor
        const float *tensor_data = outputs[0].GetTensorData<float>();
        auto data_info = outputs[0].GetTensorTypeAndShapeInfo();
        auto shape0 = data_info.GetShape();
        size_t B = (size_t)shape0[0];
        size_t N = (size_t)shape0[1];
        size_t C = (size_t)shape0[2];

        // Convert to vector of values
        std::vector<std::vector<std::vector<float>>> data;
        data.reserve(B);
        for (size_t i = 0; i < B; i++)
        {
            std::vector<std::vector<float>> item;
            item.reserve(N);
            for (size_t j = 0; j < N; j++)
            {
                std::vector<float> values;
                values.reserve(C);
                for (size_t k = 0; k < C; k++)
                {
                    values.push_back(tensor_data[i * N * C + j * C + k]);
                }
                item.push_back(values);
            }
            data.push_back(item);
        }

        return data;
    }

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

    std::vector<Ort::Value> BaseModel::call_model(const std::vector<Ort::Value> &inputs)
    {
        if (inputs.size() != input_names.size())
        {
            throw std::runtime_error("Number of input tensors does not match model's input count.");
        }

        // Actually run the model, requesting all known outputs
        return session->Run(
            Ort::RunOptions{nullptr},
            input_name_ptrs.data(),
            inputs.data(),
            inputs.size(),
            output_name_ptrs.data(),
            output_name_ptrs.size());
    }

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

    class LetterBox
    {
    public:
        LetterBox(const std::array<size_t, 2> &target_size, int fill_value = 0)
        {
            this->target_size = target_size;
            this->fill_value = fill_value;
        }

        std::tuple<std::array<int, 4>, double, std::array<size_t, 2>> calc_params(
            const cv::Mat &image) const;
        cv::Mat resize_image(const cv::Mat &image) const;

    private:
        std::array<size_t, 2> target_size;
        int fill_value;
    };

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

    std::tuple<std::array<int, 4>, double, std::array<size_t, 2>> LetterBox::calc_params(
        const cv::Mat &image) const
    {
        // Original dimensions
        int img_h = image.rows;
        int img_w = image.cols;

        // Target dimensions
        int target_h = static_cast<int>(target_size[0]);
        int target_w = static_cast<int>(target_size[1]);

        // Scale factor
        double scale = std::min(
            static_cast<double>(target_w) / static_cast<double>(img_w),
            static_cast<double>(target_h) / static_cast<double>(img_h));

        // Compute new width/height (round to nearest int)
        int new_w = static_cast<int>(std::round(img_w * scale));
        int new_h = static_cast<int>(std::round(img_h * scale));

        // Calculate padding
        int pad_w = target_w - new_w;
        int pad_h = target_h - new_h;

        int pad_left = pad_w / 2;
        int pad_top = pad_h / 2;
        int pad_right = pad_w - pad_left;
        int pad_bottom = pad_h - pad_top;

        std::array<int, 4> paddings = {pad_left, pad_right, pad_top, pad_bottom};
        std::array<size_t, 2> new_size = {(size_t)new_w, (size_t)new_h};

        return {paddings, scale, new_size};
    }

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

    cv::Mat LetterBox::resize_image(const cv::Mat &image) const
    {
        // 1) Compute letterbox params
        auto [paddings, scale, new_sz] = calc_params(image);
        int pad_left = paddings[0];
        int pad_right = paddings[1];
        int pad_top = paddings[2];
        int pad_bottom = paddings[3];

        // 2) Resize using nearest-neighbor interpolation
        cv::Mat resized_img;
        cv::Size cv_new_sz = cv::Size(new_sz[0], new_sz[1]);
        cv::resize(image, resized_img, cv_new_sz, /*fx=*/0, /*fy=*/0, cv::INTER_NEAREST);

        // 3) Pad if needed
        if (pad_left == 0 && pad_right == 0 && pad_top == 0 && pad_bottom == 0)
        {
            // No padding required
            return resized_img;
        }
        else
        {
            cv::Mat final_img;
            cv::Scalar fill_color(fill_value, fill_value, fill_value);

            cv::copyMakeBorder(
                resized_img, final_img,
                pad_top, pad_bottom,
                pad_left, pad_right,
                cv::BORDER_CONSTANT,
                fill_color);
            return final_img;
        }
    }

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

    class BoxCrop
    {
    public:
        BoxCrop(const std::array<size_t, 2> &target_size,
                float padding_scale = 1.0f,
                int fill_value = 0)
        {
            this->target_size = target_size;
            this->padding_scale = padding_scale;
            this->fill_value = fill_value;
        }

        std::tuple<std::array<int, 4>, float, std::array<int, 4>, std::array<size_t, 2>>
        calc_params(const cv::Mat &image, const std::array<float, 6> &bbox) const;

        cv::Mat crop_resize_box(const cv::Mat &image,
                                const std::array<float, 6> &bbox) const;

    private:
        std::array<size_t, 2> target_size;
        float padding_scale;
        int fill_value;
    };

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

    std::tuple<std::array<int, 4>, float, std::array<int, 4>, std::array<size_t, 2>>
    BoxCrop::calc_params(const cv::Mat &image, const std::array<float, 6> &bbox) const
    {
        // Extract some params
        int img_h = image.rows;
        int img_w = image.cols;
        int target_h = target_size[0];
        int target_w = target_size[1];

        // Round the bounding box coordinates
        float start_x = std::floor(bbox[0]);
        float start_y = std::floor(bbox[1]);
        float end_x = std::ceil(bbox[2]);
        float end_y = std::ceil(bbox[3]);

        // Calculate original bounding box center
        float center_x = 0.5f * (start_x + end_x);
        float center_y = 0.5f * (start_y + end_y);

        // Scale the bounding box by the padding_scale
        float bbox_w = end_x - start_x;
        float bbox_h = end_y - start_y;
        float scaled_w = bbox_w * padding_scale;
        float scaled_h = bbox_h * padding_scale;

        // Calculate the aspect ratios
        float bbox_aspect = scaled_w / scaled_h;
        float target_aspect = static_cast<float>(target_w) / static_cast<float>(target_h);

        // Adjust the scaled bounding box to match the target aspect ratio
        float adjusted_w, adjusted_h;
        if (bbox_aspect > target_aspect)
        {
            adjusted_w = scaled_w;
            adjusted_h = scaled_w / target_aspect;
        }
        else
        {
            adjusted_h = scaled_h;
            adjusted_w = scaled_h * target_aspect;
        }

        // Calculate scaled bounding box coordinates
        float new_start_xf = center_x - 0.5f * adjusted_w;
        float new_start_yf = center_y - 0.5f * adjusted_h;
        float new_end_xf = center_x + 0.5f * adjusted_w;
        float new_end_yf = center_y + 0.5f * adjusted_h;

        // Round the box coordinates
        int start_xi = static_cast<int>(std::floor(new_start_xf));
        int start_yi = static_cast<int>(std::floor(new_start_yf));
        int end_xi = static_cast<int>(std::ceil(new_end_xf));
        int end_yi = static_cast<int>(std::ceil(new_end_yf));

        // Define the new box coordinates
        int new_start_x = std::max(0, start_xi);
        int new_start_y = std::max(0, start_yi);
        int new_end_x = std::min(img_w - 1, end_xi);
        int new_end_y = std::min(img_h - 1, end_yi);
        std::array<int, 4> new_box{new_start_x, new_start_y, new_end_x, new_end_y};

        // Calculate resized crop size
        int clipped_w = new_box[2] - new_box[0];
        int clipped_h = new_box[3] - new_box[1];
        float scale_w = static_cast<float>(target_w) / static_cast<float>(clipped_w);
        float scale_h = static_cast<float>(target_h) / static_cast<float>(clipped_h);
        float scale = std::min(scale_w, scale_h);
        int new_w = static_cast<int>(std::round(clipped_w * scale));
        int new_h = static_cast<int>(std::round(clipped_h * scale));

        // Calculate paddings
        int pad_w = target_w - new_w;
        int pad_h = target_h - new_h;
        int pad_left = 0;
        int pad_right = 0;
        int pad_top = 0;
        int pad_bottom = 0;
        if (pad_w > 0)
        {
            if (start_xi < 0)
            {
                pad_left = pad_w;
                pad_right = 0;
            }
            else if (end_xi > img_w)
            {
                pad_left = 0;
                pad_right = pad_w;
            }
            else
            {
                // Can be caused by bbox rounding
                pad_left = pad_w / 2;
                pad_right = pad_w - pad_left;
            }
        }
        if (pad_h > 0)
        {
            if (start_yi < 0)
            {
                pad_top = pad_h;
                pad_bottom = 0;
            }
            else if (end_yi > img_h)
            {
                pad_top = 0;
                pad_bottom = pad_h;
            }
            else
            {
                // Can be caused by bbox rounding
                pad_top = pad_h / 2;
                pad_bottom = pad_h - pad_top;
            }
        }
        std::array<int, 4> paddings{pad_left, pad_right, pad_top, pad_bottom};

        return {paddings, scale, new_box, {(size_t)new_w, (size_t)new_h}};
    }

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

    cv::Mat BoxCrop::crop_resize_box(const cv::Mat &image,
                                     const std::array<float, 6> &bbox) const
    {
        auto [paddings, _, new_box, new_size] = calc_params(image, bbox);

        // Extract the bounding box
        int x1 = new_box[0];
        int y1 = new_box[1];
        int x2 = new_box[2];
        int y2 = new_box[3];
        cv::Rect roi(x1, y1, x2 - x1, y2 - y1);
        cv::Mat cropped_img = image(roi);

        // Resize the image
        cv::Size new_size_cv = cv::Size(new_size[0], new_size[1]);
        cv::Mat resized_img;
        cv::resize(cropped_img, resized_img, new_size_cv, 0, 0, cv::INTER_NEAREST);

        // Optionally pad the image
        int pad_left = paddings[0];
        int pad_right = paddings[1];
        int pad_top = paddings[2];
        int pad_bottom = paddings[3];
        if (pad_left == 0 && pad_right == 0 && pad_top == 0 && pad_bottom == 0)
        {
            // No padding
            return resized_img;
        }
        else
        {
            cv::Mat final_img;
            cv::Scalar fill_color(fill_value, fill_value, fill_value);
            cv::copyMakeBorder(
                resized_img,
                final_img,
                pad_top,
                pad_bottom,
                pad_left,
                pad_right,
                cv::BORDER_CONSTANT,
                fill_color);
            return final_img;
        }
    }

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

    class RTMDet : public BaseModel
    {
    public:
        RTMDet(const std::string &model_path,
               float conf_threshold,
               float min_area_fraction,
               int warmup = 30)
            : BaseModel(model_path, warmup)
        {
            this->target_size = {320, 320};
            this->conf_threshold = conf_threshold;

            float img_area = target_size[0] * target_size[1];
            this->min_area = img_area * min_area_fraction;

            this->letterbox = std::make_unique<LetterBox>(target_size, 114);
        }

        std::vector<std::array<float, 6>> call(const cv::Mat &image);

    private:
        std::array<size_t, 2> target_size;
        float conf_threshold;
        float min_area;
        std::unique_ptr<LetterBox> letterbox;

        cv::Mat preprocess(const cv::Mat &image);
        std::vector<std::array<float, 6>> postprocess(
            const std::vector<std::vector<std::vector<float>>> &result,
            const cv::Mat &image);

        void clip_boxes(std::vector<std::array<float, 6>> &boxes,
                        const cv::Mat &image) const;
    };

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

    std::vector<std::array<float, 6>> RTMDet::call(const cv::Mat &image)
    {
        cv::Mat preprocessed = preprocess(image);
        std::vector<cv::Mat> inputs = {preprocessed};
        auto results = call_by_image(inputs);
        auto outputs = postprocess(results, image);
        return outputs;
    }

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

    cv::Mat RTMDet::preprocess(const cv::Mat &image)
    {
        cv::Mat resized = letterbox->resize_image(image);
        return resized;
    }

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

    std::vector<std::array<float, 6>> RTMDet::postprocess(
        const std::vector<std::vector<std::vector<float>>> &result,
        const cv::Mat &image)
    {
        // Expected result shape: [B, N, 6] => (x1,y1,x2,y2,score,class)

        // Convert to vector of boxes
        std::vector<std::array<float, 6>> boxes;
        for (auto &item : result[0])
        {
            float x1 = item[0];
            float y1 = item[1];
            float x2 = item[2];
            float y2 = item[3];
            float score = item[4];
            float cls = item[5];
            boxes.push_back({x1, y1, x2, y2, score, cls});
        }

        if (boxes.empty())
        {
            return {};
        }

        // Filter by confidence + area
        for (int i = boxes.size() - 1; i >= 0; i--)
        {
            auto &b = boxes[i];
            if (b[4] < conf_threshold)
            {
                boxes.erase(boxes.begin() + i);
                continue;
            }

            float area = (b[2] - b[0]) * (b[3] - b[1]);
            if (area < min_area)
            {
                boxes.erase(boxes.begin() + i);
            }
        }

        if (boxes.empty())
        {
            return {};
        }

        // Shift by letterbox padding and scale back
        clip_boxes(boxes, image);

        return boxes;
    }

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

    void RTMDet::clip_boxes(std::vector<std::array<float, 6>> &boxes,
                            const cv::Mat &image) const
    {
        // Get paddings, scale from letterbox
        auto [paddings, scale, _] = letterbox->calc_params(image);
        int pad_left = paddings[0];
        int pad_right = paddings[1];
        int pad_top = paddings[2];
        int pad_bottom = paddings[3];

        // The effective region in the letterboxed image
        float letter_w = static_cast<float>(target_size[1] - (pad_left + pad_right));
        float letter_h = static_cast<float>(target_size[0] - (pad_top + pad_bottom));

        for (auto &b : boxes)
        {
            float &x1 = b[0];
            float &y1 = b[1];
            float &x2 = b[2];
            float &y2 = b[3];

            // Shift by left/top
            x1 -= pad_left;
            x2 -= pad_left;
            y1 -= pad_top;
            y2 -= pad_top;

            // Clamp to letterbox region
            x1 = std::max(0.f, std::min(x1, letter_w - 1.0f));
            x2 = std::max(0.f, std::min(x2, letter_w - 1.0f));
            y1 = std::max(0.f, std::min(y1, letter_h - 1.0f));
            y2 = std::max(0.f, std::min(y2, letter_h - 1.0f));

            // Scale back to original resolution
            x1 /= scale;
            x2 /= scale;
            y1 /= scale;
            y2 /= scale;
        }
    }

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

    class RTMPose : public BaseModel
    {
    public:
        RTMPose(const std::string &model_path,
                int warmup = 30)
            : BaseModel(model_path, warmup)
        {
            this->target_size = {384, 288};
            this->boxcrop = std::make_unique<BoxCrop>(target_size, 1.25f, 114);
        }

        std::vector<std::vector<std::array<float, 3>>> call(const cv::Mat &image,
                                                            const std::vector<std::array<float, 6>> &bboxes);

    private:
        std::array<size_t, 2> target_size;
        std::unique_ptr<BoxCrop> boxcrop;

        std::vector<cv::Mat> preprocess(
            const cv::Mat &image,
            const std::vector<std::array<float, 6>> &bboxes);
        std::vector<std::vector<std::array<float, 3>>> postprocess(
            const std::vector<std::vector<std::vector<float>>> &result,
            const cv::Mat &image,
            const std::vector<std::array<float, 6>> &bboxes);

        void clip_keypoints(std::vector<std::array<float, 3>> &kp,
                            const cv::Mat &image,
                            const std::array<float, 6> &bbox) const;
    };

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

    std::vector<std::vector<std::array<float, 3>>> RTMPose::call(
        const cv::Mat &image, const std::vector<std::array<float, 6>> &bboxes)
    {
        std::vector<cv::Mat> inputs = preprocess(image, bboxes);
        auto results = call_by_image(inputs);
        auto outputs = postprocess(results, image, bboxes);
        return outputs;
    }

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

    std::vector<cv::Mat> RTMPose::preprocess(
        const cv::Mat &image, const std::vector<std::array<float, 6>> &bboxes)
    {
        std::vector<cv::Mat> crops;
        for (auto &bbox : bboxes)
        {
            cv::Mat cropped = boxcrop->crop_resize_box(image, bbox);
            crops.push_back(std::move(cropped));
        }

        return crops;
    }

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

    std::vector<std::vector<std::array<float, 3>>> RTMPose::postprocess(
        const std::vector<std::vector<std::vector<float>>> &result,
        const cv::Mat &image,
        const std::vector<std::array<float, 6>> &bboxes)
    {
        // Expected result shape: [B, N, 3] => (x,y,score)
        std::vector<std::vector<std::array<float, 3>>> keypoints;
        for (size_t i = 0; i < result.size(); i++)
        {
            // Convert to vector of keypoints
            std::vector<std::array<float, 3>> kpts;
            for (auto &item : result[i])
            {
                float x = item[0];
                float y = item[1];
                float score = item[2];

                kpts.push_back({x, y, score});
            }

            if (!kpts.empty())
            {
                // Shift by boxcrop padding and scale back
                clip_keypoints(kpts, image, bboxes[i]);
                keypoints.push_back(std::move(kpts));
            }
        }

        return keypoints;
    }

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

    void RTMPose::clip_keypoints(std::vector<std::array<float, 3>> &kpts,
                                 const cv::Mat &image,
                                 const std::array<float, 6> &bbox) const
    {
        // Get paddings, scale from boxcrop
        auto [paddings, scale, box, _] = boxcrop->calc_params(image, bbox);
        int pad_left = paddings[0];
        int pad_top = paddings[2];
        int box_left = box[0];
        int box_top = box[1];

        for (auto &kp : kpts)
        {
            float &x = kp[0];
            float &y = kp[1];

            // Shift by left/top
            x -= pad_left;
            y -= pad_top;

            x /= scale;
            y /= scale;

            x += box_left;
            y += box_top;
        }
    }

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

    class TopDown
    {
    public:
        TopDown(const std::string &det_model_path,
                const std::string &pose_model_path,
                float box_conf_threshold,
                float box_min_area,
                bool batch_poses = false,
                int warmup = 30)
        {
            this->batch_poses = batch_poses;

            this->det_model = std::make_unique<RTMDet>(
                det_model_path, box_conf_threshold, box_min_area, warmup);
            this->pose_model = std::make_unique<RTMPose>(pose_model_path, warmup);
        }

        std::vector<std::vector<std::array<float, 3>>> predict(const cv::Mat &image);

    private:
        std::unique_ptr<RTMDet> det_model;
        std::unique_ptr<RTMPose> pose_model;
        bool batch_poses;

        void merge_close_poses(
            std::vector<std::vector<std::array<float, 3>>> &poses,
            std::array<size_t, 2> image_size);
    };

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

    std::vector<std::vector<std::array<float, 3>>> TopDown::predict(const cv::Mat &image)
    {
        auto bboxes = det_model->call(image);

        // Drop boxes with non-person class
        for (int i = bboxes.size() - 1; i >= 0; i--)
        {
            if (static_cast<int>(bboxes[i][5]) != 0)
            {
                bboxes.erase(bboxes.begin() + i);
            }
        }

        std::vector<std::vector<std::array<float, 3>>> poses;
        if (this->batch_poses)
        {
            if (!bboxes.empty())
            {
                poses = std::move(pose_model->call(image, bboxes));
            }
        }
        else
        {
            for (size_t i = 0; i < bboxes.size(); i++)
            {
                auto kpts = pose_model->call(image, {bboxes[i]});
                poses.push_back(std::move(kpts[0]));
            }
        }

        // Sometimes the detection model predicts multiple boxes with different shapes for the same
        // person. They then result in strongly overlapping poses, which are merged here.
        merge_close_poses(poses, {(size_t)image.cols, (size_t)image.rows});

        // Clip keypoints far outside the image
        float mask_offset = (image.cols + image.rows) / 10.0;
        for (size_t i = 0; i < poses.size(); ++i)
        {
            for (size_t j = 0; j < poses[i].size(); ++j)
            {
                auto &kp = poses[i][j];
                if (kp[0] < -mask_offset)
                {
                    kp[0] = -mask_offset;
                    kp[2] = 0.001;
                }
                if (kp[1] < -mask_offset)
                {
                    kp[1] = -mask_offset;
                    kp[2] = 0.001;
                }
                if (kp[0] >= image.cols + mask_offset)
                {
                    kp[0] = image.cols + mask_offset;
                    kp[2] = 0.001;
                }
                if (kp[1] >= image.rows + mask_offset)
                {
                    kp[1] = image.rows + mask_offset;
                    kp[2] = 0.001;
                }
            }
        }

        return poses;
    }

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

    void TopDown::merge_close_poses(
        std::vector<std::vector<std::array<float, 3>>> &poses,
        std::array<size_t, 2> image_size)
    {
        // Joint ids in COCO order
        const size_t num_overlaps = 5;
        const std::map<std::string, size_t> joint_indices = {
            {"nose", 0},
            {"left_shoulder", 5},
            {"right_shoulder", 6},
            {"left_hip", 11},
            {"right_hip", 12},
            {"left_elbow", 7},
            {"right_elbow", 8},
        };

        if (poses.size() < 2)
        {
            return;
        }

        // Calculate pixel threshold based on image size
        size_t min_dim = std::min(image_size[0], image_size[1]);
        float pixel_threshold = 0.025f * min_dim;

        // Merge poses if enough joints are close
        std::vector<std::vector<std::array<float, 3>>> merged_poses;
        merged_poses.reserve(poses.size());
        for (auto &opose : poses)
        {
            bool merged = false;
            for (auto &mpose : merged_poses)
            {
                size_t close_count = 0;
                for (auto &kv : joint_indices)
                {
                    size_t joint_id = kv.second;
                    float x1 = opose[joint_id][0];
                    float y1 = opose[joint_id][1];
                    float c1 = opose[joint_id][2];

                    float x2 = mpose[joint_id][0];
                    float y2 = mpose[joint_id][1];
                    float c2 = mpose[joint_id][2];

                    if (c1 > 0.0f && c2 > 0.0f)
                    {
                        float dx = x1 - x2;
                        float dy = y1 - y2;
                        float dist_sq = dx * dx + dy * dy;
                        if (dist_sq <= pixel_threshold * pixel_threshold)
                        {
                            ++close_count;
                        }
                    }
                }

                if (close_count >= num_overlaps)
                {
                    // Merge new pose with existing one
                    for (size_t j = 0; j < mpose.size(); ++j)
                    {
                        if (opose[j][2] > mpose[j][2])
                        {
                            mpose[j] = std::move(opose[j]);
                        }
                    }
                    merged = true;
                    break;
                }
            }
            if (!merged)
            {
                // Not mergeable, add as new pose
                merged_poses.push_back(std::move(opose));
            }
        }

        // Replace original poses with merged ones
        poses = std::move(merged_poses);
    }

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

    class PosePredictor
    {
    public:
        PosePredictor(bool whole_body,
                      float min_bbox_score,
                      float min_bbox_area,
                      bool batch_poses = false)
        {
            std::string base_path = "/RapidPoseTriangulation/extras/mmdeploy/exports/";

            std::string path_det_n1 = base_path + "rtmdet-nano_1x320x320x3_fp16_extra-steps.onnx";
            std::string path_det_m1 = base_path + "rtmdet-m_1x320x320x3_fp16_extra-steps.onnx";

            std::string path_pose_m1 = base_path + "rtmpose-m_1x384x288x3_fp16_extra-steps.onnx";
            std::string path_pose_mb = base_path + "rtmpose-m_Bx384x288x3_fp16_extra-steps.onnx";
            std::string path_pose_w1 = base_path + "rtmpose-l_wb_1x384x288x3_fp16_extra-steps.onnx";
            std::string path_pose_wb = base_path + "rtmpose-l_wb_Bx384x288x3_fp16_extra-steps.onnx";

            this->num_joints = whole_body ? 133 : 17;
            std::string path_det;
            std::string path_pose;
            if (!whole_body)
            {
                path_det = path_det_n1;
                if (!batch_poses)
                {
                    path_pose = path_pose_m1;
                }
                else
                {
                    path_pose = path_pose_mb;
                }
            }
            else
            {
                path_det = path_det_m1;
                if (!batch_poses)
                {
                    path_pose = path_pose_w1;
                }
                else
                {
                    path_pose = path_pose_wb;
                }
            }

            std::cout << "Loading 2D" + std::string(whole_body ? "-WB" : "") + " models ..."
                      << std::endl;

            this->topdown_model = std::make_unique<TopDown>(
                path_det, path_pose, min_bbox_score, min_bbox_area, batch_poses, 30);

            std::cout << "Loaded models \n"
                      << std::endl;
        }

        std::vector<std::vector<std::vector<std::array<float, 3>>>> predict(
            const std::vector<cv::Mat> &images);

    private:
        std::unique_ptr<TopDown> topdown_model;
        size_t num_joints;
    };

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

    std::vector<std::vector<std::vector<std::array<float, 3>>>> PosePredictor::predict(
        const std::vector<cv::Mat> &images)
    {
        std::vector<std::vector<std::vector<std::array<float, 3>>>> keypoints;
        for (auto &img : images)
        {
            auto kpts = topdown_model->predict(img);

            if (!kpts.empty())
            {
                keypoints.push_back(std::move(kpts));
            }
            else
            {
                // create zero keypoints
                std::vector<std::vector<std::array<float, 3>>> zero_keypoints;
                zero_keypoints.resize(1);
                zero_keypoints[0].resize(num_joints, {0.0f, 0.0f, 0.0f});
                keypoints.push_back(std::move(zero_keypoints));
            }
        }
        return keypoints;
    }
}