#include <algorithm>
#include "./truck_manned_process.h"
#include <cmath>
#include "../decoder/interface/DeviceMemory.hpp"
#include "../common/logger.hpp"
#include "../ai_platform/mvpt_process_assist.h"

#include "opencv2/opencv.hpp"
#include "opencv2/imgcodecs/legacy/constants_c.h"


namespace ai_engine_module
{
    namespace truck_manned_process
    {
        algorithm_type_t TruckMannedProcess::algor_type_ = algorithm_type_t::TRUCK_MANNED;

        int legal_person_motocycle_inarea(const std::vector<input_data_wrap_t>& person_motocycle_data, int left, int top, int right, int bottom) {
            int legal_person_motocycle_count = 0;
            for (const auto& det_result : person_motocycle_data) {
                int x1_intersection = std::max(det_result.box.left, left);
                int y1_intersection = std::max(det_result.box.top, top);
                int x2_intersection = std::min(det_result.box.right, right);
                int y2_intersection = std::min(det_result.box.bottom, bottom);
                
                if (x2_intersection <= x1_intersection || y2_intersection <= y1_intersection) continue; // 没有交集
                float area_a = float(det_result.box.right - det_result.box.left) * float(det_result.box.bottom - det_result.box.top);
                float area_b = float(right - left) * float(bottom - top);
                float area_intersection = float(x2_intersection - x1_intersection) * float(y2_intersection - y1_intersection);

                if (area_a > area_b || area_a <= 0) continue; //暂不考虑行人及非机动车目标比货车大的情况
                // float iou = area_intersection / (area_a + area_b - area_intersection);
                float small_overlap_ratio = area_intersection / area_a; //计算行人及非机动车目标相对于自身的重叠程度
                if (small_overlap_ratio > 0.2) legal_person_motocycle_count ++;
            }
            return legal_person_motocycle_count;

        }

        void vpose_transform(float pt_x, float pt_y, float center_x, float center_y, float llength, int invert, int *res) {
            //Transform pixel location to different reference.
            float tmp = float(IMG_RES) / llength;
            float t00 = tmp, t11 = tmp;
            float t02 = IMG_RES * (-center_x / llength + .5);
            float t12 = IMG_RES * (-center_y / llength + .5);
            cv::Mat t = (cv::Mat_<float>(3,3) << t00,0,t02,0,t11,t12,0,0,1);
            cv::Mat invt = (cv::Mat_<float>(3,3) << 0,0,0,0,0,0,0,0,0);
            if (invert)
                cv::invert(t, invt);

            cv::Mat new_pt = (cv::Mat_<float>(3,1) << pt_x-1, pt_y-1, 1.);
            cv::Mat arr_new = invt * new_pt;
            
            for (int i = 0; i < 3-1; i++) {
                res[i] = int(arr_new.at<float>(i,0));	
            }
        }

        // 计算当前位置与stop_check_frames间的最大位移
        float diff_car_pos(vector<sy_point> tracker_list, int stop_check_frames, float &res_max_dist) {
            float total_diff_x = 0, total_diff_y = 0, total_dist = 0;
            int count = 0;
            float max_dist = 0;
            sy_point pos_current = tracker_list[tracker_list.size() - 1];
            for (int i = tracker_list.size() - stop_check_frames; i < tracker_list.size(); i++) {
                count += 1;
                float temp_diff_x = 0, temp_diff_y = 0;
                sy_point pos_last = tracker_list[i - 1];

                temp_diff_x = pos_current.x_ - pos_last.x_;
                temp_diff_y = pos_current.y_ - pos_last.y_;
                float temp_dist = sqrt(temp_diff_x * temp_diff_x + temp_diff_y * temp_diff_y);
                if (temp_dist > max_dist) {
                    max_dist = temp_dist;
                }

                total_diff_x += temp_diff_x;
                total_diff_y += temp_diff_y;
                total_dist += temp_dist;
            }

            res_max_dist = max_dist;
            return total_dist;
        }

        TruckMannedProcess::TruckMannedProcess()
            : task_param_manager_(nullptr)
        {

        }

