#include <ocr_cls.h>

namespace ANSCENTER 
{
    void Classifier::Run(std::vector<cv::Mat> img_list,
                         std::vector<int> &cls_labels,
                         std::vector<float> &cls_scores) {

        try {
            int img_num = img_list.size();
            std::vector<int> cls_image_shape = { 3, 48, 192 };
            for (int beg_img_no = 0; beg_img_no < img_num;
                beg_img_no += this->cls_batch_num_) {
                auto preprocess_start = std::chrono::steady_clock::now();
                int end_img_no = std::min(img_num, beg_img_no + this->cls_batch_num_);
                int batch_num = end_img_no - beg_img_no;
                // preprocess
                std::vector<cv::Mat> norm_img_batch;
                for (int ino = beg_img_no; ino < end_img_no; ino++) {
                    cv::Mat srcimg;
                    img_list[ino].copyTo(srcimg);
                    cv::Mat resize_img;
                    this->resize_op_.Run(srcimg, resize_img, this->use_tensorrt_,
                        cls_image_shape);

                    this->normalize_op_.Run(&resize_img, this->mean_, this->scale_,
                        this->is_scale_);
                    if (resize_img.cols < cls_image_shape[2]) {
                        cv::copyMakeBorder(resize_img, resize_img, 0, 0, 0,
                            cls_image_shape[2] - resize_img.cols,
                            cv::BORDER_CONSTANT, cv::Scalar(0, 0, 0));
                    }
                    norm_img_batch.push_back(resize_img);
                }
                std::vector<float> input(batch_num * cls_image_shape[0] *
                    cls_image_shape[1] * cls_image_shape[2],
                    0.0f);
                this->permute_op_.Run(norm_img_batch, input.data());
      

                // inference.
                auto input_names = this->predictor_->GetInputNames();
                auto input_t = this->predictor_->GetInputHandle(input_names[0]);
                input_t->Reshape({ batch_num, cls_image_shape[0], cls_image_shape[1],
                                  cls_image_shape[2] });
                auto inference_start = std::chrono::steady_clock::now();
                input_t->CopyFromCpu(input.data());
                this->predictor_->Run();

                std::vector<float> predict_batch;
                auto output_names = this->predictor_->GetOutputNames();
                auto output_t = this->predictor_->GetOutputHandle(output_names[0]);
                auto predict_shape = output_t->shape();

                int out_num = std::accumulate(predict_shape.begin(), predict_shape.end(), 1,
                    std::multiplies<int>());
                predict_batch.resize(out_num);

                output_t->CopyToCpu(predict_batch.data());
                // postprocess
                auto postprocess_start = std::chrono::steady_clock::now();
                for (int batch_idx = 0; batch_idx < predict_shape[0]; batch_idx++) {
                    int label = int(
                        ANSOCRUtility::argmax(&predict_batch[batch_idx * predict_shape[1]],
                            &predict_batch[(batch_idx + 1) * predict_shape[1]]));
                    float score = float(*std::max_element(
                        &predict_batch[batch_idx * predict_shape[1]],
                        &predict_batch[(batch_idx + 1) * predict_shape[1]]));
                    cls_labels[beg_img_no + batch_idx] = label;
                    cls_scores[beg_img_no + batch_idx] = score;
                }
            }
        }
        catch (std::exception& e) {
            std::cout << "Classifier::Run" << e.what() << std::endl;
        }
    }

    bool Classifier::LoadModel(const std::string &model_dir) {
        try {
            paddle_infer::Config config;
            config.SetModel(model_dir + "/inference.pdmodel",
                model_dir + "/inference.pdiparams");

            if (this->use_gpu_) {
                config.EnableUseGpu(this->gpu_mem_, this->gpu_id_);
                if (this->use_tensorrt_) {
                    auto precision = paddle_infer::Config::Precision::kFloat32;
                    if (this->precision_ == "fp16") {
                        precision = paddle_infer::Config::Precision::kHalf;
                    }
                    if (this->precision_ == "int8") {
                        precision = paddle_infer::Config::Precision::kInt8;
                    }
                    config.EnableTensorRtEngine(1 << 20, 10, 3, precision, false, false);
                    if (!ANSOCRUtility::PathExists("./trt_cls_shape.txt")) {
                        config.CollectShapeRangeInfo("./trt_cls_shape.txt");
                    }
                    else {
                        config.EnableTunedTensorRtDynamicShape("./trt_cls_shape.txt", true);
                    }
                }
            }
            else {
                config.DisableGpu();
                if (this->use_mkldnn_) {
                    config.EnableMKLDNN();
                }
                config.SetCpuMathLibraryNumThreads(this->cpu_math_library_num_threads_);
            }
            // false for zero copy tensor
            config.SwitchUseFeedFetchOps(false);
            // true for multiple input
            config.SwitchSpecifyInputNames(true);
            config.SwitchIrOptim(true);
            config.EnableMemoryOptim();
            config.DisableGlogInfo();
            this->predictor_ = paddle_infer::CreatePredictor(config);
            return true;
        }
        catch (std::exception& e) {
            std::cout << "Classifier::LoadModel" << e.what() << std::endl;
            return false;
        }
      
    }
} // namespace ANSCENTER
