#include "include/ocr_rec.h"

using namespace std;

namespace PaddleOCR {
	Recognizer::Recognizer(string model_path, const string& label_path) {
		ov::Core core;
		this->model_path = model_path;
		this->model = core.read_model(this->model_path);
		// reshape the model for dynamic batch size and sentence width
		this->model->reshape({ {ov::Dimension(1, 6), this->rec_image_shape_[0], this->rec_image_shape_[1], -1} });
		//core.set_property("CPU", ov::hint::performance_mode(ov::hint::PerformanceMode::THROUGHPUT));
		this->compiled_model = core.compile_model(this->model, "CPU");
		//this->compiled_model = core.compile_model(this->model, "CPU");
		this->infer_request = this->compiled_model.create_infer_request();
		this->label_list_ = Utility::ReadDict(label_path);
		this->label_list_.insert(this->label_list_.begin(),
			"#"); // blank char for ctc
		this->label_list_.push_back(" ");
	}
	void Recognizer::SetParameters(int rec_batch_num) {
		std::lock_guard<std::recursive_mutex> lock(_mutex);
		this->rec_batch_num_ = rec_batch_num;

	}
	void Recognizer::GetParameters(int& rec_batch_num) {
		std::lock_guard<std::recursive_mutex> lock(_mutex);
		rec_batch_num = this->rec_batch_num_;

	}
	void Recognizer::Run(const std::vector<cv::Mat> &img_list, std::vector<OCRPredictResult>& ocr_results) {
		std::lock_guard<std::recursive_mutex> lock(_mutex);
		try {
			std::vector<std::string> rec_texts(img_list.size(), "");
			std::vector<float> rec_text_scores(img_list.size(), 0);
			int img_num = img_list.size();
			std::vector<float> width_list;
			for (int i = 0; i < img_num; i++) {
				width_list.push_back(float(img_list[i].cols) / img_list[i].rows);
			}
			std::vector<int> indices = Utility::argsort(width_list);

			for (int beg_img_no = 0; beg_img_no < img_num;
				beg_img_no += this->rec_batch_num_) {
				int end_img_no = std::min(img_num, beg_img_no + this->rec_batch_num_);
				size_t batch_num = end_img_no - beg_img_no;
				size_t imgH = this->rec_image_shape_[1];
				size_t imgW = this->rec_image_shape_[2];
				float max_wh_ratio = imgW * 1.0 / imgH;
				for (int ino = beg_img_no; ino < end_img_no; ino++) {
					int h = img_list[indices[ino]].rows;
					int w = img_list[indices[ino]].cols;
					float wh_ratio = w * 1.0 / h;
					max_wh_ratio = std::max(max_wh_ratio, wh_ratio);
				}

				int batch_width = imgW;
				std::vector<cv::Mat> norm_img_batch;
				for (int ino = beg_img_no; ino < end_img_no; ino++) {
					cv::Mat srcimg;
					img_list[indices[ino]].copyTo(srcimg);
					cv::Mat resize_img;
					// preprocess
					this->resize_op_.Run(srcimg, resize_img, max_wh_ratio, this->rec_image_shape_);
					this->normalize_op_.Run(&resize_img, this->mean_, this->scale_,
						this->is_scale_);
					norm_img_batch.push_back(resize_img);
					batch_width = std::max(resize_img.cols, batch_width);
				}
				// prepare input tensor
				std::vector<float> input(batch_num * 3 * imgH * batch_width, 0.0f);
				ov::Shape intput_shape = { batch_num, 3, imgH, (size_t)batch_width };
				this->permute_op_.Run(norm_img_batch, input.data());
				auto input_port = this->compiled_model.input();
				ov::Tensor input_tensor(input_port.get_element_type(), intput_shape, input.data());
				this->infer_request.set_input_tensor(input_tensor);
				// start inference
	/*			this->infer_request.start_async();
				this->infer_request.wait();*/
				this->infer_request.infer();

				auto output = this->infer_request.get_output_tensor();
				const float* out_data = output.data<const float>();
				auto predict_shape = output.get_shape();

				// predict_batch is the result of Last FC with softmax
				for (int m = 0; m < predict_shape[0]; m++) {
					std::string str_res;
					int argmax_idx;
					int last_index = 0;
					float score = 0.f;
					int count = 0;
					float max_value = 0.0f;

					for (int n = 0; n < predict_shape[1]; n++) {
						// get idx
						argmax_idx = int(Utility::argmax(
							&out_data[(m * predict_shape[1] + n) * predict_shape[2]],
							&out_data[(m * predict_shape[1] + n + 1) * predict_shape[2]]));
						// get score
						max_value = float(*std::max_element(
							&out_data[(m * predict_shape[1] + n) * predict_shape[2]],
							&out_data[(m * predict_shape[1] + n + 1) * predict_shape[2]]));

						if (argmax_idx > 0 && (!(n > 0 && argmax_idx == last_index))) {
							score += max_value;
							count += 1;
							str_res += this->label_list_[argmax_idx];
						}
						last_index = argmax_idx;
					}
					score /= count;
					if (std::isnan(score)) {
						continue;
					}
					rec_texts[indices[beg_img_no + m]] = str_res;
					rec_text_scores[indices[beg_img_no + m]] = score;
				}
			}
			// sort boex from top to bottom, from left to right
			for (int i = 0; i < rec_texts.size(); i++) {
				ocr_results[i].text = rec_texts[i];
				ocr_results[i].score = rec_text_scores[i];
			}
		}
		catch (const std::exception& e) {
			std::cerr << e.what() << std::endl;
		}
	}
}