#ifndef YOLO_HPP__
#define YOLO_HPP__
#include <vector>
#include <iomanip>
#include "common/memory.hpp"
#include "infer/infer.hpp"
#include "common/image.hpp"
#include "common/data.hpp"
#include "infer/slice/slice.hpp"
#include "infer/trt/affine.hpp"

#include "common/data.hpp"

#ifdef TRT10
#include "common/tensorrt.hpp"
namespace TensorRT = TensorRT10;
#else
#include "common/tensorrt8.hpp"
namespace TensorRT = TensorRT8;
#endif


namespace yolo
{

    class YoloModelImpl : public Infer 
    {
    public:
        ModelType model_type_;

        std::vector<std::string> class_names_;
    
        // for sahi crop image
        std::shared_ptr<slice::SliceImage> slice_;
        std::shared_ptr<TensorRT::Engine> trt_;
        std::string engine_file_;
    
        tensor::Memory<int> box_count_;
    
        tensor::Memory<float> affine_matrix_;
        tensor::Memory<float> invert_affine_matrix_;
        tensor::Memory<float> mask_affine_matrix_;
        tensor::Memory<float> input_buffer_, bbox_predict_, segment_predict_, output_boxarray_;

        std::vector<std::shared_ptr<tensor::Memory<float>>> box_segment_cache_;
    
        int network_input_width_, network_input_height_;
        affine::Norm normalize_;
        std::vector<int> bbox_head_dims_;
        std::vector<int> segment_head_dims_;
        bool isdynamic_model_ = false;
        bool has_segment_ = false;
    
        float confidence_threshold_;
        float nms_threshold_;
    
        int num_classes_ = 0;

        int device_id_ = 0;

        YoloModelImpl() = default;
    
        virtual ~YoloModelImpl() = default;
    
        void adjust_memory(int batch_size);
    
        void preprocess(int ibatch, void *stream = nullptr);

        void cal_affine_matrix(affine::LetterBoxMatrix &affine, void *stream = nullptr);
        
        std::shared_ptr<data::InstanceSegmentMap> decode_segment(int imemory, float* pbox);
    
        bool load(const std::string &engine_file, ModelType model_type, const std::vector<std::string>& names, float confidence_threshold, float nms_threshold);
    
        virtual Result forward(const tensor::Image &image, int slice_width, int slice_height, float overlap_width_ratio, float overlap_height_ratio, void *stream = nullptr);
    
        virtual Result forward(const tensor::Image &image, void *stream = nullptr);
    
        virtual Result forwards(void *stream = nullptr);
};

Infer *loadraw(const std::string &engine_file, ModelType model_type, const std::vector<std::string>& names, float confidence_threshold,
    float nms_threshold);

std::shared_ptr<Infer> load_yolo(const std::string &engine_file, ModelType model_type, const std::vector<std::string>& names, int gpu_id, float confidence_threshold, float nms_threshold);

} // namespace yolo

#endif // YOLO_HPP__