#ifndef ANSOCR_H
#define ANSOCR_H
#define ANSOCR_API __declspec(dllexport)
#pragma once
#include "opencv2/core.hpp"
#include "opencv2/imgcodecs.hpp"
#include "opencv2/imgproc.hpp"
#include <iostream>
#include <vector>
#include "LabVIEWHeader/extcode.h"
#include "ANSOCRUtility.h"
#include <ocr_cls.h>
#include <ocr_det.h>
#include <ocr_rec.h>
#include <structure_table.h>
#include <structure_layout.h>
#include "OCRLite.h"
#include "fastdeploy/vision.h"
#ifdef WIN32
const char sep = '\\';
#else
const char sep = '/';
#endif

namespace ANSCENTER {
	class ANSOCR_API ANSOCRBase {
		protected:
			bool            _licenseValid{ false };
			std::string     _licenseKey;
			std::string     _modelFolder;
			std::string     _modelConfigFile;
			OCRModelConfig  _modelConfig;
			int				_engineMode; //0: Auto detect, 1 GPU, 2 CPU
			SPDLogger& _logger = SPDLogger::instance("AOCR", false);
			void            CheckLicense();

		public:
			virtual bool Initialize(std::string licenseKey, OCRModelConfig modelConfig, const std::string& modelZipFilePath, const std::string& modelZipPassword, int engineMode);
			virtual std::vector<ANSCENTER::OCRObject> RunInference(const cv::Mat& input) = 0;
			
	};

	class ANSOCR_API ANSOCR:public ANSOCRBase {
		public:
			virtual bool Initialize(std::string licenseKey, OCRModelConfig modelConfig, const std::string& modelZipFilePath, const std::string& modelZipPassword, int engineMode) override;
			std::vector<ANSCENTER::OCRObject> RunInference(const cv::Mat& input);
		private:
		/*	std::unique_ptr<DBDetector> detector_ = nullptr;
			std::unique_ptr<Classifier> classifier_ = nullptr;
			std::unique_ptr<CRNNRecognizer> recognizer_ = nullptr;
			std::unique_ptr<StructureTableRecognizer> table_model_ = nullptr;
			std::unique_ptr<StructureLayoutRecognizer> layout_model_ = nullptr;*/

			DBDetector detector_;
			Classifier classifier_;
			CRNNRecognizer recognizer_;
			StructureTableRecognizer table_model_ ;
			StructureLayoutRecognizer layout_model_;

		private:
			std::vector<std::vector<OCRPredictResult>> OCRInference(std::vector<cv::Mat> img_list);
			std::vector<StructurePredictResult>StructureInference(cv::Mat srcimg, bool ocr = false);
			std::vector<OCRPredictResult> OCRInference(cv::Mat img);
			void DetectionInference(cv::Mat img, std::vector<OCRPredictResult>& ocr_results);
			void RecognitionInference(std::vector<cv::Mat> img_list,std::vector<OCRPredictResult>& ocr_results);
			void CLSInference(std::vector<cv::Mat> img_list,std::vector<OCRPredictResult>& ocr_results);
			void LayoutInference(cv::Mat img,std::vector<StructurePredictResult>& structure_result);
			void TableInference(cv::Mat img, StructurePredictResult& structure_result);

			std::string RebuildTable(std::vector<std::string> structure_html_tags, std::vector<std::vector<int>> structure_boxes, std::vector<OCRPredictResult>& ocr_result);
			float Dis(std::vector<int>& box1, std::vector<int>& box2);
			static bool ComparisonDis(const std::vector<float>& dis1, const std::vector<float>& dis2);
	};
	
	class ANSOCR_API ANSONNXOCR :public ANSOCRBase {
	public:
		virtual bool Initialize(std::string licenseKey, OCRModelConfig modelConfig, const std::string& modelZipFilePath, const std::string& modelZipPassword, int engineMode) override;
		std::vector<ANSCENTER::OCRObject> RunInference(const cv::Mat& input);
	private:
		OcrLite ocrLite;
		std::string modelsDir = "../models";
		std::string modelDetPath, modelClsPath, modelRecPath, keysPath;
		std::string argImgPath = "../../test_imgs/1.jpg";
		std::string imgPath, imgName;
		int numThread = 1;
		int padding = 0;
		int maxSideLen = 1024;
		float boxScoreThresh = 0.5f;
		float boxThresh = 0.3f;
		float unClipRatio = 2.0f;
		bool doAngle = true;
		int flagDoAngle = 1;
		bool mostAngle = false;
		int flagMostAngle = 0;
		int opt;
		int optionIndex = 0;
	};


