#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"


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;

        }

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

        }

        TruckMannedProcess::~TruckMannedProcess()
        {
            if (tools_) {
                hs_truck_release(&tools_);
                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" ;
            LOG_INFO("hs_truck 版本：{}  模型路径：{}", hs_truck_getversion(), model_path);

            hs_truck_param param;
            char modelNames[100];
            strcpy(modelNames, model_path.c_str());
            param.modelNames = modelNames;
            param.thresld = 0.25;
            param.devId = gpu_id;
          
            m_devId = param.devId;
            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(&tools_, param)))))
                LOG_ERROR("Init TruckMannedProcessSdk 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;
                }

            }
        }

        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);
            return res;
        }
        
        bool TruckMannedProcess::update_mstreams(const std::vector<task_id_t>& taskIds, vector<DeviceMemory*> vec_det_input_images, 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<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);
                            
                        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;

                        // 抠图
                        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;
                        }
                      
                        flattened_imgs.emplace_back(std::move(img));
                        flattened_vpc_imgs.emplace_back(std::move(img_info));
                        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];
            {
                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(tools_, flattened_imgs.data() + offset, batch_size, model_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;
                        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;
                        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;

                    // if (hs_count >= algor_param->hs_count_threshold)
                    if (hs_count - det_result.person_motocycle_cnt >= algor_param->hs_count_threshold) //头肩数量去除包含的行人和非机动车数量
                    {
                        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;
                                // 原图
                                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;
                            }
                            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:
                    {

                    }

                }
            }

            return true;
        }

    }  // namespace truck_manned_process

} // namespace ai_engine_module