/*
 * @Author: yangzilong
 * @Date: 2021-12-02 14:52:24
 * @Last Modified by: yangzilong
 * @Last Modified time: Do not edit
 * @Email: yangzilong@objecteye.com
 * @Description:
 */

#include <algorithm>
#include "sy_common.h"
#include "./fight_fall_cls.hpp"
#include "../reprocessing_module/CropImg.h"

namespace ai_engine_module
{
    namespace fight_fall_cls
    {
        long long FightfallCls::gid_ = 0;
        algorithm_type_t FightfallCls::fall_algor_type_ = algorithm_type_t::PEDESTRIAN_FALL;
        algorithm_type_t FightfallCls::fight_algor_type_ = algorithm_type_t::PEDESTRIAN_FIGHT;


        FightfallCls::FightfallCls()
        : task_param_manager_(nullptr)
        {

        }

        FightfallCls::~FightfallCls()
        {
            fight_det_release(&tools_);
            if (!tools_)
            {
                delete tools_;
                tools_ = nullptr;
            }
        }

        bool FightfallCls::init(int gpu_id, char *trt_serialize_file)
        {
            init_ = false;
            fight_det_param param;
            {
                param.mode = gpu_id >= 0 ? DEVICE_GPU : DEVICE_CPU;
                param.gpuid = gpu_id;
                param.score_thresld = 0.8;
                param.auth_license = "sy_tsl_aiplatform_sdk_2021";  // TODO:
                param.max_batch = MAX_BATCH;
	            param.serialize_file = "./serialize_file/fight_fall";
            }
            int status;
            if (!(init_ = (0 == (status = fight_det_init(&tools_, param)))))
               LOG_ERROR("Init FightFellClsSdk failed error code is {}", status);
            else
                if (!task_param_manager_)
                    task_param_manager_ = task_param_manager::getInstance();
            return init_;
        }

        bool FightfallCls::check_initied()
        {
            if (!init_)
                LOG_ERROR("[FightfallCls] call init function please.");
            return init_;
        }


        float iou(const box_t &box1, const box_t &box2)
        {
            float insect_top = max(box1.top, box2.top);
            float insect_left = max(box1.left, box2.left);
            float insect_right = min(box1.right, box2.right);
            float insect_bottom = min(box1.bottom, box2.bottom);

            if (insect_right <= insect_left || insect_bottom <= insect_top)
                return 0.0f;

            float insect_area = (insect_right - insect_left) * (insect_bottom - insect_top);
            float union_area = (box1.right - box1.left) * (box1.bottom - box1.top) +
                                (box2.right - box2.left) * (box2.bottom - box2.top) -
                                insect_area;

            return insect_area / union_area;
        }


        box_t union_box(const box_t &box1, const box_t &box2)
        {
            box_t box;
            {
                box.top = std::min(box1.top, box2.top);
                box.left = std::min(box1.left, box2.left);
                box.right = std::max(box1.right, box2.right);
                box.bottom = std::max(box1.bottom, box2.bottom);
            }
            return box;
        }


