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

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

namespace ai_engine_module
{
    namespace traffic_light_process
    {
        static std::set<algorithm_type_t> algor_type_list_ = {
            algorithm_type_t::PERSON_RUNNING_REDLIGHTS,
            algorithm_type_t::NONMOTOR_RUNNING_REDLIGHTS,
        };

        inline bool is_valid_label(const label_t &label) {
            return ((label == label_t::red)); 
        }

        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;
        }
 
        vector<tr_point> Mbuild_area(vector<int> args)
        {
            vector<tr_point> result;
            int points_check = args.size();
            //需x y作为一组，是偶数
            if (points_check % 2 == 1) {
                LOG_ERROR("points need x and y, but only get one");
                exit(0);
            }

            for (int i = 0; i < points_check; i++) {
                if (i % 2 == 0) {
                    tr_point tmp_point;
                    tmp_point.x = args[i];
                    result.push_back(tmp_point);
                }
                else {
                    result[result.size() - 1].y = args[i];
                }
            }
            
            return result;
        }
        
        bool McheckPointPolygon(vector<tr_point> tr_boxes, tr_point pCur) {
            //任意四边形有4个顶点
            int nCount = tr_boxes.size();
            tr_point RectPoints[MAX_LANE_NUM];

            // std::cout << "ponits: ";
            for (int i = 0; i < nCount; i++) {
                RectPoints[i] = tr_boxes[i];
                // std::cout << tr_boxes[i].x << " " << tr_boxes[i].y << " ";
            }
            // std::cout << std::endl;

            int nCross = 0;
            pCur.y += 0.1;
            for (int i = 0; i < nCount; i++) {
                //依次取相邻的两个点
                tr_point pStart = RectPoints[i];
                tr_point pEnd = RectPoints[(i + 1) % nCount];

                //相邻的两个点是平行于x轴的，肯定不相交，忽略
                if (pStart.y == pEnd.y) continue;

                //交点在pStart,pEnd的延长线上，pCur肯定不会与pStart.pEnd相交，忽略
                if (pCur.y < min(pStart.y, pEnd.y) || pCur.y > max(pStart.y, pEnd.y))   continue;

                //求当前点和x轴的平行线与pStart,pEnd直线的交点的x坐标
                double x = (double)(pCur.y - pStart.y) * (double)(pEnd.x - pStart.x) / (double)(pEnd.y - pStart.y) + pStart.x;

                //若x坐标大于当前点的坐标，则有交点
                if (x > pCur.x)    nCross++;
            }

            // 单边交点为偶数，点在多边形之外
            return (nCross % 2 == 1);
        }

#if 0
        void CvtYuvToBgr(sy_img src, sy_img dst) {
            const int uvStart = src.w_ * src.h_;
            for (int h = 0; h < dst.h_; h++) {
                for (int w = 0; w < dst.w_; w++) {
                    int uvIndex = h / 2 * src.w_ + w - w % 2;
                    unsigned char y = *(src.data_ + w + h * src.w_);
                    unsigned char u = *(src.data_ + uvStart + uvIndex);
                    unsigned char v = *(src.data_ + uvStart + uvIndex + 1);
            
                    int r = (int)(y + 1.402 * (v - 128));
                    int g = (int)(y - 0.34414 * (u - 128) - 0.71414 * (v - 128));
                    int b = (int)(y + 1.772 * (u -128));
                    
                    if (r > 255) r = 255;
                    if (g > 255) g = 255;
                    if (b > 255) b = 255;
                    if (r < 0) r = 0;
                    if (g < 0) g = 0;
                    if (b < 0) b = 0;
                    *(dst.data_ + h * dst.w_ * 3 + w * 3)     = b;
                    *(dst.data_ + h * dst.w_ * 3 + w * 3 + 1) = g;
                    *(dst.data_ + h * dst.w_ * 3 + w * 3 + 2) = r;
                }
            }
        }
#endif

        TrafficLightProcess::TrafficLightProcess()
            : task_param_manager_(nullptr)
        {

        }

        TrafficLightProcess::~TrafficLightProcess()
        {
            if (tools_) {
                traffic_light_release(&tools_);
                tools_ = nullptr;
            }
            if (m_algorthim_ctx) {
                aclrtDestroyContext(m_algorthim_ctx);
            }
        }

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

            string model_path = models_dir + "/models/village/trlight_det_310p.om" ;
            LOG_INFO("traffic_light 版本：{}  模型路径：{}", traffic_light_getversion(), model_path);

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


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


