#include <set>

#include "vp_openpose_detector_node.h"

namespace vp_nodes {
        
    vp_openpose_detector_node::vp_openpose_detector_node(std::string node_name, 
                                                        std::string model_path, 
                                                        std::string model_config_path, 
                                                        std::string labels_path, 
                                                        int input_width, 
                                                        int input_height, 
                                                        int batch_size,
                                                        int class_id_offset,
                                                        float score_threshold,
                                                        vp_objects::vp_pose_type type,
                                                        float scale,
                                                        cv::Scalar mean,
                                                        cv::Scalar std,
                                                        bool swap_rb):
                                                        vp_primary_infer_node(node_name, 
                                                                            model_path, 
                                                                            model_config_path, 
                                                                            labels_path, 
                                                                            input_width, 
                                                                            input_height, 
                                                                            batch_size, 
                                                                            class_id_offset, 
                                                                            scale, mean, 
                                                                            std, 
                                                                            swap_rb),
                                                        score_threshold(score_threshold), type(type) {
        this->initialized();
    }
    
    vp_openpose_detector_node::~vp_openpose_detector_node() {
        deinitialized();
    }
    
    void vp_openpose_detector_node::postprocess(const std::vector<cv::Mat>& raw_outputs, const std::vector<std::shared_ptr<vp_objects::vp_frame_meta>>& frame_meta_with_batch) {
        // make sure heads of output are not zero
        assert(raw_outputs.size() > 0);
        // as it just one head of output
        auto& netOutputBlobs = raw_outputs[0];
        // batch, number of heatmaps, h, w
        assert(netOutputBlobs.dims == 4);

        auto start1 = std::chrono::system_clock::now();

        int netOutputBlob_dims[3] = {netOutputBlobs.size[1], netOutputBlobs.size[2], netOutputBlobs.size[3]};
        // scan batch
        for (int b = 0; b < netOutputBlobs.size[0]; b++) {
            auto& frame_meta = frame_meta_with_batch[b];
            //auto h_scale = frame_meta->frame.rows * 1.0 / netOutputBlobs.size[2];
            //auto w_scale = frame_meta->frame.cols * 1.0 / netOutputBlobs.size[3];

            cv::Mat netOutputBlob = cv::Mat(3, netOutputBlob_dims, CV_32F, const_cast<uchar*>(netOutputBlobs.ptr(b)));
            
            int keyPointId = 0;
            std::vector<std::vector<KeyPoint>> detectedKeypoints;
            std::vector<KeyPoint> keyPointsList;

            auto start2 = std::chrono::system_clock::now();
            std::vector<cv::Mat> netOutputParts;
            splitNetOutputBlobToParts(netOutputBlob, cv::Size(frame_meta->frame.cols, frame_meta->frame.rows), netOutputParts);
            
            auto start3 = std::chrono::system_clock::now();
            for(int i = 0; i < points_map.at(type); ++i) {
                std::vector<KeyPoint> keyPoints;
                getKeyPoints(netOutputParts[i], score_threshold, keyPoints);

                for(int i = 0; i< keyPoints.size(); ++i, ++keyPointId) {
                    keyPoints[i].id = keyPointId;
                }

                detectedKeypoints.push_back(keyPoints);
                keyPointsList.insert(keyPointsList.end(), keyPoints.begin(), keyPoints.end());
            }

            std::vector<std::vector<ValidPair>> validPairs;
            std::set<int> invalidPairs;
            getValidPairs(netOutputParts, detectedKeypoints, validPairs, invalidPairs);

            std::vector<std::vector<int>> personwiseKeypoints;
            getPersonwiseKeypoints(validPairs, invalidPairs, personwiseKeypoints);

            // insert pose targets back into frame meta
            for(int n = 0; n < personwiseKeypoints.size(); ++n) {
                std::vector<vp_objects::vp_pose_keypoint> kps;
                for (int i = 0; i < personwiseKeypoints[n].size(); i++) {
                    auto index = personwiseKeypoints[n][i];
                    if (index != -1) {                
                        auto& p = keyPointsList[index];
                        kps.push_back(vp_objects::vp_pose_keypoint {i, p.point.x, p.point.y, p.probability});
                    }
                    else {
                        // point not detected
                        kps.push_back(vp_objects::vp_pose_keypoint {i, -1, -1, 0});
                    }
                }
                
                auto pose_target = std::make_shared<vp_objects::vp_frame_pose_target>(type, kps);
                frame_meta->pose_targets.push_back(pose_target);
            }
        }
    }
}