#include <algorithm>
#include "./motocycle_hs_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 motocycle_hs_process
    {
        static std::set<algorithm_type_t> algor_type_list_ = {
            algorithm_type_t::NONMOTOR_VEHICLE_NOHELMET,
            algorithm_type_t::NONMOTOR_VEHICLE_OVERMAN,
        };

        inline bool is_valid_label(const label_t &label) {
            return ((label == label_t::helmet_major) ||
                    (label == label_t::nohelmet_major)); 
        }

        std::set<algorithm_type_t> task_id_to_algorithm_type_seq(const task_id_t &task_id,
                                                                 task_param_manager *const task_param) {
            std::set<algorithm_type_t> seq;
            auto &&algor_map = task_param->get_task_other_param(task_id);
            if (algor_map) {
                // LOG_TRACE("task id is {} size algor type {}", task_id, algor_map->size());
                for (auto iter = algor_map->begin(); iter != algor_map->end(); ++iter) {
                    if (algor_type_list_.count(iter->first) > 0)
                        seq.emplace(iter->first);
                }
            }
            return seq;  // N(RVO)
        }

        bool is_valid_box(const int top, const int left, const int right, const int bottom, const float score,
                          const algorithm_type_t &algor_type,
                          const task_param_manager::algo_param_type_t_ *params_ptr = nullptr) {
            if (!params_ptr)
                return false;

            if (!snapshot_legal_inarea(params_ptr->basic_param->algor_valid_rect, left, top, right, bottom))
                return false;

            if (params_ptr->algor_param == nullptr)
                return false;


            const unsigned width = right - left;
            const unsigned height = bottom - top;

            if (width == 0 || height == 0)
                return false;

            //! TODO: use switch to replace.
            using data_t = algor_config_param_manned_incident;
            data_t *algor_params_ptr = (data_t *) (params_ptr->algor_param);

            if ((width < algor_params_ptr->obj_min_width || height < algor_params_ptr->obj_min_height || score < algor_params_ptr->obj_confidence_threshold))
                return false;

            return true;
        }


        MotorHsProcess::MotorHsProcess()
            : task_param_manager_(nullptr)
        {

        }

        MotorHsProcess::~MotorHsProcess()
        {
            if (tools_) {
                hs_motor_release(&tools_);
                tools_ = nullptr;
            }
            if (m_algorthim_ctx) {
                aclrtDestroyContext(m_algorthim_ctx);
            }
        }

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

            // string model_path = models_dir + "/models/hs/hs_motor_310p.om" ;
            string model_path = models_dir + "/models/hs/hs_motor_b8_310p.om" ;
            LOG_INFO("hs_motor 版本：{}  模型路径：{}", hs_motor_getversion(), model_path);

            hs_motor_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;
          
            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_motor_init(&tools_, param)))))
                LOG_ERROR("Init MotorHsProcessSdk failed error code is {}", status);
            else
                if (!task_param_manager_)
                    task_param_manager_ = task_param_manager::getInstance();
            return init_;
        }


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


        void MotorHsProcess::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> MotorHsProcess::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 MotorHsProcess::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;
            }

            struct stream_idx_and_algor_seq_t {
                unsigned stream_idx;
                std::set<algorithm_type_t> algors;
            };

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

            unsigned flattened_idx = 0;
            std::map<int, int> flattened_idx_to_batch_idx;
            //! 记录每个box对应的算法以及流id.
            std::map<unsigned, stream_idx_and_algor_seq_t> flattened_idx_to_algor_seq;

            /* 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;
                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)
                    {   
                        auto& taskId = *taskId_iter;
                        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);

                        //! loop per algor from set.
                        stream_idx_and_algor_seq_t stream_idx_and_algor_seq{n, {}};
                        std::set<algorithm_type_t> algorithm_type_seq = task_id_to_algorithm_type_seq(taskId,
                                                                                                  task_param_manager_);  // N(RVO).

                        for (auto algor_iter = algorithm_type_seq.begin();algor_iter != algorithm_type_seq.end(); ++algor_iter) {
                            const algorithm_type_t algor_type = *algor_iter;
                            auto &&algor_param_wrap = task_param_manager_->get_task_other_param(taskId, algor_type);
                            if (!algor_param_wrap) {
                                LOG_ERROR("{} is nullptr when get algor param from task_param", taskId.c_str());
                                continue;
                            }

                            if (!is_valid_box(top, left, right, bottom, box.confidence, algor_type, algor_param_wrap))
                                continue;

                            stream_idx_and_algor_seq.algors.emplace(algor_type);
                        }

                        if (stream_idx_and_algor_seq.algors.empty())
                            continue;

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

                        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 };

                        // 抠图
                        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_algor_seq[flattened_idx] = std::move(stream_idx_and_algor_seq);
                        flattened_idx_to_batch_idx[flattened_idx++] = n;
                    }
                }
                ++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_motor_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_motor_process_batch(tools_, flattened_imgs.data() + offset, batch_size, model_results + offset);
                    hs_motor_process_batchV2(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;
                    auto& task_id = det_result.taskId;

                    auto &stream_idx_and_algor_seq = flattened_idx_to_algor_seq[n];
                    auto &algors = stream_idx_and_algor_seq.algors;
                    // auto &steram_idx = stream_idx_and_algor_seq.stream_idx;
                    
                    const auto& src_img = vec_det_input_images[flattened_idx_to_batch_idx[n]];
                    auto &model_result = model_results[n];
                    int person_cnt = 0, helmet_cnt = 0; //统计主要车驾乘人员数量及戴盔数量
                    for (unsigned i = 0; i < model_result.objcount; ++i) {
                        auto &box = model_result.objinfo[i];
                        const label_t label = static_cast<label_t>(box.index);
                        // LOG_TRACE("task id is {} obj_id {} label {} score {}", task_id, obj_id, label, box.obj_score);
                        if (!is_valid_label(label))
                            continue;
                        person_cnt ++;
                        if (box.index == 0)
                            helmet_cnt ++;
                    }

                    // if (person_cnt == 0) {
                    //     VPCUtil::vpc_img_release(flattened_vpc_imgs[n]); //flattened_imgs[n].data_
                    //     flattened_imgs[n].data_ = nullptr;
                    //     continue;
                    // }

                    for (auto algor_type_iter = algors.begin();algor_type_iter != algors.end(); ++algor_type_iter) {
                        const algorithm_type_t algor_type = *algor_type_iter; 
                        
                        auto &&algor_param_wrap = task_param_manager_->get_task_other_param(task_id, algor_type);
                        if (!algor_param_wrap) {
                            LOG_ERROR("{} is nullptr when get algor param from task_param", task_id);
                            continue;
                        }
                        auto algor_param = ((algor_param_type)algor_param_wrap->algor_param);
                
                        id_t obj_key = obj_key_t{ objId, task_id, algor_type};
                        if (id_to_result_.find(obj_key) != id_to_result_.end())
                            continue;
                        
                        // LOG_TRACE("task id is {} algor type is {} obj_id {}", task_id, int(algor_type), objId);
                        
                        auto& e = id_to_mn_[obj_key];
                        ++e.m_frame;

                        // 超载：数量小于设定阈值不报警
                        if (algor_type == algorithm_type_t::NONMOTOR_VEHICLE_OVERMAN && person_cnt < algor_param->hs_count_threshold)
                            continue;

                        // 未戴盔：戴盔数量不小于驾乘数量则不报警
                        if (algor_type == algorithm_type_t::NONMOTOR_VEHICLE_NOHELMET && person_cnt <= helmet_cnt)
                            continue;
                        
                        {
                            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;
                                    result.algor_type = algor_type; 
#if 0 /*暂不保存报警时刻的抓拍图，有需要再启用*/
                                    // 原图
                                    vpc_img_info src_img_info = VPCUtil::vpc_devMem2vpcImg(src_img);
                                    result.origin_img_desc = src_img_info.pic_desc;
                                    // 抠图--拷贝后赋值
                                    void *outputDataDev = acldvppGetPicDescData(flattened_vpc_imgs[n].pic_desc);
                                    int nBufferSize = acldvppGetPicDescSize(flattened_vpc_imgs[n].pic_desc);

                                    void *devBuffer = nullptr;
                                    auto ret = acldvppMalloc(&devBuffer, nBufferSize);
                                    if (ret != ACL_SUCCESS) {
                                        LOG_ERROR("acldvppMalloc failed, size = %u, errorCode = %d.", nBufferSize, static_cast<int32_t>(ret));
                                        return false;
                                    }
                                    aclrtMemcpy(devBuffer, nBufferSize, outputDataDev, nBufferSize, ACL_MEMCPY_DEVICE_TO_DEVICE);

                                    acldvppPicDesc *vpcInputDesc_= acldvppCreatePicDesc();
                                    acldvppSetPicDescData(vpcInputDesc_, devBuffer); 
                                    acldvppSetPicDescFormat(vpcInputDesc_, PIXEL_FORMAT_YUV_SEMIPLANAR_420);
                                    acldvppSetPicDescWidth(vpcInputDesc_, acldvppGetPicDescWidth(flattened_vpc_imgs[n].pic_desc));
                                    acldvppSetPicDescHeight(vpcInputDesc_, acldvppGetPicDescHeight(flattened_vpc_imgs[n].pic_desc));
                                    acldvppSetPicDescWidthStride(vpcInputDesc_, acldvppGetPicDescWidthStride(flattened_vpc_imgs[n].pic_desc));
                                    acldvppSetPicDescHeightStride(vpcInputDesc_, acldvppGetPicDescHeightStride(flattened_vpc_imgs[n].pic_desc));
                                    acldvppSetPicDescSize(vpcInputDesc_, nBufferSize);

                                    result.roi_img_desc = vpcInputDesc_; //需复制
#endif
                                }
                                id_to_result_.emplace(obj_key, std::move(result));
                            }
                        }

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

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

            return true;
        }

    }  // namespace motocycle_hs_process

} // namespace ai_engine_module