#ifndef ANSONNXSEG_H
#define ANSONNXSEG_H
#pragma once
#include "ANSEngineCommon.h"
#include "engine.h"

namespace ANSCENTER {
    class ANSENGINE_API ANSONNXSEG :public ANSODBase {
    public:
        virtual bool Initialize(std::string licenseKey, ModelConfig modelConfig, const std::string& modelZipFilePath, const std::string& modelZipPassword, std::string& labelMap) override;
        virtual bool LoadModel(const std::string& modelZipFilePath, const std::string& modelZipPassword)override;
        virtual bool LoadModelFromFolder(std::string licenseKey, ModelConfig modelConfig, std::string modelName, std::string className, const std::string& modelFolder, std::string& labelMap)override;
        virtual bool OptimizeModel(bool fp16, std::string& optimizedModelFolder);
        std::vector<Object> RunInference(const cv::Mat& input);
        std::vector<Object> RunInference(const cv::Mat& input, const std::string& camera_id);
        bool Destroy();
        ~ANSONNXSEG();
        int getInstanceId() const { return instanceId_; }
    private:
        std::string     _modelFilePath;
        bool            _modelLoadValid;
        bool            _fp16{ false };
        size_t vectorProduct(const std::vector<int64_t>& vector);
        void letterBox(const cv::Mat& image,
            cv::Mat& outImage,
            const cv::Size& newShape,
            const cv::Scalar& color = cv::Scalar(114, 114, 114),
            bool auto_ = true,
            bool scaleFill = false,
            bool scaleUp = true,
            int stride = 32);
        BoundingBox scaleCoords(const cv::Size& letterboxShape,
            const BoundingBox& coords,
            const cv::Size& originalShape,
            bool p_Clip = true);
        std::vector<cv::Scalar> generateColors(const std::vector<std::string>& classNames, int seed = 42);
        cv::Mat sigmoid(const cv::Mat& src);
        void NMSBoxes(const std::vector<BoundingBox>& boxes,
            const std::vector<float>& scores,
            float scoreThreshold,
            float nmsThreshold,
            std::vector<int>& indices);
        void drawSegmentations(cv::Mat& image,
            const std::vector<Object>& results,
            float maskAlpha = 0.5f) const;
        void drawSegmentationsAndBoxes(cv::Mat& image,
            const std::vector<Object>& results,
            float maskAlpha) const;
        const std::vector<cv::Scalar>& getClassColors() const { return classColors; }

        void warmupModel();
        bool Init(const std::string& modelPath, bool useGPU=true, int deviceId = 0);
        cv::Mat preprocess(const cv::Mat& image,float*& blobPtr,std::vector<int64_t>& inputTensorShape);
        std::vector<Object> postprocess(const cv::Size& origSize,const cv::Size& letterboxSize,
                                       const std::vector<Ort::Value>& outputs, const std::string& camera_id);
        std::vector<Object> segment(const cv::Mat& image, const std::string& camera_id);
        std::vector<cv::Point2f> maskToPolygon(const cv::Mat& binaryMask,
            const cv::Rect& boundingBox,
            float simplificationEpsilon = 2.0f,
            int minContourArea = 10);
    private:
        static std::atomic<int> instanceCounter_;  // Thread-safe counter
        int instanceId_;
        int deviceId_ = 0;
        Ort::Env env{ nullptr };                         // ONNX Runtime environment
        Ort::SessionOptions sessionOptions{ nullptr };   // Session options for ONNX Runtime
        Ort::Session session{ nullptr };                 // ONNX Runtime session for running inference
        bool isDynamicInputShape{};                     // Flag indicating if input shape is dynamic
        cv::Size inputImageShape;                       // Expected input image shape for the model

        // Vectors to hold allocated input and output node names
        std::vector<Ort::AllocatedStringPtr> inputNameAllocs;
        std::vector<const char*> inputNames;
        std::vector<Ort::AllocatedStringPtr> outputNameAllocs;
        std::vector<const char*> outputNames;
        size_t numInputNodes, numOutputNodes;          // Number of input and output nodes in the model
        std::vector<cv::Scalar>  classColors;

        float m_imgWidth = 0;
        float m_imgHeight = 0;
    };
}
#endif
#pragma once