        TruckMannedProcess::~TruckMannedProcess()
        {
            if (hs_tools_) {
                hs_truck_release(&hs_tools_);
                hs_tools_ = nullptr;
            }
            if (vpose_tools_) {
                vehicle_pose_release(&vpose_tools_);
                vpose_tools_ = nullptr;
            }
            if (m_algorthim_ctx) {
                aclrtDestroyContext(m_algorthim_ctx);
            }
        }

        bool TruckMannedProcess::init(int gpu_id, string models_dir)
        {
            init_ = false;

            // string model_path = models_dir + "/models/hs/hs_truck_310p.om" ;
            string model_path = models_dir + "/models/hs/hs_truck_b8_310p.om" ;
            LOG_INFO("hs_truck 版本：{}  模型路径：{}", hs_truck_getversion(), model_path);
            string model_vpose_path = models_dir + "/models/village/vpose_310p.om" ;
            LOG_INFO("vpose 版本：{}  模型路径：{}", vehicle_pose_get_version(), model_vpose_path);

            hs_truck_param param;
            char modelNames[100];
            strcpy(modelNames, model_path.c_str());
            // param.modelNames = modelNames;
            param.modelNames_b = modelNames;
            param.thresld = 0.25;
            param.devId = gpu_id;
            param.max_batch = 8;

            vehicle_pose_param vparam;
            char vpose_modelNames[100];
            strcpy(vpose_modelNames, model_vpose_path.c_str());
            vparam.modelNames = vpose_modelNames; 
            vparam.ldmk_thres = 0.0;
            vparam.devId = gpu_id;
          
            m_devId = gpu_id;
            ACL_CALL(aclrtSetDevice(m_devId), ACL_SUCCESS, -1);
            ACL_CALL(aclrtCreateContext(&m_algorthim_ctx, m_devId), ACL_SUCCESS, -1);

            int status;
            if (!(init_ = (0 == (status = hs_truck_init(&hs_tools_, param)))))
                LOG_ERROR("Init TruckMannedProcessSdk failed error code is {}", status);
            else if (!(init_ = (0 == (status = vehicle_pose_init(&vpose_tools_, vparam)))))
                LOG_ERROR("Init VehiclePoseProcessSdk failed error code is {}", status);
            else
                if (!task_param_manager_)
                    task_param_manager_ = task_param_manager::getInstance();
            return init_;
        }


        bool TruckMannedProcess::check_initied()
        {
            if (!init_)
                LOG_ERROR("[%s:%d] call init function please.", __FILE__, __LINE__);
            return init_;
        }


        void TruckMannedProcess::force_release_result(const task_id_t& task_id) {
            for (auto iter = id_to_result_.begin(); iter != id_to_result_.end();) {
                const auto& key = iter->first;
                if (key.task_id == task_id) {
                    auto& value = iter->second;
                    if (value.origin_img_desc != nullptr) {
                        VPCUtil::vpc_pic_desc_release(value.origin_img_desc);
                    }

                    if (value.roi_img_desc != nullptr) {
                        VPCUtil::vpc_pic_desc_release(value.roi_img_desc);
                    }
                    iter = id_to_result_.erase(iter);
                }
                else {
                    ++iter;
                }

            }
            for (auto iter = id_to_mn_.begin(); iter != id_to_mn_.end();) {
                const auto& key = iter->first;
                if (key.task_id == task_id) { iter = id_to_mn_.erase(iter);}
                else { ++iter; }
            }
        }

        std::shared_ptr<results_data_t> TruckMannedProcess::get_result_by_objectid(const id_t& id, bool do_erase)
        {
            auto it = id_to_result_.find(id);
            if (it == id_to_result_.end())
                return std::shared_ptr<results_data_t>(nullptr);
            std::shared_ptr<results_data_t> res = std::make_shared<results_data_t>(it->second);
            if (do_erase) {
                id_to_result_.erase(id);
                if (id_to_mn_.count(id)) id_to_mn_.erase(id);
            }
            return res;
        }
        
