#pragma once
#include "paddle_api.h"
#include "paddle_inference_api.h"
#include <preprocess_op.h>
#include <ANSOCRUtility.h>
namespace ANSCENTER {
    class Classifier {
    public:
      //explicit
        Classifier() {}
      bool Initialize(const std::string &model_dir, const bool &use_gpu,
                          const int &gpu_id, const int &gpu_mem,
                          const int &cpu_math_library_num_threads,
                          const bool &use_mkldnn, const double &cls_thresh,
                          const bool &use_tensorrt, const std::string &precision,
                          const int &cls_batch_num) {
          try {
              this->use_gpu_ = use_gpu;
              this->gpu_id_ = gpu_id;
              this->gpu_mem_ = gpu_mem;
              this->cpu_math_library_num_threads_ = cpu_math_library_num_threads;
              this->use_mkldnn_ = use_mkldnn;

              this->cls_thresh = cls_thresh;
              this->use_tensorrt_ = use_tensorrt;
              this->precision_ = precision;
              this->cls_batch_num_ = cls_batch_num;
              return LoadModel(model_dir);
          }
          catch (std::exception& e) {
              std::cout<<"Classifier::Classifier"<<e.what()<<std::endl;
              return false;
          }

      }
      double cls_thresh = 0.9;
      // Load Paddle inference model
      bool LoadModel(const std::string &model_dir);
      void Run(std::vector<cv::Mat> img_list, std::vector<int> &cls_labels,
               std::vector<float> &cls_scores);
    private:
      std::shared_ptr<paddle_infer::Predictor> predictor_;
      bool use_gpu_ = false;
      int gpu_id_ = 0;
      int gpu_mem_ = 4000;
      int cpu_math_library_num_threads_ = 4;
      bool use_mkldnn_ = false;

      std::vector<float> mean_ = {0.5f, 0.5f, 0.5f};
      std::vector<float> scale_ = {1 / 0.5f, 1 / 0.5f, 1 / 0.5f};
      bool is_scale_ = true;
      bool use_tensorrt_ = false;
      std::string precision_ = "fp32";
      int cls_batch_num_ = 1;
      // pre-process
      ClsResizeImg resize_op_;
      Normalize normalize_op_;
      PermuteBatch permute_op_;
    }; // class Classifier
} // namespace ANSCENTER