	class ANSOCR_API ANSFDOCR :public ANSOCRBase {
	public:
		virtual bool Initialize(std::string licenseKey, OCRModelConfig modelConfig, const std::string& modelZipFilePath, const std::string& modelZipPassword, int engineMode) override;
		std::vector<ANSCENTER::OCRObject> RunInference(const cv::Mat& input);
	private:
		void InitAndInfer(const std::string& det_model_dir,
			const std::string& cls_model_dir,
			const std::string& rec_model_dir,
			const std::string& rec_label_file,
			const std::string& image_file,
			const fastdeploy::RuntimeOption& option) {
			try {
				auto det_model_file = det_model_dir + sep + "inference.pdmodel";
				auto det_params_file = det_model_dir + sep + "inference.pdiparams";

				auto cls_model_file = cls_model_dir + sep + "inference.pdmodel";
				auto cls_params_file = cls_model_dir + sep + "inference.pdiparams";

				auto rec_model_file = rec_model_dir + sep + "inference.pdmodel";
				auto rec_params_file = rec_model_dir + sep + "inference.pdiparams";

				auto det_option = option;
				auto cls_option = option;
				auto rec_option = option;

				// The cls and rec model can inference a batch of images now.
				// User could initialize the inference batch size and set them after create
				// PP-OCR model.
				int cls_batch_size = 1;
				int rec_batch_size = 6;

				// If use TRT backend, the dynamic shape will be set as follow.
				// We recommend that users set the length and height of the detection model to
				// a multiple of 32.
				// We also recommend that users set the Trt input shape as follow.
				det_option.SetTrtInputShape("x", { 1, 3, 64, 64 }, { 1, 3, 640, 640 },
					{ 1, 3, 960, 960 });
				cls_option.SetTrtInputShape("x", { 1, 3, 48, 10 }, { cls_batch_size, 3, 48, 320 },
					{ cls_batch_size, 3, 48, 1024 });
				rec_option.SetTrtInputShape("x", { 1, 3, 48, 10 }, { rec_batch_size, 3, 48, 320 },
					{ rec_batch_size, 3, 48, 2304 });

				// Users could save TRT cache file to disk as follow.
				// det_option.SetTrtCacheFile(det_model_dir + sep + "det_trt_cache.trt");
				// cls_option.SetTrtCacheFile(cls_model_dir + sep + "cls_trt_cache.trt");
				// rec_option.SetTrtCacheFile(rec_model_dir + sep + "rec_trt_cache.trt");

				auto det_model = fastdeploy::vision::ocr::DBDetector(
					det_model_file, det_params_file, det_option);
				auto cls_model = fastdeploy::vision::ocr::Classifier(
					cls_model_file, cls_params_file, cls_option);
				auto rec_model = fastdeploy::vision::ocr::Recognizer(
					rec_model_file, rec_params_file, rec_label_file, rec_option);

				assert(det_model.Initialized());
				assert(cls_model.Initialized());
				assert(rec_model.Initialized());

				// Parameters settings for pre and post processing of Det/Cls/Rec Models.
				// All parameters are set to default values.
				det_model.GetPreprocessor().SetMaxSideLen(960);
				det_model.GetPostprocessor().SetDetDBThresh(0.3);
				det_model.GetPostprocessor().SetDetDBBoxThresh(0.6);
				det_model.GetPostprocessor().SetDetDBUnclipRatio(1.5);
				det_model.GetPostprocessor().SetDetDBScoreMode("slow");
				det_model.GetPostprocessor().SetUseDilation(0);
				cls_model.GetPostprocessor().SetClsThresh(0.9);

				// The classification model is optional, so the PP-OCR can also be connected
				// in series as follows
				// auto ppocr_v3 = fastdeploy::pipeline::PPOCRv3(&det_model, &rec_model);
				auto ppocr_v3 =
					fastdeploy::pipeline::PPOCRv3(&det_model, &cls_model, &rec_model);

				// Set inference batch size for cls model and rec model, the value could be -1
				// and 1 to positive infinity.
				// When inference batch size is set to -1, it means that the inference batch
				// size
				// of the cls and rec models will be the same as the number of boxes detected
				// by the det model.
				ppocr_v3.SetClsBatchSize(cls_batch_size);
				ppocr_v3.SetRecBatchSize(rec_batch_size);

				if (!ppocr_v3.Initialized()) {
					std::cerr << "Failed to initialize PP-OCR." << std::endl;
					return;
				}

				auto im = cv::imread(image_file);
				auto im_bak = im.clone();

				fastdeploy::vision::OCRResult result;
				if (!ppocr_v3.Predict(&im, &result)) {
					std::cerr << "Failed to predict." << std::endl;
					return;
				}

				std::cout << result.Str() << std::endl;

				auto vis_im = fastdeploy::vision::VisOcr(im_bak, result);
				cv::imwrite("C:\\Projects\\AICamera\\Documentation\\TestOCR\\vis_result.jpg", vis_im);
				//std::cout << "Visualized result saved in ./vis_result.jpg" << std::endl;

			}
			catch (std::exception& e) {
				this->_logger.LogFatal("ANSOCR::RunInference", e.what());
			}
		}
	};

}


extern "C" ANSOCR_API int		   CreateANSOCRHandle(ANSCENTER::ANSOCRBase * *Handle, const char* licenseKey, const char* modelFilePath, const char* modelFileZipPassword, int language, int useDetector, int useRecognizer, int useCLS, int useLayout, int useTable, int engineMode, int useTensorRT, int useMKLDNN);
extern "C" ANSOCR_API std::string  RunInference(ANSCENTER::ANSOCRBase * *Handle, unsigned char* jpeg_string, unsigned int bufferLength);
extern "C" ANSOCR_API std::string  RunInferenceBinary(ANSCENTER::ANSOCRBase * *Handle, unsigned char* jpeg_bytes, unsigned int width, unsigned int height);
extern "C" ANSOCR_API int          ReleaseANSOCRHandle(ANSCENTER::ANSOCRBase * *Handle);
extern "C" ANSOCR_API std::string  RunInferenceImagePath(ANSCENTER::ANSOCRBase * *Handle, const char* imageFilePath);
//// For LabVIEW API
extern "C" ANSOCR_API int          RunInference_LV(ANSCENTER::ANSOCRBase * *Handle, unsigned char* jpeg_string, unsigned int bufferLength, LStrHandle detectionResult);
extern "C" ANSOCR_API int		   RunInferenceImagePath_LV(ANSCENTER::ANSOCRBase * *Handle, const char* imageFilePath, LStrHandle detectionResult);
extern "C" ANSOCR_API int		   ANSOCRUnitTest(const char* modelFilePath, const char* imageFilePath, LStrHandle detectionResult);
extern "C" ANSOCR_API int		   ANSONNXOCRUnitTest(const char* modelFilePath, const char* imageFilePath);
extern "C" ANSOCR_API int		   ANSFDOCRUnitTest();

#endif