        bool TruckMannedProcess::update_mstreams(const std::vector<task_id_t>& taskIds, vector<DeviceMemory*> vec_det_input_images, map<OBJ_KEY, OBJ_INDEX> &m_total_obj_info, const std::vector<onelevel_det_result> &det_results) 
        {
            if (!check_initied())
                return false;

            if (det_results.empty())
            {
                LOG_DEBUG("detection result is empty.");
                return false;
            }

            int n_images = det_results.size();  // or n_stream

            unsigned flattened_idx = 0;
            std::map<int, int> flattened_idx_to_batch_idx;

            /* 1. Crop & keep some interest class. */
            auto taskId_iter = taskIds.begin();
            std::vector<sy_img> flattened_imgs(0);
            std::vector<vpc_img_info> flattened_vpc_imgs(0);
            std::vector<sy_img> flattened_vpose_imgs(0);
            std::vector<vpc_img_info> flattened_vpc_vpose_imgs(0);
            std::vector<input_data_wrap_t> flattened_interest_data(0);  //
            VPCUtil* pVpcUtil = VPCUtil::getInstance();
            for (int n = 0; n < n_images; ++n)
            {
                int n_interest_obj = 0;
                auto& src_img = vec_det_input_images[n];
                int src_img_w = src_img->getWidth();
                int src_img_h = src_img->getHeight();

                auto& boxes_of_one_image = det_results[n].obj;

                //汇总当前帧的行人/非机动车目标
                std::vector<input_data_wrap_t> person_motocycle_data(0);
                for (int i = 0; i < det_results[n].obj_count; ++i)
                {
                    auto& box = boxes_of_one_image[i];
                    if (static_cast<det_class_label_t>(box.index) == det_class_label_t::MOTOCYCLE || static_cast<det_class_label_t>(box.index) == det_class_label_t::BICYCLE || 
                    static_cast<det_class_label_t>(box.index) == det_class_label_t::HUMAN || static_cast<det_class_label_t>(box.index) == det_class_label_t::TRICYCLE)
                    {
                        auto& taskId = *taskId_iter;
                        input_data_wrap_t data;
                        int top = std::max(int(box.top), 0);                int left = std::max(int(box.left), 0);
                        int right = std::min(int(box.right), src_img_w);    int bottom = std::min(int(box.bottom), src_img_h);
                        int width = right - left;                           int height = bottom - top;
 
                        if (static_cast<det_class_label_t>(box.index) == det_class_label_t::HUMAN && width >= 0.9*height) //基于载人只能看到半身的假设
                            continue;

                        data.box.top = top;                 data.box.left = left;
                        data.box.right = right;             data.box.bottom = bottom;
                        data.box.score = box.confidence;    data.box.cls = box.index;
                        data.taskId = taskId;               data.objId = box.id;

                        person_motocycle_data.emplace_back(std::move(data));
                    }
                }


                for (int i = 0; i < det_results[n].obj_count; ++i)
                {
                    auto& box = boxes_of_one_image[i];
                    if (static_cast<det_class_label_t>(box.index) == det_class_label_t::TRUCK)
                    {
                        auto& taskId = *taskId_iter;
                        auto algor_param_wrap = task_param_manager_->get_task_other_param(taskId, this->algor_type_);
                        if (!algor_param_wrap)
                        {
                            LOG_ERROR("{} is nullptr when get algor param from task_param", taskId.c_str());
                            continue;
                        }
                        auto algor_param = ((algor_param_type)algor_param_wrap->algor_param);

                        input_data_wrap_t data;
                        int top = std::max(int(box.top - (IMAGE_CROP_EXPAND_RATIO * box.top)), 0);
                        int left = std::max(int(box.left - (IMAGE_CROP_EXPAND_RATIO * box.left)), 0);
                        int right = std::min(int(box.right + (IMAGE_CROP_EXPAND_RATIO * box.right)), src_img_w);
                        int bottom = std::min(int(box.bottom + (IMAGE_CROP_EXPAND_RATIO * box.bottom)), src_img_h);

                        int width = right - left;
                        int height = bottom - top;

                        if ((width < algor_param->obj_min_width || height < algor_param->obj_min_height || box.confidence < algor_param->obj_confidence_threshold) ||
                            !snapshot_legal_inarea(algor_param_wrap->basic_param->algor_valid_rect, left, top, right, bottom))
                            continue;

                        //统计货车检测框内行人及非机动车的数量
                        int person_motocycle_inarea = legal_person_motocycle_inarea(person_motocycle_data, left, top, right, bottom);
                        //vpose=========================================================================================================
                        int x_center = round(float((right + left))/2);
                        int y_center = round(float((bottom + top))/2);
                        float llength = std::max(right -left, bottom - top);    
                        
                        // Crop image according to the supplied bounding box.
                        int ul[2]={0};   int br[2]={0};
                        vpose_transform(1, 1, x_center, y_center, llength, 1, ul); // Upper left point
                        vpose_transform(IMG_RES+1, IMG_RES+1, x_center, y_center, llength, 1, br); // Bottom right point
                   
                        int vpose_left = std::max(-ul[0],0);    int vpose_top = std::max(-ul[1],0);
                        int vpose_right = std::min(br[0], src_img_w) - ul[0];
                        int vpose_bottom = std::min(br[1], src_img_h) - ul[1];
                        int vpose_width = vpose_right - vpose_left;
                        int vpose_height = vpose_bottom - vpose_top;
                        if (vpose_width != vpose_height) continue; // 宽高不相等说明处于边缘位置，不处理 
                        //=============================================================================================================
                        
                        data.box.top = top;                     data.box.left = left;
                        data.box.right = right;                 data.box.bottom = bottom;
                        data.box.score = box.confidence;
                        data.taskId = taskId;                   data.objId = box.id;
                        data.id = obj_key_t{ box.id, taskId, algorithm_type_t::TRUCK_MANNED };
                        data.person_motocycle_cnt = person_motocycle_inarea;
                        data.vpose_box.top = vpose_top;         data.vpose_box.left = vpose_left;
                        data.vpose_box.right = vpose_right;     data.vpose_box.bottom = vpose_bottom;
                        data.x_center = x_center;   data.y_center = y_center;   data.llength = llength;

                        // 抠图
                        video_object_info obj;
                        strcpy(obj.task_id, taskId.c_str());
                        obj.object_id = box.id;
                        obj.left = left;    obj.top = top;  obj.right = right;  obj.bottom = bottom;

                        vpc_img_info img_info = pVpcUtil->crop(src_img, obj);

                        sy_img img;
                        img.w_ = width;     img.h_ = height;    img.c_ = src_img->getChannel();
                        
                        if (img_info.pic_desc != nullptr) {
                            void *outputDataDev = acldvppGetPicDescData(img_info.pic_desc);
                            img.data_ = reinterpret_cast<unsigned char*>(outputDataDev);
                        }
                        else {
                            LOG_ERROR("Crop image NPU failed wh is [{}, {}] ltrb is [{} {} {} {}]",
                                src_img_w, src_img_h, data.box.left, data.box.top, data.box.right, data.box.bottom);
                            continue;
                        }

                        // vpose抠图
                        video_object_info vpose_obj;
                        strcpy(vpose_obj.task_id, taskId.c_str());
                        vpose_obj.object_id = box.id;
                        vpose_obj.left = vpose_left;    vpose_obj.top = vpose_top;  vpose_obj.right = vpose_right;  vpose_obj.bottom = vpose_bottom;

                        vpc_img_info vpose_img_info = pVpcUtil->crop(src_img, vpose_obj);

                        sy_img vpose_img;   
                        vpose_img.w_ = vpose_width;     vpose_img.h_ = vpose_height;    vpose_img.c_ = src_img->getChannel();

                        if (vpose_img_info.pic_desc != nullptr) {
                            void *outputDataDev = acldvppGetPicDescData(vpose_img_info.pic_desc);
                            vpose_img.data_ = reinterpret_cast<unsigned char*>(outputDataDev);
                        }
                        else {
                            LOG_ERROR("Crop image NPU failed wh is [{}, {}] ltrb is [{} {} {} {}]",
                                src_img_w, src_img_h, data.vpose_box.left, data.vpose_box.top, data.vpose_box.right, data.vpose_box.bottom);
                            VPCUtil::vpc_img_release(img_info); 
                            continue;
                        }
                      
                        flattened_imgs.emplace_back(std::move(img));
                        flattened_vpc_imgs.emplace_back(std::move(img_info));
                        flattened_vpose_imgs.emplace_back(std::move(vpose_img)); //vpose
                        flattened_vpc_vpose_imgs.emplace_back(std::move(vpose_img_info)); //vpose
                        flattened_interest_data.emplace_back(std::move(data));
                        flattened_idx_to_batch_idx[flattened_idx++] = n;
                    }
                }
                std::vector<input_data_wrap_t>().swap(person_motocycle_data);
                ++taskId_iter;
            }

            int ret = aclrtSetCurrentContext(m_algorthim_ctx);
            if (ACL_SUCCESS != ret) {
                return false;
            }
            /* 2. collection result. */
            int n_input_image = flattened_imgs.size();
            hs_truck_result model_results[n_input_image];
            vehicle_pose_result vpose_results[n_input_image];
            {
                int steps = (n_input_image + MAX_BATCH - 1) / MAX_BATCH;
                for (int step = 0; step < steps; ++step)
                {
                    int offset = step * MAX_BATCH;
                    int batch_size = (step == steps - 1) ? n_input_image - offset : MAX_BATCH;
                    // hs_truck_process_batch(hs_tools_, flattened_imgs.data() + offset, batch_size, model_results + offset);
                    hs_truck_process_batchV2(hs_tools_, flattened_imgs.data() + offset, batch_size, model_results + offset);
                    vehicle_pose_batch(vpose_tools_, flattened_vpose_imgs.data() + offset, batch_size, vpose_results + offset);
                }
            }

            /* 3. postprocess. */
            {
                for (int n = 0; n < n_input_image; ++n)
                {
                    auto& det_result = flattened_interest_data[n];
                    auto& objId = det_result.objId;
                    if (id_to_result_.find(det_result.id) != id_to_result_.end())
                    {
                        VPCUtil::vpc_img_release(flattened_vpc_imgs[n]); //flattened_imgs[n].data_
                        flattened_imgs[n].data_ = nullptr;
                        VPCUtil::vpc_img_release(flattened_vpc_vpose_imgs[n]); 
                        flattened_vpose_imgs[n].data_ = nullptr;
                        continue;
                    }

                    const auto& src_img = vec_det_input_images[flattened_idx_to_batch_idx[n]];

                    auto algor_param_wrap = task_param_manager_->get_task_other_param(det_result.taskId, this->algor_type_);
                    if (!algor_param_wrap)
                    {
                        LOG_ERROR("{} nullptr when get algor param from task_param", det_result.taskId.c_str());
                        VPCUtil::vpc_img_release(flattened_vpc_imgs[n]); //flattened_imgs[n].data_
                        flattened_imgs[n].data_ = nullptr;
                        VPCUtil::vpc_img_release(flattened_vpc_vpose_imgs[n]); 
                        flattened_vpose_imgs[n].data_ = nullptr;
                        continue;
                    }

                    auto algor_param = ((algor_param_type)algor_param_wrap->algor_param);

                    int hs_count = model_results[n].objcount;

                    obj_key_t obj_key{ det_result.objId, det_result.taskId, algorithm_type_t::TRUCK_MANNED };

                    auto& e = id_to_mn_[obj_key];
                    ++e.m_frame;

                    // 过滤停车的货车
                    {
                        OBJ_KEY trace_obj_key = {det_result.taskId, objId};
                        int stop_check_frames = 30; // 6s
                        float stop_min_move = 25;
                        // 时间太短，忽略
                        if (m_total_obj_info[trace_obj_key].center_points.size() > stop_check_frames) {
                            float res_max_dist = 0;
                            float total_dist = diff_car_pos(m_total_obj_info[trace_obj_key].center_points, stop_check_frames, res_max_dist);    
                            if (res_max_dist < stop_min_move) { // 停车
                                VPCUtil::vpc_img_release(flattened_vpc_imgs[n]); //flattened_imgs[n].data_
                                flattened_imgs[n].data_ = nullptr;
                                VPCUtil::vpc_img_release(flattened_vpc_vpose_imgs[n]); 
                                flattened_vpose_imgs[n].data_ = nullptr;
                                continue;   
                            }    
                        }
                        else {
                            VPCUtil::vpc_img_release(flattened_vpc_imgs[n]); //flattened_imgs[n].data_
                            flattened_imgs[n].data_ = nullptr;
                            VPCUtil::vpc_img_release(flattened_vpc_vpose_imgs[n]); 
                            flattened_vpose_imgs[n].data_ = nullptr;
                            continue;
                        }
                    }
                    
                    //vpose===========================================================
                    auto &vpose_result = vpose_results[n]; // vpose
                    int ldmk_x1 = src_img->getWidth(), ldmk_y1 = src_img->getHeight(), ldmk_x2 = 0, ldmk_y2 = 0;
                    for (int p = 0; p < LDMK_COUNT; p++) {
                        int ul[2]={0};
                        // 映射到大图
                        vpose_transform(vpose_result.vehicle_ldmk[p].x_-1, vpose_result.vehicle_ldmk[p].y_-1, det_result.x_center, det_result.y_center, det_result.llength, 1, ul);

                        ldmk_x1 = std::min(ul[0] + 1, ldmk_x1);
                        ldmk_y1 = std::min(ul[1] + 1, ldmk_y1);
                        ldmk_x2 = std::max(ul[0] + 1, ldmk_x2);
                        ldmk_y2 = std::max(ul[1] + 1, ldmk_y2);
                    }
                    // 统计关键点映射出的四边形中的头肩数量
                    int hs_inpose_count = 0;
                    if (ldmk_x2 > ldmk_x1 && ldmk_y2 > ldmk_y1) {
                        for (int i = 0; i < model_results[n].objcount; i++) {
                            int x1_intersection = std::max(det_result.box.left + model_results[n].objinfo[i].left, ldmk_x1);
                            int y1_intersection = std::max(det_result.box.top + model_results[n].objinfo[i].top, ldmk_y1);
                            int x2_intersection = std::min(det_result.box.left + model_results[n].objinfo[i].right, ldmk_x2);
                            int y2_intersection = std::min(det_result.box.top + model_results[n].objinfo[i].bottom, ldmk_y2);
                            if (x2_intersection <= x1_intersection || y2_intersection <= y1_intersection) continue; // 没有交集
                            hs_inpose_count ++;
                        }

                    }
                    //=================================================================
         
                    // if (hs_count >= algor_param->hs_count_threshold)
                    LOG_INFO("hs_count：{}  hs_inpose_count: {} person_motocycle_cnt：{}  threshold: {}", hs_count, hs_inpose_count, det_result.person_motocycle_cnt, algor_param->hs_count_threshold);
                    // if (hs_count - det_result.person_motocycle_cnt >= algor_param->hs_count_threshold) //头肩数量去除包含的行人和非机动车数量
                    // if (hs_count - det_result.person_motocycle_cnt >= algor_param->hs_count_threshold || hs_inpose_count >= algor_param->hs_count_threshold)
                    if (hs_inpose_count >= algor_param->hs_count_threshold) //1025
                    { 
                        if (++e.n_frame == algor_param->n)
                        {
                            results_data_t result;
                            {
                                result.box = det_result.box;
                                result.taskId = det_result.taskId;
                                result.objId = det_result.objId;
#if 0 /*暂不保存报警时刻的抓拍图，有需要再启用*/
                                // 原图
                                vpc_img_info src_img_info = VPCUtil::vpc_devMem2vpcImg(src_img);
                                result.origin_img_desc = src_img_info.pic_desc;
                                // 抠图
                                result.roi_img_desc = flattened_vpc_imgs[n].pic_desc;
#else
                                VPCUtil::vpc_img_release(flattened_vpc_imgs[n]);
#endif
                            }
                            id_to_result_.emplace(obj_key, std::move(result));
                            goto _continue;
                        }
                    }

                    if (e.m_frame == algor_param->m)
                        e.reset();

                    VPCUtil::vpc_img_release(flattened_vpc_imgs[n]); //flattened_imgs[n].data_
                _continue:
                    {
                        VPCUtil::vpc_img_release(flattened_vpc_vpose_imgs[n]);
                    }

                }
            }

            return true;
        }

    }  // namespace truck_manned_process

} // namespace ai_engine_module