/*
 * AXERA is pleased to support the open source community by making ax-samples available.
 * 
 * Copyright (c) 2022, AXERA Semiconductor (Shanghai) Co., Ltd. All rights reserved.
 * 
 * Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
 * in compliance with the License. You may obtain a copy of the License at
 * 
 * https://opensource.org/licenses/BSD-3-Clause
 * 
 * Unless required by applicable law or agreed to in writing, software distributed
 * under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
 * CONDITIONS OF ANY KIND, either express or implied. See the License for the
 * specific language governing permissions and limitations under the License.
 */

/*
 * Author: ls.wang
 */

#pragma once

#include <cstdint>
#include <vector>
#include <algorithm>
#include <cmath>
#include <string>

#ifndef CHECK_SCORE
#define CHECK_SCORE(class_attributes, class_names, score) \
    ((class_attributes.find(class_names) != class_attributes.end() && (score) < class_attributes.find(class_names)->second) || \
    ((score) < class_attributes.find("all")->second))
#endif

#ifndef UNSIGMOID
#define UNSIGMOID(value)(-1.0f * (float)std::log((1.0f / value) - 1.0f))
#define UNSIGMOID_STR "_unsigmoid"
#endif

namespace detection
{
    typedef struct
    {
        int grid0;
        int grid1;
        int stride;
    } GridAndStride;

    struct Box
    {
        float xyxy[4] = {0, 0, 0, 0};
        float xywh[4] = {0, 0, 0, 0};
        float object_score = 0;
        size_t index = 0;
        float score = 0;
        float area = 0;
    };

    static inline float sigmoid(float x)
    {
        return static_cast<float>(1.f / (1.f + exp(-x)));
    }

    static float softmax(const float* src, float* dst, int length)
    {
        const float alpha = *std::max_element(src, src + length);
        float denominator = 0;
        float dis_sum = 0;
        for (int i = 0; i < length; ++i)
        {
            dst[i] = exp(src[i] - alpha);
            denominator += dst[i];
        }
        for (int i = 0; i < length; ++i)
        {
            dst[i] /= denominator;
            dis_sum += i * dst[i];
        }
        return dis_sum;
    }

    template<typename T>
    static inline float intersection_area(const T& a, const T& b)
    {
        float xx1 = std::max(a.xyxy[0], b.xyxy[0]);
        float yy1 = std::max(a.xyxy[1], b.xyxy[1]);
        float xx2 = std::min(a.xyxy[2], b.xyxy[2]);
        float yy2 = std::min(a.xyxy[3], b.xyxy[3]);
        float w = std::max(0.0f, xx2 - xx1 + 1.0f);
        float h = std::max(0.0f, yy2 - yy1 + 1.0f);

        float inter_area = w * h;
        return inter_area;
    }

    template<typename T>
    static void qsort_descent_inplace(std::vector<T>& faceobjects, int left, int right)
    {
        int i = left;
        int j = right;
        float p = faceobjects[(left + right) / 2].score;

        while (i <= j)
        {
            while (faceobjects[i].score > p)
                i++;

            while (faceobjects[j].score < p)
                j--;

            if (i <= j)
            {
                // swap
                std::swap(faceobjects[i], faceobjects[j]);

                i++;
                j--;
            }
        }
#pragma omp parallel sections
        {
#pragma omp section
            {
                if (left < j) qsort_descent_inplace(faceobjects, left, j);
            }
#pragma omp section
            {
                if (i < right) qsort_descent_inplace(faceobjects, i, right);
            }
        }
    }

    template<typename T>
    static void qsort_descent_inplace(std::vector<T>& faceobjects)
    {
        if (faceobjects.empty())
            return;

        qsort_descent_inplace(faceobjects, 0, faceobjects.size() - 1);
    }

    template<typename T>
    static void nms_sorted_bboxes(const std::vector<T>& faceobjects, std::vector<int>& picked, float nms_threshold)
    {
        picked.clear();

        const int n = faceobjects.size();

        std::vector<float> areas(n);
        for (int i = 0; i < n; i++)
        {
            areas[i] = faceobjects[i].area;
        }

        for (int i = 0; i < n; i++)
        {
            const T& a = faceobjects[i];

            int keep = 1;
            for (int j = 0; j < (int)picked.size(); j++)
            {
                const T& b = faceobjects[picked[j]];

                // intersection over union
                float inter_area = intersection_area(a, b);
                float union_area = areas[i] + areas[picked[j]] - inter_area;
                // float IoU = inter_area / union_area
                if (inter_area / union_area > nms_threshold)
                    keep = 0;
            }

            if (keep)
                picked.push_back(i);
        }
    }

