/*
 * Utility function headers for the postprocess functions.
 *
 * Copyright (C) 2021 Kneron, Inc. All rights reserved.
 *
 */
#ifndef DEMO_POST_UTILS_H
#define DEMO_POST_UTILS_H

#include <stdlib.h>
#include <stdio.h>
#include "utils.h"

#define YOLO_BOX_FIX_CH                 (5)

#define YOLO_V5_O0_GRID_W               (60)
#define YOLO_V5_O0_GRID_H               (32)
#define YOLO_V5_O0_GRID_MAX             (YOLO_V5_O0_GRID_W * YOLO_V5_O0_GRID_H)
#define YOLO_V5_O1_GRID_W               (30)
#define YOLO_V5_O1_GRID_H               (16)
#define YOLO_V5_O1_GRID_MAX             (YOLO_V5_O1_GRID_W * YOLO_V5_O1_GRID_H)
#define YOLO_V5_O2_GRID_W               (15)
#define YOLO_V5_O2_GRID_H               (8)
#define YOLO_V5_O2_GRID_MAX             (YOLO_V5_O2_GRID_W * YOLO_V5_O2_GRID_H)
#define YOLO_V5_CELL_BOX_NUM            (3)
#define YOLO_V5_MAX_BOX_NUM             MIN(500, MAX(MAX(YOLO_V5_O1_GRID_MAX, YOLO_V5_O2_GRID_MAX), YOLO_V5_O0_GRID_MAX) * YOLO_V5_CELL_BOX_NUM)

#define KDP_COL_MIN                     16 /* Bytes, i.e. 128 bits */

#define NMS_ALL_CLASS                   (0)
#define NMS_GROUP_CLASS                 (1)
#define NMS_SINGLE_CLASS                (2)

// thresholds for various solutions
// prob_thresh is the score threshold
// nms_thresh is the iou threshold

extern float nms_thresh_yolov5;
extern float prob_thresh_yolov5;
extern const float anchors_yolov5[3][3][2];

// how a output node layout
struct output_node {
    int8_t *base_ptr;
    uint32_t ch;
    uint32_t row;
    uint32_t col;
    uint32_t col_len;
    uint32_t radix;
    float scale;
};

/* Shared global variable area among models */
struct yolo_v5_post_globals_s {
    struct bounding_box_s bboxes[YOLO_V5_MAX_BOX_NUM];
    struct bounding_box_s result_tmp_s[YOLO_V5_MAX_BOX_NUM];
};

struct yolo_v5_post_globals_s     *yolov5_gp;

struct yolo_v5_post_globals_s *get_yolov5_gp();
void free_gp(void *gp);

// conversion functions
float do_div_scale(float v, int div, float scale);
float do_div_scale_optim(float v, float scale);
uint32_t round_up(uint32_t num);

// bounding box helpers
int int_comparator(const void *pa, const void *pb);
int float_comparator(float a, float b);
int box_score_comparator(const void *pa, const void *pb);
int box_lm_score_comparator(const void *pa, const void *pb);
float box_intersection(struct bounding_box_s *a, struct bounding_box_s *b);
float box_union(struct bounding_box_s *a, struct bounding_box_s *b);
float box_iou(struct bounding_box_s *a, struct bounding_box_s *b);
float box_area(struct bounding_box_s *box);
float box_lm_area(struct bounding_box_landmark_s *box);
float overlap(float l1, float r1, float l2, float r2);

// functions to get parameters from 520/720 specific structures to common structure
uint32_t get_index(struct output_node node, uint32_t ch_idx, uint32_t row_idx, uint32_t col_idx);
int8_t *get_data(struct output_node node, uint32_t ch_idx, uint32_t row_idx, uint32_t col_idx);
void get_output_node(struct output_node *out_node, struct kdp_image_s *image_p, int node_num);

// more bounding box helpers
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);
void remap_bbox(struct kdp_image_s *image_p, int index, struct bounding_box_s *box, int need_scale);

#endif
