/*
 * Utility functions for the postprocess functions.
 *
 * Copyright (C) 2021 Kneron, Inc. All rights reserved.
 *
 */
#include <math.h>
#include <stdbool.h>
#include <string.h>


#include "ncpu_gen_struct.h"
#include "demo_post_utils.h"

float prob_thresh_yolov5 = 0.3;
float nms_thresh_yolov5 = 0.5;

float unpass_score = -999.0;    // used as box filter

// model 211
const float anchors_yolov5[3][3][2] = {
    {{10, 13}, {16, 30}, {33, 23}},
    {{30, 61}, {62, 45}, {59, 119}},
    {{116 ,90}, {156, 198}, {373, 326}}};

void *get_gp(void **gp, size_t len)
{
    if (*gp == NULL) {
        *gp = malloc(len);

        if (*gp == NULL)
            printf("ncpu: ERROR: malloc %u in get_gp()\n", len);
    }

    return *gp;
}

struct yolo_v5_post_globals_s *get_yolov5_gp()
{
    return (struct yolo_v5_post_globals_s *)get_gp((void **)&yolov5_gp, sizeof(struct yolo_v5_post_globals_s));
}

void free_gp(void *gp)
{
    free(gp);
    gp = NULL;
}

float do_div_scale(float v, int div, float scale) {
    return ((v / div) / scale);
}

float do_div_scale_optim(float v, float scale) {
    return (v * scale);
}

uint32_t round_up(uint32_t num) {
    return ((num + (KDP_COL_MIN - 1)) & ~(KDP_COL_MIN - 1));
}

int int_comparator(const void *pa, const void *pb) {
    int diff = *(int *)pa - *(int *)pb;

    if (diff < 0)
        return 1;
    else if (diff > 0)
        return -1;
    return 0;
}

int float_comparator(float a, float b) {
    float diff = a - b;

    if (diff < 0)
        return 1;
    else if (diff > 0)
        return -1;
    return 0;
}

int box_score_comparator(const void *pa, const void *pb) {
    float a, b;

    a = ((struct bounding_box_s *) pa)->score;
    b = ((struct bounding_box_s *) pb)->score;

    if (a == b) {   // take box with bigger area
        float area_a = box_area((struct bounding_box_s *)pa);
        float area_b = box_area((struct bounding_box_s *)pb);
        return float_comparator(area_a, area_b);
    }

    return float_comparator(a, b);
}

int box_lm_score_comparator(const void *pa, const void *pb) {
    float a, b;

    a = ((struct bounding_box_landmark_s *) pa)->score;
    b = ((struct bounding_box_landmark_s *) pb)->score;

    if (a == b) {   // take box with bigger area
        float area_a = box_lm_area((struct bounding_box_landmark_s *)pa);
        float area_b = box_lm_area((struct bounding_box_landmark_s *)pb);
        return float_comparator(area_a, area_b);
    }

    return float_comparator(a, b);
}

float overlap(float l1, float r1, float l2, float r2) {
    /**
     * legacy function of box_area(), but better performance
    */
    float left = l1 > l2 ? l1 : l2;
    float right = r1 < r2 ? r1 : r2;
    return right - left;
}

/**
 * Calculate the area of a bounding box.
 */
float box_area(struct bounding_box_s *box) {
    return fmax(0, box->y2 - box->y1) * fmax(0, box->x2 - box->x1);
}

/**
 * Calculate the area of a bounding box.
 */
float box_lm_area(struct bounding_box_landmark_s *box) {
    return fmax(0, box->y2 - box->y1) * fmax(0, box->x2 - box->x1);
}

float box_intersection(struct bounding_box_s *a, struct bounding_box_s *b) {
    /**
     * e2e version (commit: fc61d8a05bf481374341b6a86e6d2c2cea90451a)
    struct bounding_box_s overlap;

    overlap.x1 = fmax(a->x1, b->x1);
    overlap.y1 = fmax(a->y1, b->y1);
    overlap.x2 = fmin(a->x2, b->x2);
    overlap.y2 = fmin(a->y2, b->y2);

    float area = box_area(&overlap);
    return area;
    */

    /* optimized version */
    float w, h, area;

    w = overlap(a->x1, a->x2, b->x1, b->x2);
    h = overlap(a->y1, a->y2, b->y1, b->y2);

    if (w < 0 || h < 0)
        return 0;

    area = w * h;
    return area;
}