    static void generate_grids_and_stride(const int target_w, const int target_h, std::vector<int>& strides, std::vector<GridAndStride>& grid_strides)
    {
        for (auto stride : strides)
        {
            int num_grid_w = target_w / stride;
            int num_grid_h = target_h / stride;
            for (int g1 = 0; g1 < num_grid_h; g1++)
            {
                for (int g0 = 0; g0 < num_grid_w; g0++)
                {
                    GridAndStride gs;
                    gs.grid0 = g0;
                    gs.grid1 = g1;
                    gs.stride = stride;
                    grid_strides.push_back(gs);
                }
            }
        }
    }

    // static void generate_proposals_yolov7(int stride, const float* feat, float prob_threshold, std::vector<Object>& objects,
    //                                       int letterbox_cols, int letterbox_rows, const float* anchors, int cls_num = 80)
    // {
    //     int feat_w = letterbox_cols / stride;
    //     int feat_h = letterbox_rows / stride;

    //     auto feat_ptr = feat;

    //     for (int h = 0; h <= feat_h - 1; h++)
    //     {
    //         for (int w = 0; w <= feat_w - 1; w++)
    //         {
    //             for (int a_index = 0; a_index < 3; ++a_index)
    //             {
    //                 float box_objectness = feat_ptr[4];
    //                 if (box_objectness < prob_threshold)
    //                 {
    //                     feat_ptr += cls_num + 5;
    //                     continue;
    //                 }

    //                 //process cls score
    //                 int class_index = 0;
    //                 float class_score = -FLT_MAX;
    //                 for (int s = 0; s <= cls_num - 1; s++)
    //                 {
    //                     float score = feat_ptr[s + 5];
    //                     if (score > class_score)
    //                     {
    //                         class_index = s;
    //                         class_score = score;
    //                     }
    //                 }

    //                 float box_prob = box_objectness * class_score;

    //                 if (box_prob > prob_threshold)
    //                 {
    //                     float x_center = (feat_ptr[0] * 2 - 0.5f + (float)w) * (float)stride;
    //                     float y_center = (feat_ptr[1] * 2 - 0.5f + (float)h) * (float)stride;
    //                     float box_w = (feat_ptr[2] * 2) * (feat_ptr[2] * 2) * anchors[a_index * 2];
    //                     float box_h = (feat_ptr[3] * 2) * (feat_ptr[3] * 2) * anchors[a_index * 2 + 1];
    //                     float x0 = x_center - box_w * 0.5f;
    //                     float y0 = y_center - box_h * 0.5f;

    //                     Object obj;
    //                     obj.rect.x = x0;
    //                     obj.rect.y = y0;
    //                     obj.rect.width = box_w;
    //                     obj.rect.height = box_h;
    //                     obj.label = class_index;
    //                     obj.prob = box_prob;

    //                     objects.push_back(obj);
    //                 }

    //                 feat_ptr += cls_num + 5;
    //             }
    //         }
    //     }
    // }
    
    static void generate_proposals_yolov5(const float* feat, std::vector<Box>& objects,
                                          std::array<size_t, 4U> output_size, int stride,
                                          const unsigned int *anchors, const std::vector<std::string> class_names,
                                          std::unordered_map<std::string, float> class_attributes)
    {                           
        int cls_num = class_names.size();
        int anchor_num = 3;
        int feat_h = output_size[1];
        int feat_w = output_size[2];

        auto feature_ptr = feat;

        for (int h = 0; h <= feat_h - 1; h++)
        {
            for (int w = 0; w <= feat_w - 1; w++)
            {
                for (int a = 0; a <= anchor_num - 1; a++)
                {
                    //process cls score
                    int class_index = 0;
                    float class_score = -1;
                    for (int s = 0; s <= cls_num - 1; s++)
                    {
                        float score = feature_ptr[s + 5];
                        if (score > class_score)
                        {
                            class_index = s;
                            class_score = score;
                        }
                    }
                    // process box score
                    float box_score = feature_ptr[4];
                    if (CHECK_SCORE(class_attributes, class_names[class_index] + UNSIGMOID_STR, box_score))
                    {
                        feature_ptr += (cls_num + 5);
                        continue;
                    }

                    float final_score = sigmoid(box_score) * sigmoid(class_score);
                    if (CHECK_SCORE(class_attributes, class_names[class_index], final_score))
                    {
                        feature_ptr += (cls_num + 5);
                        continue;
                    }

                    float dx = sigmoid(feature_ptr[0]);
                    float dy = sigmoid(feature_ptr[1]);
                    float dw = sigmoid(feature_ptr[2]);
                    float dh = sigmoid(feature_ptr[3]);
                    float pred_cx = (dx * 2.0f - 0.5f + w) * stride;
                    float pred_cy = (dy * 2.0f - 0.5f + h) * stride;
                    float anchor_w = (float)anchors[a * 2];
                    float anchor_h = (float)anchors[a * 2 + 1];
                    float pred_w = dw * dw * 4.0f * anchor_w;
                    float pred_h = dh * dh * 4.0f * anchor_h;
                    float x0 = pred_cx - pred_w * 0.5f;
                    float y0 = pred_cy - pred_h * 0.5f;
                    float x1 = pred_cx + pred_w * 0.5f;
                    float y1 = pred_cy + pred_h * 0.5f;

                    Box obj;
                    obj.xyxy[0] = x0;
                    obj.xyxy[1] = y0;
                    obj.xyxy[2] = x1;
                    obj.xyxy[3] = y1;
                    obj.index = class_index;
                    obj.score = final_score;
                    obj.area = (obj.xyxy[2] - obj.xyxy[0] + 1) * (obj.xyxy[3] - obj.xyxy[1] + 1);
                    objects.push_back(obj);

                    feature_ptr += (cls_num + 5);
                }
            }
        }
    }

