#include "ANSLIB.h"
#include <deque>
#include <unordered_set>

#define RETAINFRAMES 80
#define FILTERFRAMES 10
//#define FNS_DEBUG
class CUSTOM_API ANSCustomFS : public IANSCustomClass
{
    typedef std::pair<cv::Scalar, cv::Scalar> Range;
    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)>;

    int                     engineType{ 0 };
    ANSLIBPtr _motiondetector{ nullptr, &ANSCENTER::ANSLIB::Destroy };
    ANSLIBPtr _detector{ nullptr, &ANSCENTER::ANSLIB::Destroy };
    ANSLIBPtr _filter{ nullptr, &ANSCENTER::ANSLIB::Destroy };

    std::recursive_mutex    _mutex;
    cv::Rect                _detectedArea;
    cv::Rect                _previousDetectedArea;
    int                     _retainDetectedArea{ 0 };
    bool                    _isFireNSmokeDetected{ false };
    float                   _detectionScoreThreshold{ 0.5 };
    std::vector<cv::Mat>    _frameHistory;
    std::deque <cv::Mat>    _frameQueue;
    std::string             _filterLabels;
    int                     _realFireCheck{ 0 };
    bool                    _isRealFireFrame{ false };
    double                  _hsvThreshold;
    std::vector<Range>      _smoke_colour;
    std::vector<Range>      _fire_colour;

    // ----- Model config -----
    int _detectorModelType;
    int _detectorDetectionType;
    int _filterModelType;
    int _filterDetectionType;
    int _motionModelType;
    int _motionDetectionType;

    // Add these constants at class level
    int MIN_ROI_SIZE = 50;
    float SMOKE_CONFIDENCE_THRESHOLD = 0.7f;
    float DETECTION_CONFIDENCE_THRESHOLD = 0.5f;
    float MAX_AREA_OVERLAP = 0.5f;
    int MAX_MOTION_TRACKING = 10;
    int FILTER_VERIFICATION_FRAMES = 10;
    int MIN_IMAGE_SIZE = 10;
    float INCLUSIVE_IOU_THRESHOLD = 0.001f;
    float MOTION_SENSITIVITY = 0.5f;

    // Excluded filter class IDs
    const std::unordered_set<int> EXCLUDED_FILTER_CLASSES =
    {
        13,  // bench
        59,  // bed
        60,  // dining table
        68,  // microwave
        69,  // oven
        70   // toaster
    };

    cv::Size previousImageSize = cv::Size(0, 0);
    std::vector<ImageSection> cachedSections;
    int                     _currentPriority{ 0 };
    bool                    _readROIs{ false };
    std::vector<cv::Rect>   _exclusiveROIs;
    float                   _smokeDetetectionThreshold{ 0 };
    float                   _motionSpecificity{ 0 };

    cv::Rect                GenerateMinimumSquareBoundingBox(const std::vector<ANSCENTER::Object>& detectedObjects, int minSize = 640);
    void                    UpdateNoDetectionCondition();
    bool                    detectStaticFire(std::deque<cv::Mat>& frameQueue);
    bool                    detectScreenFlicker(const std::vector<cv::Mat>& frames);
    bool                    detectReflection(const cv::Mat& frame);
    bool                    MajorityColourInFrame(cv::Mat frame, Range range, float area_threshold);
    bool                    MajorityColourInFrame(cv::Mat frame, std::vector<Range> ranges, float area_threshold);
    bool                    DetectFireNSmokeColourInFrame(const cv::Mat& frame, const cv::Rect& bBox, const std::vector<Range>& ranges, float area_threshold);
    bool                    IsFireDetected(const cv::Mat image, const cv::Rect bBox);
    bool                    IsSmokeDetected(const cv::Mat image, const cv::Rect bBox);
    float                   calculateIoU(const cv::Rect& box1, const cv::Rect& box2);
    bool                    IsOverlapping(const ANSCENTER::Object& obj, const std::vector<ANSCENTER::Object>& objectList, float iouThreshold = 0.5);
    bool                    IsOverlapping(const ANSCENTER::Object& obj, const std::vector<cv::Rect>& objectList, float iouThreshold = 0.5);

    bool                    IsROIOverlapping(const cv::Rect& obj, const std::vector<cv::Rect>& objectList, float iouThreshold = 0.5);
    bool                    IsDetectedAreaValid(const cv::Rect& area) const;
    bool                    IsFireOrSmoke(int classId, float confidence) const;
    bool                    ValidateMotionCorrelation(const std::vector<cv::Rect>& fireNSmokeRects) const;
    void                    UpdatePositiveDetection();
    void                    ResetDetectedArea();
    void                    UpdatePriorityRegion(const cv::Mat& frame);
    int                     CalculateOptimalCropSize(const cv::Mat& frame) const;
    void                    RefineDetectedArea(const cv::Mat& frame, const ANSCENTER::Object& detection);
    std::vector<CustomObject> ConvertToCustomObjects(const std::vector<ANSCENTER::Object>& objects);
    void                    ResetDetectionState();
    void                    GetModelParameters();
    std::vector<ANSCENTER::Object> ProcessExistingDetectedArea(const cv::Mat& frame, const std::string& camera_id, cv::Mat& draw);
    bool ProcessValidDetection(const cv::Mat& frame, const std::string& camera_id, cv::Mat& draw, ANSCENTER::Object& detectedObj, std::vector<ANSCENTER::Object>& output);
    void AddConfirmedDetection(ANSCENTER::Object& detectedObj, std::vector<ANSCENTER::Object>& output);
#ifdef FNS_DEBUG
    void DisplayDebugFrame(cv::Mat& draw) {
        cv::resize(draw, draw, cv::Size(1920, (1920 * draw.rows) / draw.cols));
        cv::imshow("Combined Detected Areas", draw);
        cv::waitKey(1);
    }
#endif
    // Function to separate 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);

    std::vector<ANSCENTER::Object> ProcessExistingDetectedArea(
        const cv::Mat& frame,
        const std::string& camera_id,
        const std::vector<cv::Rect>& fireNSmokeRects, cv::Mat& draw);
    bool ProcessDetectedObject(
        const cv::Mat& frame,
        ANSCENTER::Object& detectedObj,
        const std::string& camera_id,
        const std::vector<cv::Rect>& fireNSmokeRects,
        std::vector<ANSCENTER::Object>& output, cv::Mat& draw);
    bool ValidateWithFilter(
        const cv::Mat& frame,
        const ANSCENTER::Object& detectedObj,
        const std::string& camera_id,
        std::vector<ANSCENTER::Object>& output, cv::Mat& draw);
    std::vector<ANSCENTER::Object> FindNewDetectedArea(
        const cv::Mat& frame,
        const std::string& camera_id, cv::Mat& draw);

    std::vector<ANSCENTER::Object> FindMovementObjects(
        const cv::Mat& frame,
        const std::string& camera_id, cv::Mat& draw);

    std::vector<ANSCENTER::Object> FindExcludedObjects(
        const cv::Mat& frame,
        const std::string& camera_id, cv::Mat& draw);

public:
    bool Initialize(const std::string& modelDiretory, float detectionScoreThreshold, std::string& labelMap) override;
    bool OptimizeModel(bool fp16) override;
    bool ConfigureParameters(CustomParams& param) 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;
    ANSCustomFS();
    ~ANSCustomFS();
};