        bool FightfallCls::process_mstreams(const std::set<taskid_t> &taskIds, const sy_img *det_input_images, const std::vector<onelevel_det_result> &det_results,
                                            std::vector<result_data_t> &results)
        {
            if (!check_initied())
                return false;

            if (det_results.empty())
            {
                LOG_ERROR("[FightfallCls] call init function please.");
                LOG_ERROR("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. */
            /* 1. preprocess. */
            auto taskId_iter = taskIds.begin();
            std::vector<sy_img> flattened_imgs(0);
            std::vector<input_data_wrap_t> flattened_interest_data(0);  //
            for (int n = 0; n < n_images; ++n)
            {
                int n_interest_obj = 0;
                const sy_img &src_img = det_input_images[n];
                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::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_);

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

                        auto fall_algor_param_wrap = task_param_manager_->get_task_other_param(taskId, this->fall_algor_type_);
                        auto fall_algor_param = (fall_algor_param_wrap != nullptr) ? ((fall_algor_param_type)fall_algor_param_wrap->algor_param) : nullptr;
                        auto fall_basic_param = (fall_algor_param_wrap != nullptr) ? (fall_algor_param_wrap->basic_param) : nullptr;

                        auto fight_algor_param_wrap = task_param_manager_->get_task_other_param(taskId, this->fight_algor_type_);
                        auto fight_algor_param = (fight_algor_param_wrap != nullptr) ? ((fight_algor_param_type)fight_algor_param_wrap->algor_param) : nullptr;
                        auto fight_basic_param = (fight_algor_param_wrap != nullptr) ? (fight_algor_param_wrap->basic_param) : nullptr;


                        auto minimum_width = std::min((fall_algor_param == nullptr ? DEFAULT_MIN_WIDTH : fall_algor_param->pedestrian_min_width),
                                                      (fight_algor_param == nullptr ? DEFAULT_MIN_WIDTH : fight_algor_param->pedestrian_min_width));
                        auto minimum_height = std::min((fall_algor_param == nullptr ? DEFAULT_MIN_HEIGHT : fall_algor_param->pedestrian_min_height),
                                                       (fight_algor_param == nullptr ? DEFAULT_MIN_HEIGHT : fight_algor_param->pedestrian_min_height));
                        auto minimum_threshold = std::min((fall_algor_param == nullptr ? DEFAULT_PTHRESHOLD : fall_algor_param->threshold),
                                                          (fight_algor_param == nullptr ? DEFAULT_PTHRESHOLD : fight_algor_param->threshold));

                        sy_rect intel_rect;
                        if (fight_basic_param == nullptr)
                        {
                            if (fall_basic_param != nullptr)
                                intel_rect = fall_basic_param->algor_valid_rect;
                        }
                        else
                        {
                            intel_rect = fight_basic_param->algor_valid_rect;
                        }


                        if ((width < minimum_width || height < minimum_height || box.confidence < minimum_threshold) ||
                            !snapshot_legal_inarea(intel_rect, left, top, right, bottom))
                            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;


                        sy_img img;
                        img.w_ = width;
                        img.h_ = height;
                        img.c_ = src_img.c_;

                        cudaError_t cuda_status;
                        const unsigned nbytes = img.c_ * img.h_ * img.w_ * sizeof(unsigned char);
                        if (CUDA_SUCCESS != (cuda_status = cudaMalloc((void**)&img.data_, nbytes)))
                        {
		                    LOG_ERROR("cudaMalloc failed: {} malloc nbytes is {} mb is {} ", cudaGetErrorString(cuda_status), nbytes, nbytes / (1024 * 1024));
                            continue;
                        }

                        if (CUDA_SUCCESS != ( cuda_status = cudacommon::CropImgGpu(src_img.data_, src_img.w_, src_img.h_, img.data_, left, top, width, height)))
                        {
                            LOG_ERROR("Crop image GPU failed error is %s wh is [{}, {}] ltrb is [{} {} {} {}]",
                                      cudaGetErrorString(cuda_status), src_img.w_, src_img.h_, data.box.left, data.box.top, data.box.right, data.box.bottom);
                            CHECK(cudaFree(img.data_));
                            continue;
                        }
                        flattened_imgs.emplace_back(std::move(img));
                        flattened_interest_data.emplace_back(std::move(data));
                        flattened_idx_to_batch_idx[flattened_idx++] = n;
                    }
                }
                ++taskId_iter;
            }