    inline static float clamp(
            float val,
            float min = 0.f,
            float max = 1536.f)
        {
            return val > min ? (val < max ? val : max) : min;
        }
        
    // static void generate_proposals_yolov8(int stride, const float* dfl_feat, const float* cls_feat, const float* cls_idx, float prob_threshold, std::vector<Object>& objects,
    //                                       int letterbox_cols, int letterbox_rows, int cls_num = 80)
    // {
    //     int feat_w = letterbox_cols / stride;
    //     int feat_h = letterbox_rows / stride;
    //     int reg_max = 16;

    //     auto dfl_ptr = dfl_feat;
    //     auto cls_ptr = cls_feat;
    //     auto cls_idx_ptr = cls_idx;

    //     std::vector<float> dis_after_sm(reg_max, 0.f);
    //     for (int h = 0; h <= feat_h - 1; h++)
    //     {
    //         for (int w = 0; w <= feat_w - 1; w++)
    //         {
    //             //process cls score
    //             int class_index = static_cast<int>(cls_idx_ptr[h * feat_w + w]);
    //             float class_score = cls_ptr[h * feat_w * cls_num + w * cls_num + class_index];

    //             float box_prob = sigmoid(class_score);

    //             if (box_prob > prob_threshold)
    //             {
    //                 float pred_ltrb[4];
    //                 for (int k = 0; k < 4; k++)
    //                 {
    //                     float dis = softmax(dfl_ptr + k * reg_max, dis_after_sm.data(), reg_max);
    //                     pred_ltrb[k] = dis * stride;
    //                 }

    //                 float pb_cx = (w + 0.5f) * stride;
    //                 float pb_cy = (h + 0.5f) * stride;

    //                 float x0 = pb_cx - pred_ltrb[0];
    //                 float y0 = pb_cy - pred_ltrb[1];
    //                 float x1 = pb_cx + pred_ltrb[2];
    //                 float y1 = pb_cy + pred_ltrb[3];

    //                 x0 = std::max(std::min(x0, (float)(letterbox_cols - 1)), 0.f);
    //                 y0 = std::max(std::min(y0, (float)(letterbox_rows - 1)), 0.f);
    //                 x1 = std::max(std::min(x1, (float)(letterbox_cols - 1)), 0.f);
    //                 y1 = std::max(std::min(y1, (float)(letterbox_rows - 1)), 0.f);

    //                 Object obj;
    //                 obj.rect.x = x0;
    //                 obj.rect.y = y0;
    //                 obj.rect.width = x1 - x0;
    //                 obj.rect.height = y1 - y0;
    //                 obj.label = class_index;
    //                 obj.prob = box_prob;

    //                 objects.push_back(obj);
    //             }
    //             dfl_ptr += (4 * reg_max);
    //         }
    //     }
    // }