float box_union(struct bounding_box_s *a, struct bounding_box_s *b) {
    /**
     * e2e version (commit: fc61d8a05bf481374341b6a86e6d2c2cea90451a)

    float i, u;

    i = box_intersection(a, b);
    u = box_area(a) + box_area(b) - i;

    return u;
    */

    /* optimized version */
    float i, u;

    i = box_intersection(a, b);
    u = (a->y2 - a->y1) * (a->x2 - a->x1) + (b->y2 - b->y1) * (b->x2 - b->x1) - i;

    return u;
}

float box_iou(struct bounding_box_s *a, struct bounding_box_s *b) {
    /**
     * e2e version (commit: fc61d8a05bf481374341b6a86e6d2c2cea90451a)

    float c;
    float intersection_a_b = box_intersection(a, b);
    float union_a_b = box_union(a, b);

    c = intersection_a_b / union_a_b;

    return c;
    */

    /* optimized version */
    return box_intersection(a, b) / box_union(a, b);
}


/*
 * Get the index corresponding to given channel, row, and column indices (unpadded).
 */
uint32_t _get_index(struct output_node node, uint32_t ch_idx, uint32_t row_idx, uint32_t col_idx, bool padded_col) {
    int col = padded_col ? node.col_len: node.col;
#ifdef KL520
    uint32_t index = row_idx * node.ch * col + ch_idx * col + col_idx;
#else
    uint32_t index = ch_idx * node.row * col + row_idx * col + col_idx;
#endif
    return index;
}

/*
 * Get the index corresponding to given channel, row, and column indices (padded).
 */
uint32_t get_index(struct output_node node, uint32_t ch_idx, uint32_t row_idx, uint32_t col_idx) {
    return _get_index(node, ch_idx, row_idx, col_idx, true);
}

/**
 * Get the data pointer corresponding to given channel, row, and column indices.
 */
int8_t *get_data(struct output_node node, uint32_t ch_idx, uint32_t row_idx, uint32_t col_idx) {
    uint32_t index = get_index(node, ch_idx, row_idx, col_idx);
    return node.base_ptr + index;
}

/**
 * Get the output node information.
 */
void get_output_node(struct output_node *out_node, struct kdp_image_s *image_p, int node_num) {
#if defined(KL530) || defined(KL630)
    tensor_hdl_t handle = mdl_parse_get_outputs_tensor_hdl(image_p->pParsedModelFlatbuffer->out_node_hdl);

    mdl_shape_t shape;
    mdl_quant_factor_t quant;
    mdl_parse_get_shape(handle, node_num, &shape);
    mdl_parse_get_quant_vector_cell(handle, node_num, 0, &quant);   // check 0 input

#ifdef VATICS_PLATFORM
    out_node->base_ptr = (int8_t *)(image_p->postproc.output_mem_addr + mdl_parse_get_out_addr(image_p->pParsedModelFlatbuffer->out_node_hdl, handle, node_num));
#else
    out_node->base_ptr = (int8_t *)mdl_parse_get_out_addr(image_p->pParsedModelFlatbuffer->parser_session_hdl, handle, node_num);
#endif
    out_node->ch = shape.ch;
    out_node->row = shape.h;
    out_node->col = shape.w;
    out_node->radix = quant.radix;
    out_node->scale = quant.scale;
#else
#ifdef KL520
    struct out_node_s *out_p;
    out_p = (struct out_node_s *)((kdp_size_t)POSTPROC_OUT_NODE(image_p) + node_num * sizeof(struct out_node_s));
#else // KL720
    struct Out_Node *out_p;
    out_p = fw_get_out_node(image_p->pParsedModel, node_num);
#endif
    out_node->base_ptr = (int8_t *)OUT_NODE_ADDR(out_p);
    out_node->ch = OUT_NODE_CH(out_p);
    out_node->row = OUT_NODE_ROW(out_p);
    out_node->col = OUT_NODE_COL(out_p);
    out_node->radix = OUT_NODE_RADIX(out_p);
    out_node->scale = *(float *)&OUT_NODE_SCALE(out_p);
#endif
    out_node->col_len = round_up(out_node->col);
}

/**
 * Performs NMS on the potential boxes.
 */