        void TrafficLightProcess::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> TrafficLightProcess::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 TrafficLightProcess::update_mstreams(const std::vector<task_id_t>& taskIds, vector<sy_img> src_interest_imgs, 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 || 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)
                    {   
                        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;

                            //DEBUG增加多边形区域选择及判断================================================================
                            tr_point curpos;
						    curpos.x = (left + right) * 0.5;    curpos.y = bottom;
                            vector<int> args = {1552,435,1756,537,1915,537,1915,499,1725,435}; //区域(泰兴黄桥锦润福府路口)
                            // vector<int> args = {0,0,0,1080,1920,1080,1920,0}; //区域
                            vector<tr_point> tr_boxes = Mbuild_area(args);
                            int cur_flag = McheckPointPolygon(tr_boxes, curpos); //cur_flag==true表示在区域内
                            if (!cur_flag) continue;
                            //DEBUG END==================================================================================
                            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.box.cls = box.index;
                        data.taskId = taskId;               data.objId = box.id;

                        // 抠图
                        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();
            int n_input_src_image = src_interest_imgs.size();
            traffic_light_result model_results[n_input_src_image];
            {
                int steps = (n_input_src_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_src_image - offset : MAX_BATCH;
                    traffic_light_process_batch(tools_, src_interest_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[flattened_idx_to_batch_idx[n]];
                    int red_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);
                        if (!is_valid_label(label))
                            continue;
                        if (box.left < 1200 || box.top < 159 || box.right > 1307 || box.bottom > 212) continue; // 限制红绿灯的出现位置(泰兴黄桥锦润福府路口)
                        LOG_TRACE("task id is {} obj_id {} label {} index {} score {}", task_id, objId, label, box.index, box.confidence);
                        red_cnt ++;   
                    }
#if 0
                    if (red_cnt > 0) {
                        sy_img img = src_interest_imgs[flattened_idx_to_batch_idx[n]];
						int dataSize = img.w_ * img.h_ * 3 / 2;
                        uint8_t* buffer = new uint8_t[dataSize];
                        uint8_t* imgBgr = new uint8_t[dataSize];
                        aclError aclRet = aclrtMemcpy(buffer, dataSize, img.data_, dataSize, ACL_MEMCPY_DEVICE_TO_HOST);

						sy_img imgyuv_local, tmpbgr;
                        imgyuv_local.data_ = (unsigned char*)buffer;    imgyuv_local.w_ = img.w_;	imgyuv_local.h_ = img.h_;
						tmpbgr.w_ = img.w_;	    tmpbgr.h_ = img.h_;     tmpbgr.data_ = imgBgr;
						CvtYuvToBgr(imgyuv_local, tmpbgr);

                        cv::Mat cvImg = cv::Mat(tmpbgr.h_, tmpbgr.w_, CV_8UC3, tmpbgr.data_);
                        string jpgSaveName1 = "result/oral/" + task_id + "_" + to_string(objId) + ".jpg";
                        cv::imwrite(jpgSaveName1, cvImg);
                        for (unsigned i = 0; i < model_result.objcount; ++i) {
                            cv::Point lt(model_result.objinfo[i].left, model_result.objinfo[i].top);
                            cv::Point rb(model_result.objinfo[i].right, model_result.objinfo[i].bottom);
                            cv::rectangle(cvImg, lt, rb, cv::Scalar(0, 0, 255), 3);

                            char buffer[50]; int fontface = cv::FONT_HERSHEY_SIMPLEX; double scale = 0.8; int thickness = 2; int baseline = 0;
                            snprintf(buffer, sizeof(buffer), "%d:%.2f", model_result.objinfo[i].index, model_result.objinfo[i].confidence);
                            cv::Size text = cv::getTextSize(buffer, fontface, scale, thickness, &baseline);
                            cv::putText(cvImg, buffer, lt - cv::Point(0, baseline), fontface, scale, cv::Scalar(0, 0, 255), thickness, 8);
                        }
                        string jpgSaveName2 = "result/draw/" + task_id + "_" + to_string(objId) + ".jpg";
                        cv::imwrite(jpgSaveName2, cvImg);

                        delete[] buffer;    buffer = NULL;
						delete[] imgBgr;    imgBgr = NULL;
                    }
#endif

                    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);
                        if (algor_type == algorithm_type_t::PERSON_RUNNING_REDLIGHTS && static_cast<det_class_label_t>(det_result.box.cls) != det_class_label_t::HUMAN)
                            continue;
                        // if (algor_type == algorithm_type_t::NONMOTOR_RUNNING_REDLIGHTS && (static_cast<det_class_label_t>(det_result.box.cls) != det_class_label_t::MOTOCYCLE ||
                        //     static_cast<det_class_label_t>(det_result.box.cls) != det_class_label_t::BICYCLE))
                        //     continue;

                        if (algor_type == algorithm_type_t::NONMOTOR_RUNNING_REDLIGHTS && static_cast<det_class_label_t>(det_result.box.cls) == det_class_label_t::HUMAN)
                            continue;
                        
                        auto& e = id_to_mn_[obj_key];
                        ++e.m_frame;

                        // 数量小于设定阈值不报警
                        if (red_cnt < algor_param->hs_count_threshold)
                            continue;
                        
                        {
                            if (++e.n_frame == algor_param->n)
                            {
                                results_data_t result;
                                {   LOG_TRACE("task id is {} obj_id {} red_cnt {}", task_id, det_result.objId, red_cnt);
                                    result.box = det_result.box;
                                    result.taskId = det_result.taskId;
                                    result.objId = det_result.objId;
                                    result.algor_type = algor_type; 
                                    // 原图
                                    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_; //需复制
                                }
                                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 traffic_light_process

} // namespace ai_engine_module