            /* 2. inference. */
            int n_input_image = flattened_imgs.size();
            fight_det_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;
                    fight_det_process_batch(tools_, flattened_imgs.data() + offset, batch_size, model_results + offset);
                }
            }


            /* 3. postprocess. */
            {
                /* a. review to 2d format. */
                std::unordered_map<taskid_t, std::vector<fight_det_result>> taskid_to_sdk_result;
                {
                    for (int n = 0; n < n_input_image; ++n)
                        taskid_to_sdk_result[flattened_interest_data[n].taskId].emplace_back(std::move(model_results[n]));
                }

                /* b. post process */
                int n = 0;
                for (auto iter = taskIds.begin(); iter != taskIds.end(); ++iter)
                {
                    result_data_t result_data;
                    auto &taskId = *iter;

                    auto fall_algor_param_wrap = task_param_manager_->get_task_other_param(taskId, this->fall_algor_type_);
                    auto fall_algor_param = (fall_algor_param_wrap != nullptr) ? ((fall_algor_param_type)fall_algor_param_wrap->algor_param) : nullptr;

                    auto fight_algor_param_wrap = task_param_manager_->get_task_other_param(taskId, this->fight_algor_type_);
                    auto fight_algor_param = (fight_algor_param_wrap != nullptr) ? ((fight_algor_param_type)fight_algor_param_wrap->algor_param) : nullptr;

                    auto &model_result = taskid_to_sdk_result[taskId];
                    const unsigned model_result_size = model_result.size();

                    bool has_fight_pair[model_result_size];
                    memset(has_fight_pair, false, model_result_size * sizeof(bool));

                    for (int i = 0; i < model_result_size; ++i, ++n)
                    {
                        const sy_img& src_img = det_input_images[flattened_idx_to_batch_idx[n]];
                        auto &cropped_img = flattened_imgs[n];  // croped image.
                        auto &preprocessed_data = flattened_interest_data[n];
                        auto &fight_result_iter = model_result[i];
                        for (int j = 0; j < FIGHT_MODEL_NUM; ++j)
                        {
                            auto &res = fight_result_iter.fight_infos[j];

                            if (fight_algor_param != nullptr)
                                if (!has_fight_pair[i])
                                    for (int z = i + 1; z < model_result_size; ++z)
                                        if (!has_fight_pair[z] && (res.fight_score > fight_algor_param->threshold ||
                                            model_result[z].fight_infos[j].fight_score > fight_algor_param->threshold))
                                            if (iou(preprocessed_data.box, flattened_interest_data[n + (z - i)].box) > fight_algor_param->iou_threshold)
                                            {
                                                // std::printf("[Debug] fight Pair ");
                                                has_fight_pair[i] = has_fight_pair[z] = true;

                                                auto &preprocessed_data2 = flattened_interest_data[n + (z - i)];
                                                result_fight_data_t data;
                                                {
                                                    data.taskid = taskId;
                                                    data.objectids.insert(preprocessed_data.objId);
                                                    data.objectids.insert(preprocessed_data2.objId);
                                                    data.box = union_box(preprocessed_data.box, preprocessed_data2.box);
                                                    // data.box.score = (res.fight_score + model_result[z].fight_infos[j].fight_score) * 0.5;
                                                    data.box.score = std::max(res.fight_score, model_result[z].fight_infos[j].fight_score);

                                                    // std::printf("\tltrb is [%d %d %d %d] wh of src img is [%d %d]", data.box.left, data.box.top, data.box.right, data.box.bottom, src_img.w_, src_img.h_);

                                                    sy_img img;
                                                    {
                                                        img.c_ = src_img.c_;
                                                        img.w_ = data.box.width();
                                                        img.h_ = data.box.height();
                                                    }
                                                    cudaError_t cuda_status;
                                                    const unsigned nbytes = img.c_ * img.h_ * img.w_ * sizeof(unsigned char);
                                                    if (CUDA_SUCCESS != (cuda_status = cudaMalloc((void**)&img.data_, nbytes)))
                                                    {
                            		                    LOG_ERROR("cudaMalloc failed: {} malloc nbytes is {} mb is {} ", cudaGetErrorString(cuda_status), nbytes, nbytes / (1024 * 1024));
                                                        continue;
                                                    }

                                                    if (CUDA_SUCCESS != ( cuda_status = cudacommon::CropImgGpu(src_img.data_, src_img.w_, src_img.h_, img.data_, data.box.left, data.box.top, img.w_, img.h_)))
                                                    {
                                                        LOG_ERROR("Crop image GPU failed error is {} wh is [{} {}] ltrb is [{} {} {} {}]",
                                                                  cudaGetErrorString(cuda_status), src_img.w_, src_img.h_, data.box.left, data.box.top, data.box.right, data.box.bottom);
                                                        CHECK(cudaFree(img.data_));
                                                        continue;
                                                    }
                                                    data.roi_img = img;
                                                    data.ori_img = src_img;
                                                    data.id = gid_++;
                                                }
                                                result_data.fight_data.emplace_back(std::move(data));
                                            }

                            if (fall_algor_param != nullptr)
                                if (res.fall_score > fall_algor_param->threshold)
                                {
                                    if ((float)(src_img.h_ - preprocessed_data.box.bottom) > (float)(src_img.h_ * 0.05))
                                    {
                                        result_fall_data_t data;
                                        {
                                            data.box = preprocessed_data.box;
                                            data.objectid = preprocessed_data.objId;
                                            data.taskid = taskId;
                                            data.roi_img = cropped_img;
                                            data.ori_img = src_img;
                                            data.id = gid_++;
                                        }
                                        result_data.fall_data.emplace_back(std::move(data));
                                        goto _continue;
                                    }
                                }


                        }
                        cudaFree(cropped_img.data_);
            _continue:
                        {

                        }
                    }
                    results.emplace_back(std::move(result_data));
                }
            }
            return true;
        }  // func end
    }  // namespace fight_fall_cls
} // namespace ai_engine_module