    // static void get_out_bbox(std::vector<Box>& objects, int letterbox_rows, int letterbox_cols, int src_rows, int src_cols)
    // {
    //     /* yolov5 draw the result */
    //     float scale_letterbox;
    //     int resize_rows;
    //     int resize_cols;
    //     if ((letterbox_rows * 1.0 / src_rows) < (letterbox_cols * 1.0 / src_cols))
    //     {
    //         scale_letterbox = letterbox_rows * 1.0 / src_rows;
    //     }
    //     else
    //     {
    //         scale_letterbox = letterbox_cols * 1.0 / src_cols;
    //     }
    //     resize_cols = int(scale_letterbox * src_cols);
    //     resize_rows = int(scale_letterbox * src_rows);

    //     int tmp_h = (letterbox_rows - resize_rows) / 2;
    //     int tmp_w = (letterbox_cols - resize_cols) / 2;

    //     float ratio_x = (float)src_rows / resize_rows;
    //     float ratio_y = (float)src_cols / resize_cols;

    //     int count = objects.size();

    //     objects.resize(count);
    //     for (int i = 0; i < count; i++)
    //     {
    //         float x0 = (objects[i].xyxy[0]);
    //         float y0 = (objects[i].xyxy[1]);
    //         float x1 = (objects[i].xyxy[2]);
    //         float y1 = (objects[i].xyxy[3]);

    //         x0 = (x0 - tmp_w) * ratio_x;
    //         y0 = (y0 - tmp_h) * ratio_y;
    //         x1 = (x1 - tmp_w) * ratio_x;
    //         y1 = (y1 - tmp_h) * ratio_y;

    //         // for (int l = 0; l < 5; l++)
    //         // {
    //         //     auto lx = objects[i].landmark[l].x;
    //         //     auto ly = objects[i].landmark[l].y;
    //         //     objects[i].landmark[l] = cv::Point2f((lx - tmp_w) * ratio_x, (ly - tmp_h) * ratio_y);
    //         // }

    //         x0 = std::max(std::min(x0, (float)(src_cols - 1)), 0.f);
    //         y0 = std::max(std::min(y0, (float)(src_rows - 1)), 0.f);
    //         x1 = std::max(std::min(x1, (float)(src_cols - 1)), 0.f);
    //         y1 = std::max(std::min(y1, (float)(src_rows - 1)), 0.f);

    //         objects[i].xyxy[0] = x0;
    //         objects[i].xyxy[1] = y0;
    //         objects[i].xyxy[2] = x1;
    //         objects[i].xyxy[3] = y1;
    //     }
    // }

    static void get_out_bbox(std::vector<Box> proposals, std::vector<BoundingBox> *objects, 
                             float scaleInfo, int src_rows, int src_cols,
                             std::vector<std::string> classNames,
                             std::unordered_map<std::string, float> classAttributes)
    {
        qsort_descent_inplace(proposals);
        std::vector<int> picked;
        nms_sorted_bboxes(proposals, picked, classAttributes.find("iou_thresh")->second);

        int count = picked.size();

        objects->resize(count);
        for (int i = 0; i < count; i++)
        {
            if (i >= classAttributes.find("topk")->second)
            {
                break;
            }

            BoundingBox post_box;
            float x0 = (proposals[picked[i]].xyxy[0] / scaleInfo);
            float y0 = (proposals[picked[i]].xyxy[1] / scaleInfo);
            float x1 = (proposals[picked[i]].xyxy[2] / scaleInfo);
            float y1 = (proposals[picked[i]].xyxy[3] / scaleInfo);

            // for (int l = 0; l < 5; l++)
            // {
            //     auto lx = objects[i].landmark[l].x;
            //     auto ly = objects[i].landmark[l].y;
            //     objects[i].landmark[l] = cv::Point2f((lx - tmp_w) * ratio_x, (ly - tmp_h) * ratio_y);
            // }

            x0 = std::max(std::min(x0, (float)(src_cols - 1)), 0.f);
            y0 = std::max(std::min(y0, (float)(src_rows - 1)), 0.f);
            x1 = std::max(std::min(x1, (float)(src_cols - 1)), 0.f);
            y1 = std::max(std::min(y1, (float)(src_rows - 1)), 0.f);

            post_box.x1 = x0;
            post_box.y1 = y0;
            post_box.x2 = x1;
            post_box.y2 = y1;
            post_box.category = classNames[proposals[picked[i]].index];
            post_box.detect_confidence = proposals[picked[i]].score;
            objects->push_back(post_box);
            printf("[%s Line:%d] ------------final-box x1:%.0f y1:%.0f x2:%.0f y2:%.0f conf:%.2f class_name:%s------------\n", __FUNCTION__, __LINE__, post_box.x1, post_box.y1, post_box.x2, post_box.y2, post_box.detect_confidence, post_box.category.c_str());
        }
    }
} // namespace detection