#include "ANSLIB.h"
#define RETAINFRAMES 80
#define FILTERFRAMES 10

class CUSTOM_API ANSCustomWD : public IANSCustomClass
{
    struct ImageSection {
        cv::Rect region;
        int priority;
        ImageSection(const cv::Rect& r) : region(r), priority(0) {}
    };
    private:
        using ANSLIBPtr = std::unique_ptr<ANSCENTER::ANSLIB, decltype(&ANSCENTER::ANSLIB::Destroy)>;
        ANSLIBPtr _detector{ nullptr, &ANSCENTER::ANSLIB::Destroy };
        ANSLIBPtr _filter{ nullptr, &ANSCENTER::ANSLIB::Destroy };     
        int	engineType;
        std::string _detectorModelName;
        std::string _detectorClassName;
        int _detectorModelType; // Assuming 4 represents TensorRT YoloV11 
        int _detectorDetectionType; // Assuming 1 represents object detection

        std::string _filterModelName;
        std::string _filterClassName;
        std::string _filterLabelMap;

        int _filterModelType; // Assuming 1 represents OpenVINO YoloV12
        int _filterDetectionType; // Assuming 1 represents object detection

        std::recursive_mutex	_mutex;
        cv::Rect                _detectedArea;// Area where weapon is detected
        int                     _retainDetectedArea{ 0 };
        bool				    _isWeaponDetected{ false };
		float                   _detectionScoreThreshold{ 0.5 };
        cv::Size previousImageSize = cv::Size(0, 0);
        std::vector<ImageSection> cachedSections;
        int                     _currentPriority{ 0 }; // None 
        int                     _realWeaponCheck{ 0 };
        bool 				    _isRealWeaponFrame{ false };
        void                    UpdateNoDetectionCondition();
        std::vector<ANSCENTER::Object> RunFilterGetPersons(const cv::Mat& frame);
        // Function to seperate screen size
        double calculateDistanceToCenter(const cv::Point& center, const cv::Rect& rect);
        std::vector<ImageSection> divideImage(const cv::Mat& image);
        int getHighestPriorityRegion();
        int getLowestPriorityRegion();
        cv::Rect getRegionByPriority(int priority);

        // Utilities
        float calculateIoU(const cv::Rect& box1, const cv::Rect& box2);
		bool IsOverlapping(const ANSCENTER::Object& obj, const std::vector<ANSCENTER::Object>& objectList, float iouThreshold);
        void UpdateActiveROI(const cv::Mat& frame, ANSCENTER::Object detectedObj);
        std::vector<ANSCENTER::Object> ProcessExistingDetectedArea(
            const cv::Mat& frame, const std::string& camera_id
#ifdef FNS_DEBUG
            , cv::Mat& draw
#endif
        );
        bool ProcessWeaponDetection(
            const cv::Mat& frame,
            ANSCENTER::Object& detectedObj,
            std::vector<ANSCENTER::Object>& output,
            const std::vector<ANSCENTER::Object>& personDetections,
            bool filterHadResults
#ifdef FNS_DEBUG
            , cv::Mat& draw
#endif
        );
        void AddConfirmedWeaponDetection(
            const cv::Mat& frame,
            ANSCENTER::Object& detectedObj,
            std::vector<ANSCENTER::Object>& output);

        void ProcessNewDetectedArea(
            const cv::Mat& frame,
            const std::string& camera_id,
            std::vector<ANSCENTER::Object>& output
#ifdef FNS_DEBUG
            , cv::Mat& draw
#endif
        );
        void UpdateDetectedAreaFromObject(const cv::Mat& frame, const ANSCENTER::Object& detectedObj);
    public:
        bool Initialize(const std::string& modelDiretory,float detectionScoreThreshold, std::string& labelMap)override;
        bool OptimizeModel(bool fp16)override;
        std::vector<CustomObject> RunInference(const cv::Mat& input)override;
        std::vector<CustomObject> RunInference(const cv::Mat& input, const std::string& camera_id)override;
        bool Destroy()override;
        virtual bool ConfigureParameters(CustomParams& param) override;
        ANSCustomWD();
        ~ANSCustomWD();
};