#include <include/ocr_export.h>

using namespace std;
using namespace cv;
using namespace PaddleOCR;
using namespace assistant;

char* TextOCR::PaddleOCRText(cv::Mat& image)
{
	DBDetector det(_det_model_dir, _use_gpu,
		_gpu_id, _gpu_mem, _cpu_threads,
		_enable_mkldnn, _max_side_len, _det_db_thresh,
		_det_db_box_thresh, _det_db_unclip_ratio,
		_use_polygon_score, _visualize,
		_use_tensorrt, _precision);

	Classifier* cls = nullptr;
	if (_use_angle_cls)
	{
		cls = new Classifier(_cls_model_dir, _use_gpu,
			_gpu_id, _gpu_mem, _cpu_threads, _enable_mkldnn,
			_cls_thresh, _use_tensorrt, _precision);
	}

	std::string char_list_file = _pchar_list_file;

	CRNNRecognizer rec(_rec_model_dir, _use_gpu,
		_gpu_id, _gpu_mem, _cpu_threads, _enable_mkldnn,
		char_list_file, _use_tensorrt, _precision,
		_rec_batch_num);

	if (!image.data)
	{
		std::cout << "[ERROR] cannot open the image" << std::endl;
		exit(0);
	}

	std::vector<std::vector<std::vector<int>>> boxes;
	std::vector<double> det_times;
	std::vector<double> rec_times;

	det.Run(image, boxes, &det_times);

	std::vector<cv::Mat> img_list;
	for (int j = 0; j < boxes.size(); j++) {
		cv::Mat crop_img;
		crop_img = Utility::GetRotateCropImage(image, boxes[j]);
		if (cls != nullptr) {
			crop_img = cls->Run(crop_img);
		}
		img_list.push_back(crop_img);
	}

	std::vector<std::vector<std::string>> str_res;
	std::string tmpstr;
	str_res = rec.GetOCR(img_list);

	char* reschar;
	for (int i = 0; i < str_res.size(); i++)
	{
		for (auto item : str_res[i]) {
			tmpstr += item;
		}
		tmpstr += "\n";
	}
	reschar = new char[tmpstr.length() + 1];
	tmpstr.copy(reschar, std::string::npos);

	return reschar;
}

void TextOCR::InitTextOCR(const char* save_log_path, bool use_gpu,
	int gpu_id, int gpu_mem, int cpu_threads,
	bool enable_mkldnn, bool use_tensorrt, const char* precision,
	bool benchmark, const char* det_model_dir, int max_side_len,
	double det_db_thresh, double det_db_box_thresh, double det_db_unclip_ratio,
	bool use_polygon_score, bool visualize, bool use_angle_cls,
	const char* cls_model_dir, double cls_thresh, const char* rec_model_dir,
	int rec_batch_num, const char* pchar_list_file)
{
	_save_log_path = save_log_path;
	_use_gpu = use_gpu;
	_gpu_id = gpu_id;
	_gpu_mem = gpu_mem;
	_cpu_threads = cpu_threads;
	_enable_mkldnn = enable_mkldnn;
	_use_tensorrt = use_tensorrt;
	_precision = precision;
	_benchmark = benchmark;
	_det_model_dir = det_model_dir;
	_max_side_len = max_side_len;
	_det_db_thresh = det_db_thresh;
	_det_db_box_thresh = det_db_box_thresh;
	_det_db_unclip_ratio = det_db_unclip_ratio;
	_use_polygon_score = use_polygon_score;
	_visualize = visualize;
	_use_angle_cls = use_angle_cls;
	_cls_model_dir = cls_model_dir;
	_cls_thresh = cls_thresh;
	_rec_model_dir = rec_model_dir;
	_rec_batch_num = rec_batch_num;
	_pchar_list_file = pchar_list_file;
}