#ifndef _FACE_DETECTION_H
#define _FACE_DETECTION_H

#include <iostream>
#include <vector>
#include "k230_math.h"
#include "ai_utils.h"
#include "kpu.h"
#include "ai2d.h"

using std::vector;

#define LOC_SIZE 4
#define CONF_SIZE 2
#define LAND_SIZE 10
#define PI 3.1415926

/**
 * @brief Structure representing a bounding box.
 */
typedef struct Bbox
{
    float x;  ///< Top-left x-coordinate
    float y;  ///< Top-left y-coordinate
    float w;  ///< Width of the bounding box
    float h;  ///< Height of the bounding box
} Bbox;

/**
 * @brief Structure representing sparse facial landmarks (typically 5 points × 2 coordinates).
 */
typedef struct SparseLandmarks
{
    float points[10];  ///< Array of 10 float values (x1, y1, x2, y2, ..., x5, y5)
} SparseLandmarks;

/**
 * @brief Structure representing an object candidate used in NMS (Non-Maximum Suppression).
 */
typedef struct NMSRoiObj
{
    int index;          ///< Object index
    float confidence;   ///< Confidence score
} NMSRoiObj;

/**
 * @brief Structure representing final face detection result information.
 */
typedef struct FaceDetectionInfo
{
    Bbox bbox;                ///< Bounding box of the detected face
    SparseLandmarks sparse_kps;  ///< Facial landmarks (keypoints)
    float score;              ///< Confidence score for the detection
} FaceDetectionInfo;

/**
 * @class FaceDetection
 * @brief Class implementing face detection using a pre-trained KPU model.
 *
 * This class handles the entire face detection pipeline, including preprocessing,
 * model inference, post-processing, and visualization of detection results.
 */
class FaceDetection
{
public:
    /**
     * @brief Constructor for FaceDetection class.
     * @param kmodel_file Path to the KModel file.
     * @param obj_thresh Object confidence threshold.
     * @param nms_thresh Non-Maximum Suppression (NMS) threshold.
     * @param image_size Input image frame size (channels, height, width).
     * @param debug_mode Debug flag (0 = off, 1 = on).
     */
    FaceDetection(std::string &kmodel_file, float obj_thresh, float nms_thresh,
                  FrameCHWSize image_size, int debug_mode);

    /**
     * @brief Destructor for FaceDetection class.
     */
    ~FaceDetection();

    /**
     * @brief Preprocess input tensor for model inference.
     * @param input_tensor Input tensor reference.
     */
    void pre_process(runtime_tensor &input_tensor);

    /**
     * @brief Perform model inference using the KPU.
     */
    void inference();

    /**
     * @brief Post-process model outputs and extract face detection results.
     * @param results Vector to store output detection results.
     */
    void post_process(vector<FaceDetectionInfo> &results);

    /**
     * @brief Draw detected faces and landmarks on the image frame.
     * @param draw_frame Input/output frame to draw results.
     * @param results Vector of detection results to render.
     */
    void draw_results(cv::Mat &draw_frame, vector<FaceDetectionInfo> &results);

    /**
     * @brief Run the complete face detection pipeline.
     * @param input_tensor Input tensor for inference.
     * @param results Vector to store final detection results.
     */
    void run(runtime_tensor &input_tensor, vector<FaceDetectionInfo> &results);

protected:
    /**
     * @brief Apply local softmax to an array.
     * @param x Input array.
     * @param dx Output array after softmax.
     * @param len Length of the array.
     */
    void local_softmax(float *x, float *dx, uint32_t len);

    /**
     * @brief Process confidence scores and select candidate objects.
     * @param conf Confidence score array.
     * @param so Output array of selected NMS candidates.
     * @param size Input array size.
     * @param obj_cnt Output number of valid objects.
     */
    void deal_conf(float *conf, NMSRoiObj *so, int size, int &obj_cnt);

    /**
     * @brief Process and decode location information (bounding boxes).
     * @param loc Input location array.
     * @param boxes Output bounding boxes.
     * @param size Number of detections.
     * @param obj_cnt Object count.
     */
    void deal_loc(float *loc, float *boxes, int size, int &obj_cnt);

    /**
     * @brief Process and decode landmark information.
     * @param landms Input landmark array.
     * @param landmarks Output landmark array.
     * @param size Number of detections.
     * @param obj_cnt Object count.
     */
    void deal_landms(float *landms, float *landmarks, int size, int &obj_cnt);

    /**
     * @brief Retrieve a bounding box by index.
     * @param boxes Pointer to bounding box data.
     * @param obj_index Object index.
     * @return Bounding box structure.
     */
    Bbox get_box(float *boxes, int obj_index);

    /**
     * @brief Retrieve facial landmarks by index.
     * @param landmarks Pointer to landmarks data.
     * @param obj_index Object index.
     * @return SparseLandmarks structure.
     */
    SparseLandmarks get_landmark(float *landmarks, int obj_index);

    /**
     * @brief Compute overlap between two 1D segments.
     * @param x1 First segment start position.
     * @param w1 First segment width.
     * @param x2 Second segment start position.
     * @param w2 Second segment width.
     * @return Overlap size.
     */
    float overlap(float x1, float w1, float x2, float w2);

    /**
     * @brief Compute intersection area between two bounding boxes.
     * @param a First box.
     * @param b Second box.
     * @return Intersection area.
     */
    float box_intersection(Bbox a, Bbox b);

    /**
     * @brief Compute union area between two bounding boxes.
     * @param a First box.
     * @param b Second box.
     * @return Union area.
     */
    float box_union(Bbox a, Bbox b);

    /**
     * @brief Compute Intersection over Union (IoU) between two boxes.
     * @param a First box.
     * @param b Second box.
     * @return IoU value (0.0 to 1.0).
     */
    float box_iou(Bbox a, Bbox b);

    /**
     * @brief Convert detection results from model coordinates to original frame coordinates.
     * @param frame_size Original frame size.
     * @param results Vector of detection results to adjust.
     */
    void get_final_box(FrameCHWSize &frame_size, vector<FaceDetectionInfo> &results);

private:
    runtime_tensor model_input_tensor_;  ///< KPU model input tensor
    FrameCHWSize image_size_;            ///< Input image size (CHW)
    FrameCHWSize input_size_;            ///< Model input size (CHW)

    int min_size_;                       ///< Minimum object size
    float obj_thresh_;                   ///< Object confidence threshold
    float nms_thresh_;                   ///< NMS IoU threshold
    int objs_num_;                       ///< Number of detected objects
    NMSRoiObj *so_;                      ///< Temporary array for NMS candidates
    float *boxes_;                       ///< Array of decoded bounding boxes
    float *landmarks_;                   ///< Array of decoded landmarks
    int debug_mode_;                     ///< Debug flag

    KPU *kpu = nullptr;                  ///< Pointer to KPU (Kendryte Processing Unit) instance
    AI2D *ai2d = nullptr;                ///< Pointer to AI2D preprocessing module
};

#endif
