#ifndef RETINAFACE_H
#define RETINAFACE_H
#pragma once
#include "ANSFRCommon.h"
#define CLIP(a, b, c) (MAX(MIN(a, c), b)) // MIN, MAX defined in opencv
namespace ANSCENTER {
    // Retina will be FD subclass
    class RetinaFaceTRT {
    public:
        RetinaFaceTRT();
        ~RetinaFaceTRT();
        std::vector<struct Bbox> FindFace(cv::Mat& img);
        bool Initialize(const std::string engineFile, 
                        int frameWidth, 
                        int frameHeight, 
                        std::string inputName, 
                        std::vector<std::string> outputNames,
                        std::vector<int> inputShape, 
                        int maxBatchSize, 
                        int maxFacesPerScene, 
                        float nms_threshold, 
                        float bbox_threshold);

    private:
        ANSCENTER::SPDLogger& _logger = ANSCENTER::SPDLogger::GetInstance("ANSFD", true);
        TRTLogger           m_logger;

        int m_frameWidth, m_frameHeight, m_INPUT_C, m_INPUT_H, m_INPUT_W, m_INPUT_SIZE, m_OUTPUT_SIZE_BASE, m_maxBatchSize, m_maxFacesPerScene;
        float m_nms_threshold, m_bbox_threshold, m_scale_h, m_scale_w;
        cv::Mat m_input;
        float* m_output0, * m_output1;
        std::vector<struct Bbox> m_outputBbox;

        nvinfer1::ICudaEngine* m_engine;
        nvinfer1::IExecutionContext* m_context;
        cudaStream_t stream;
        void* buffers[3];
        int inputIndex, outputIndex0, outputIndex1;
    private:
        void LoadEngine(const std::string engineFile);   
        void PreInference(std::string inputNames, std::vector<std::string> outputNames);
        void RunInference(float* input, float* output0, float* output1);
        void PreProcess(cv::Mat& img);
        void PostProcessing(float* bbox, float* conf);
        void CreateAnchorRetinaface(std::vector<AnchorBox>& anchor, int w, int h);
        static inline bool MCMP(Bbox a, Bbox b);
        void NMS(std::vector<Bbox>& input_boxes, float NMS_THRESH);
    };
}

#endif