int nms_bbox(struct bounding_box_s *potential_boxes, struct bounding_box_s *temp_results,
             int class_num, int good_box_count, int max_boxes, int single_class_max_boxes, struct bounding_box_s *results,
             float score_thresh, float iou_thresh, float nms_mode) {
    int good_result_count = 0;

    // check overlap between all boxes and not just those from same class
    if (nms_mode == NMS_ALL_CLASS) {
        if (good_box_count == 1) {
            memcpy(&results[good_result_count], &potential_boxes[0], sizeof(struct bounding_box_s));
            good_result_count++;
        } else if (good_box_count >= 2) {
            // sort boxes based on the score
            qsort(potential_boxes, good_box_count, sizeof(struct bounding_box_s), box_score_comparator);
            for (int j = 0; j < good_box_count; j++) {
                // if the box score is too low or is already filtered by previous box
                if (potential_boxes[j].score < score_thresh)
                    continue;

                // filter out overlapping, lower score boxes
                for (int k = j + 1; k < good_box_count; k++)
                    if (box_iou(&potential_boxes[j], &potential_boxes[k]) > iou_thresh)
                        potential_boxes[k].score = unpass_score;

                // keep boxes with highest scores, up to a certain amount
                memcpy(&results[good_result_count], &potential_boxes[j], sizeof(struct bounding_box_s));
                good_result_count++;
                if (good_result_count == max_boxes)
                    break;
            }
        }
    } else {    // check overlap between only boxes from same class
        for (int i = 0; i < class_num; i++) {
            int class_good_result_count = 0;
            if (good_result_count == max_boxes) // break out of outer loop as well for future classes
                break;

            int class_good_box_count = 0;

            // find all boxes of a specific class
            for (int j = 0; j < good_box_count; j++) {
                if (potential_boxes[j].class_num == i) {
                    memcpy(&temp_results[class_good_box_count], &potential_boxes[j], sizeof(struct bounding_box_s));
                    class_good_box_count++;
                }
            }

            if (class_good_box_count == 1) {
                memcpy(&results[good_result_count], temp_results, sizeof(struct bounding_box_s));
                good_result_count++;
            } else if (class_good_box_count >= 2) {
                // sort boxes based on the score
                qsort(temp_results, class_good_box_count, sizeof(struct bounding_box_s), box_score_comparator);
                for (int j = 0; j < class_good_box_count; j++) {
                    // if the box score is too low or is already filtered by previous box
                    if (temp_results[j].score < score_thresh)
                        continue;

                    // filter out overlapping, lower score boxes
                    for (int k = j + 1; k < class_good_box_count; k++)
                        if (box_iou(&temp_results[j], &temp_results[k]) > iou_thresh)
                            temp_results[k].score = unpass_score;

                    // keep boxes with highest scores, up to a certain amount
                    if ((good_result_count == max_boxes) || (class_good_result_count == single_class_max_boxes))
                        break;
                    memcpy(&results[good_result_count], &temp_results[j], sizeof(struct bounding_box_s));
                    good_result_count++;
                    class_good_result_count++;
                }
            }
        }
    }

    return good_result_count;
}

/**
 * Remap one bounding box to original image coordinates.
 */
void remap_bbox(struct kdp_image_s *image_p, int index, struct bounding_box_s *box, int need_scale) {
    // original box values are percentages, scale to model sizes
    if (need_scale) {
        box->x1 *= DIM_INPUT_COL(image_p, index);
        box->y1 *= DIM_INPUT_ROW(image_p, index);
        box->x2 *= DIM_INPUT_COL(image_p, index);
        box->y2 *= DIM_INPUT_ROW(image_p, index);
    }

    // scale from model sizes to original input sizes
    box->x1 = (box->x1 - RAW_PAD_LEFT(image_p, index)) * RAW_SCALE_WIDTH(image_p, index) + RAW_CROP_LEFT(image_p, index);
    box->y1 = (box->y1 - RAW_PAD_TOP(image_p, index)) * RAW_SCALE_HEIGHT(image_p, index) + RAW_CROP_TOP(image_p, index);
    box->x2 = (box->x2 - RAW_PAD_LEFT(image_p, index)) * RAW_SCALE_WIDTH(image_p, index) + RAW_CROP_LEFT(image_p, index);
    box->y2 = (box->y2 - RAW_PAD_TOP(image_p, index)) * RAW_SCALE_HEIGHT(image_p, index) + RAW_CROP_TOP(image_p, index);

    // clip to boundaries of image
    box->x1 = (int)((box->x1 < 0 ? 0 : box->x1) + (float)0.5);
    box->y1 = (int)((box->y1 < 0 ? 0 : box->y1) + (float)0.5);
    box->x2 = (int)((box->x2 > RAW_INPUT_COL(image_p, index) ? RAW_INPUT_COL(image_p, index) : box->x2) + (float)0.5);
    box->y2 = (int)((box->y2 > RAW_INPUT_ROW(image_p, index) ? RAW_INPUT_ROW(image_p, index) : box->y2) + (float)0.5);
}
