#include "ANSOPENVINOOD.h"
#include "Utility.h"
namespace ANSCENTER
{
	bool OPENVINOOD::OptimizeModel(bool fp16, std::string& optimizedModelFolder) {
		std::lock_guard<std::recursive_mutex> lock(_mutex);
		if (!ANSODBase::OptimizeModel(fp16, optimizedModelFolder)) {
			return false;
		}
		if (FileExist(_modelFilePath)) {
			std::string modelName = GetFileNameWithoutExtension(_modelFilePath);
			std::string binaryModelName = modelName + ".bin";
			std::string modelFolder = GetParentFolder(_modelFilePath);
			std::string optimizedModelPath = CreateFilePath(modelFolder, binaryModelName);
			if (FileExist(optimizedModelPath)) {
				this->_logger.LogDebug("OPENVINOOD::OptimizeModel", "This model is optimized. No need other optimization.", __FILE__, __LINE__);
				optimizedModelFolder = modelFolder;
				return true;
			}
			else {
				this->_logger.LogFatal("OPENVINOOD::OptimizeModel", "This model can not be optimized.", __FILE__, __LINE__);
				optimizedModelFolder = modelFolder;
				return false;
			}
		}
		else {
			this->_logger.LogFatal("OPENVINOOD::OptimizeModel", "This model is not exist. Please check the model path again.", __FILE__, __LINE__);
			optimizedModelFolder = "";
			return false;
		}
	}
	bool OPENVINOOD::LoadModel(const std::string& modelZipFilePath, const std::string& modelZipPassword) {
		std::lock_guard<std::recursive_mutex> lock(_mutex);
		try {
			bool result = ANSODBase::LoadModel(modelZipFilePath, modelZipPassword);
			if (!result) return false;
			// 0. Check if the configuration file exist
			if (FileExist(_modelConfigFile)) {
				ModelType modelType;
				std::vector<int> inputShape;
				_classes = ANSUtilityHelper::GetConfigFileContent(_modelConfigFile, modelType, inputShape);
				if (inputShape.size() == 2) {
					if (inputShape[0] > 0)_modelConfig.inpHeight = inputShape[0];
					if (inputShape[1] > 0)_modelConfig.inpWidth = inputShape[1];
				}
			}
			else {// This is old version of model zip file
				std::string onnxfile = CreateFilePath(_modelFolder, "train_last.xml");//yolov8n.xml
				if (std::filesystem::exists(onnxfile)) {
					_modelFilePath = onnxfile;
					_classFilePath = CreateFilePath(_modelFolder, "classes.names");
					this->_logger.LogDebug("OPENVINOOD::Initialize.  Loading OpenVINO weight", _modelFilePath, __FILE__, __LINE__);
				}
				else {
					this->_logger.LogError("OPENVINOOD::Initialize.  Model file is not exist", _modelFilePath, __FILE__, __LINE__);
					return false;
				}
				std::ifstream isValidFileName(_classFilePath);
				if (!isValidFileName)
				{
					this->_logger.LogDebug("OPENVINOOD::Initialize.  Load classes from string", _classFilePath, __FILE__, __LINE__);
					LoadClassesFromString();
				}
				else {
					this->_logger.LogDebug("OPENVINOOD::Initialize.  Load classes from file", _classFilePath, __FILE__, __LINE__);
					LoadClassesFromFile();
				}
			}

			// Load Model from Here
			InitialModel();
			_isInitialized = true;
			return true;
		}
		catch (std::exception& e) {
			this->_logger.LogFatal("OPENVINOOD::LoadModel", e.what(), __FILE__, __LINE__);
			return false;
		}
	}
	bool OPENVINOOD::LoadModelFromFolder(std::string licenseKey, ModelConfig modelConfig, std::string modelName, std::string className, const std::string& modelFolder, std::string& labelMap) {
		std::lock_guard<std::recursive_mutex> lock(_mutex);
		try {
			bool result = ANSODBase::LoadModelFromFolder(licenseKey, modelConfig, modelName, className, modelFolder, labelMap);
			if (!result) return false;
			std::string _modelName = modelName;
			if (_modelName.empty()) {
				_modelName = "train_last";
			}
			std::string modelFullName = _modelName + ".xml";
			_modelConfig = modelConfig;
			_modelConfig.detectionType = ANSCENTER::DetectionType::DETECTION;
			_modelConfig.modelType = ModelType::OPENVINO;
			_modelConfig.inpHeight = 640;
			_modelConfig.inpWidth = 640;

			if (_modelConfig.modelMNSThreshold < 0.2)
				_modelConfig.modelMNSThreshold = 0.5;
			if (_modelConfig.modelConfThreshold < 0.2)
				_modelConfig.modelConfThreshold = 0.5;
			// 0. Check if the configuration file exist
			if (FileExist(_modelConfigFile)) {
				ModelType modelType;
				std::vector<int> inputShape;
				_classes = ANSUtilityHelper::GetConfigFileContent(_modelConfigFile, modelType, inputShape);
				if (inputShape.size() == 2) {
					if (inputShape[0] > 0)_modelConfig.inpHeight = inputShape[0];
					if (inputShape[1] > 0)_modelConfig.inpWidth = inputShape[1];
				}
			}
			else {// This is old version of model zip file
				std::string onnxfile = CreateFilePath(_modelFolder, modelFullName);//yolov8n.xml
				if (std::filesystem::exists(onnxfile)) {
					_modelFilePath = onnxfile;
					_classFilePath = CreateFilePath(_modelFolder, className);
					this->_logger.LogDebug("OPENVINOOD::Initialize.  Loading OpenVINO weight", _modelFilePath, __FILE__, __LINE__);
				}
				else {
					this->_logger.LogError("OPENVINOOD::Initialize.  Model file is not exist", _modelFilePath, __FILE__, __LINE__);
					return false;
				}
				std::ifstream isValidFileName(_classFilePath);
				if (!isValidFileName)
				{
					this->_logger.LogDebug("OPENVINOOD::Initialize.  Load classes from string", _classFilePath, __FILE__, __LINE__);
					LoadClassesFromString();
				}
				else {
					this->_logger.LogDebug("OPENVINOOD::Initialize.  Load classes from file", _classFilePath, __FILE__, __LINE__);
					LoadClassesFromFile();
				}
			}

			// 1. Load labelMap and engine
			labelMap.clear();
			if (!_classes.empty())
				labelMap = VectorToCommaSeparatedString(_classes);

			// Load Model from Here
			InitialModel();
			_isInitialized = true;
			return true;
		}
		catch (std::exception& e) {
			this->_logger.LogFatal("OPENVINOOD::LoadModel", e.what(), __FILE__, __LINE__);
			return false;
		}
	}
	cv::Mat OPENVINOOD::PreProcessing(const cv::Mat& source) {
		try {
			if (source.empty()) {
				this->_logger.LogFatal("OPENVINOOD::PreProcessing", "Empty image provided", __FILE__, __LINE__);
				return cv::Mat();
			}

			// Convert grayscale to 3-channel BGR if needed
			cv::Mat processedImage;
			if (source.channels() == 1) {
				cv::cvtColor(source, processedImage, cv::COLOR_GRAY2BGR);
			}
			else {
				processedImage = source;
			}

			int col = processedImage.cols;
			int row = processedImage.rows;
			int maxSize = std::max(col, row);

			// Create a square padded image with a black background
			cv::Mat result = cv::Mat::zeros(maxSize, maxSize, CV_8UC3);

			// Copy the original image to the top-left corner of the square matrix
			processedImage.copyTo(result(cv::Rect(0, 0, col, row)));

			return result;
		}
		catch (const std::exception& e) {
			this->_logger.LogFatal("OPENVINOOD::PreProcessing", e.what(), __FILE__, __LINE__);
			return cv::Mat();
		}
	}
	bool OPENVINOOD::Initialize(std::string licenseKey, ModelConfig modelConfig, const std::string& modelZipFilePath, const std::string& modelZipPassword, std::string& labelMap) {
		std::lock_guard<std::recursive_mutex> lock(_mutex);
		try {
			bool result = ANSODBase::Initialize(licenseKey, modelConfig, modelZipFilePath, modelZipPassword, labelMap);
			if (!result) return false;
			// Parsing for YOLO only here
			_modelConfig = modelConfig;
			_modelConfig.detectionType = ANSCENTER::DetectionType::DETECTION;
			_modelConfig.modelType = ModelType::OPENVINO;
			_modelConfig.inpHeight = 640;
			_modelConfig.inpWidth = 640;
			if (_modelConfig.modelMNSThreshold < 0.2)
				_modelConfig.modelMNSThreshold = 0.5;
			if (_modelConfig.modelConfThreshold < 0.2)
				_modelConfig.modelConfThreshold = 0.5;
			model_input_shape_ = cv::Size2f(_modelConfig.inpWidth, _modelConfig.inpHeight);
			// 0. Check if the configuration file exist
			if (FileExist(_modelConfigFile)) {
				ModelType modelType;
				std::vector<int> inputShape;
				_classes = ANSUtilityHelper::GetConfigFileContent(_modelConfigFile, modelType, inputShape);
				if (inputShape.size() == 2) {
					if (inputShape[0] > 0)_modelConfig.inpHeight = inputShape[0];
					if (inputShape[1] > 0)_modelConfig.inpWidth = inputShape[1];
				}
			}
			else {// This is old version of model zip file
				std::string onnxfile = CreateFilePath(_modelFolder, "train_last.xml");//yolov8n.xml
				if (std::filesystem::exists(onnxfile)) {
					_modelFilePath = onnxfile;
					_classFilePath = CreateFilePath(_modelFolder, "classes.names");
					this->_logger.LogDebug("OPENVINOOD::Initialize.  Loading OpenVINO weight", _modelFilePath, __FILE__, __LINE__);
				}
				else {
					this->_logger.LogError("OPENVINOOD::Initialize.  Model file is not exist", _modelFilePath, __FILE__, __LINE__);
					return false;
				}
				std::ifstream isValidFileName(_classFilePath);
				if (!isValidFileName)
				{
					this->_logger.LogDebug("OPENVINOOD::Initialize.  Load classes from string", _classFilePath, __FILE__, __LINE__);
					LoadClassesFromString();
				}
				else {
					this->_logger.LogDebug("OPENVINOOD::Initialize.  Load classes from file", _classFilePath, __FILE__, __LINE__);
					LoadClassesFromFile();
				}
			}

			// 1. Load labelMap and engine
			labelMap.clear();
			if (!_classes.empty())
				labelMap = VectorToCommaSeparatedString(_classes);

			// Load Model from Here
			InitialModel();
			_isInitialized = true;
			return true;

		}
		catch (std::exception& e) {
			this->_logger.LogFatal("OPENVINOOD::Initialize", e.what(), __FILE__, __LINE__);
			return false;
		}
	}
	std::vector<Object> OPENVINOOD::RunInference(const cv::Mat& input) {
		return RunInference(input, "CustomCam");
	}
	std::vector<Object> OPENVINOOD::RunInference(const cv::Mat& input, const std::string& camera_id)
	{
		std::lock_guard<std::recursive_mutex> lock(_mutex);

		// Early validation
		if (!_licenseValid) {
			_logger.LogError("OPENVINOOD::RunInference", "Invalid License",
				__FILE__, __LINE__);
			return {};
		}

		if (!_isInitialized) {
			_logger.LogError("OPENVINOOD::RunInference", "Model is not initialized",
				__FILE__, __LINE__);
			return {};
		}

		if (input.empty() || input.cols < 10 || input.rows < 10) {
			_logger.LogError("OPENVINOOD::RunInference", "Invalid input image",
				__FILE__, __LINE__);
			return {};
		}

		try {
			// Step 1: Preprocessing
			cv::Mat letterbox_img = PreProcessing(input);
			if (letterbox_img.empty()) {
				_logger.LogError("OPENVINOOD::RunInference", "PreProcessing failed",
					__FILE__, __LINE__);
				return {};
			}

			// Step 2: Create blob
			constexpr int imageSize = 640;
			const float scale = static_cast<float>(letterbox_img.rows) / imageSize;

			cv::Mat blob = cv::dnn::blobFromImage(letterbox_img, 1.0 / 255.0,
				cv::Size(imageSize, imageSize),
				cv::Scalar(), true);

			// Step 3: Run inference
			auto input_port = compiled_model_.input();
			ov::Tensor input_tensor(input_port.get_element_type(),
				input_port.get_shape(), blob.ptr(0));
			inference_request_.set_input_tensor(input_tensor);
			inference_request_.infer();

			// Step 4: Get output
			auto output = inference_request_.get_output_tensor(0);
			auto output_shape = output.get_shape();

			if (output_shape.size() != 3) {
				_logger.LogError("OPENVINOOD::RunInference", "Unexpected output shape",
					__FILE__, __LINE__);
				return {};
			}

			const int rows = static_cast<int>(output_shape[2]);        // 8400
			const int dimensions = static_cast<int>(output_shape[1]);  // 84
			float* data = output.data<float>();

			// Step 5: Parse detections (avoid transpose!)
			std::vector<int> class_ids;
			std::vector<float> class_scores;
			std::vector<cv::Rect> boxes;

			// Pre-allocate for efficiency
			class_ids.reserve(rows / 10);      // Estimate ~10% pass threshold
			class_scores.reserve(rows / 10);
			boxes.reserve(rows / 10);

			const cv::Rect imageBounds(0, 0, input.cols, input.rows);

			// Process detections without transpose
			for (int i = 0; i < rows; i++) {
				// Get class scores starting at index 4
				float max_score = -1.0f;
				int max_class_id = 0;

				for (int j = 4; j < dimensions; j++) {
					const float score = data[j * rows + i];  // Column-major access
					if (score > max_score) {
						max_score = score;
						max_class_id = j - 4;
					}
				}

				if (max_score > _modelConfig.detectionScoreThreshold) {
					// Extract box coordinates
					const float cx = data[0 * rows + i];
					const float cy = data[1 * rows + i];
					const float w = data[2 * rows + i];
					const float h = data[3 * rows + i];

					// Convert to pixel coordinates and clamp
					cv::Rect box(
						static_cast<int>((cx - 0.5f * w) * scale),
						static_cast<int>((cy - 0.5f * h) * scale),
						static_cast<int>(w * scale),
						static_cast<int>(h * scale)
					);

					box &= imageBounds;  // Clamp to image bounds

					if (box.area() > 0) {
						class_ids.push_back(max_class_id);
						class_scores.push_back(max_score);
						boxes.push_back(box);
					}
				}
			}

			// Step 6: Apply NMS
			if (boxes.empty()) {
				return {};
			}

			std::vector<int> indices;
			cv::dnn::NMSBoxes(boxes, class_scores,
				_modelConfig.detectionScoreThreshold,
				_modelConfig.modelMNSThreshold, indices);

			// Step 7: Build output objects
			std::vector<Object> outputs;
			outputs.reserve(indices.size());

			const int classNameSize = static_cast<int>(_classes.size());

			for (int idx : indices) {
				Object result;
				result.classId = class_ids[idx];
				result.confidence = class_scores[idx];
				result.box = boxes[idx];
				result.cameraId = camera_id;

				// Set class name
				if (!_classes.empty()) {
					result.className = (result.classId < classNameSize)
						? _classes[result.classId]
						: _classes.back();
				}
				else {
					result.className = "Unknown";
				}

				// Set polygon
				result.polygon = ANSUtilityHelper::RectToNormalizedPolygon(
					result.box, input.cols, input.rows
				);

				outputs.push_back(std::move(result));
			}

			if (_trackerEnabled) {
				outputs = ApplyTracking(outputs, camera_id);
				if (_stabilizationEnabled) outputs = StabilizeDetections(outputs, camera_id);
			}
			return outputs;

		}
		catch (const std::exception& e) {
			_logger.LogFatal("OPENVINOOD::RunInference", e.what(), __FILE__, __LINE__);
			return {};
		}
		catch (...) {
			_logger.LogFatal("OPENVINOOD::RunInference", "Unknown error",
				__FILE__, __LINE__);
			return {};
		}
	}
	OPENVINOOD::~OPENVINOOD() {
		try {
			if (FolderExist(_modelFolder)) {
				if (!DeleteFolder(_modelFolder)) {
					this->_logger.LogError("OPENVINOOD::~OPENVINOOD", "Failed to delete OpenVINO Models", __FILE__, __LINE__);
				}
			}
		}
		catch (std::exception& e) {
			this->_logger.LogError("OPENVINOOD::~OPENVINOOD()", "Failed to release OPENVINO Models", __FILE__, __LINE__);
		}
	}
	bool OPENVINOOD::Destroy() {
		try {
			if (FolderExist(_modelFolder)) {
				DeleteFolder(_modelFolder);
			}
			return true;
		}
		catch (std::exception& e) {
			this->_logger.LogError("OPENVINOOD::Destroy()", "Failed to release OPENVINO Models", __FILE__, __LINE__);
			return false;
		}
	}
	//private
	void OPENVINOOD::InitialModel() {
		try {
			// Step 1: Initialize OpenVINO Runtime Core
			ov::Core core;
			// Step 2: Get Available Devices and Log
			std::vector<std::string> available_devices = core.get_available_devices();
			// Define device priority: NPU > GPU > CPU
			std::vector<std::string> priority_devices = { "NPU", "GPU" };
			bool device_found = false;

			// Iterate over prioritized devices
			for (const auto& device : priority_devices) {
				if (std::find(available_devices.begin(), available_devices.end(), device) != available_devices.end()) {
					if (device == "NPU") {
						core.set_property("NPU", ov::hint::performance_mode(ov::hint::PerformanceMode::CUMULATIVE_THROUGHPUT));
						core.set_property("GPU", ov::hint::performance_mode(ov::hint::PerformanceMode::CUMULATIVE_THROUGHPUT));
						compiled_model_ = core.compile_model(_modelFilePath, "MULTI:NPU, GPU");
					}
					else {
						// Configure and compile for individual device
						//core.set_property(device, ov::hint::performance_mode(ov::hint::PerformanceMode::THROUGHPUT));
						compiled_model_ = core.compile_model(_modelFilePath, device);
					}
					device_found = true;
					break;
				}
			}
			// Fallback: Default to CPU if no devices found
			if (!device_found) {
				//core.set_property("CPU", ov::hint::performance_mode(ov::hint::PerformanceMode::THROUGHPUT));
				compiled_model_ = core.compile_model(_modelFilePath, "CPU");
			}
			// Step 3: Create Inference Request
			inference_request_ = compiled_model_.create_infer_request();
		}
		catch (const std::exception& e) {
			// Log any errors
			this->_logger.LogFatal("OPENVINOOD::InitialModel", e.what(), __FILE__, __LINE__);
		}
	}
}

//std::vector<Object> OPENVINOOD::RunInference(const cv::Mat& input, const std::string& camera_id) {
	//	std::lock_guard<std::recursive_mutex> lock(_mutex);
	//	std::vector<Object> outputs;
	//	if (!_licenseValid) {
	//		_logger.LogError("OPENVINOOD::RunInference", "Invalid License", __FILE__, __LINE__);
	//		return outputs;
	//	}
	//	if (!_isInitialized) {
	//		_logger.LogError("OPENVINOOD::RunInference", "Model is not initialized", __FILE__, __LINE__);
	//		return outputs;
	//	}
	//	if (input.empty()) {
	//		_logger.LogError("OPENVINOOD::RunInference", "Input image is empty", __FILE__, __LINE__);
	//		return outputs;
	//	}
	//	try {
	//		// Step 0: Prepare input
	//		if (input.empty()) return outputs;
	//		if ((input.cols < 10) || (input.rows < 10)) return outputs;
	//		cv::Mat letterbox_img = PreProcessing(input);
	//		if (letterbox_img.empty()) {
	//			_logger.LogError("OPENVINOOD::RunInference", "PreProcessing failed", __FILE__, __LINE__);
	//			return outputs;
	//		}
	//		int imageSize = 640;
	//		int maxImageSize = std::max(letterbox_img.cols, letterbox_img.rows);
	//		//if (maxImageSize < imageSize)imageSize = maxImageSize;
	//		float scale = static_cast<float>(letterbox_img.rows) / imageSize;
	//		cv::Mat blob = cv::dnn::blobFromImage(letterbox_img, 1.0 / 255.0, 
	//											  cv::Size(imageSize, imageSize),
	//											  cv::Scalar(), true);

	//		// Step 1: Feed blob to the network
	//		auto input_port = compiled_model_.input();
	//		ov::Tensor input_tensor(input_port.get_element_type(), input_port.get_shape(), blob.ptr(0));
	//		inference_request_.set_input_tensor(input_tensor);
	//		inference_request_.infer();

	//		// Step 4: Get output
	//		auto output = inference_request_.get_output_tensor(0);
	//		auto output_shape = output.get_shape();

	//		if (output_shape.size() != 3) {
	//			_logger.LogError("OPENVINOOD::RunInference", "Unexpected output shape", __FILE__, __LINE__);
	//			return outputs;
	//		}

	//		int rows = output_shape[2];        // 8400
	//		int dimensions = output_shape[1]; // 84: box[cx, cy, w, h]+80 class scores
	//		float* data = output.data<float>();
	//		cv::Mat output_buffer(dimensions, rows, CV_32F, data);
	//		transpose(output_buffer, output_buffer); // [8400, 84]

	//		std::vector<int> class_ids;
	//		std::vector<float> class_scores;
	//		std::vector<cv::Rect> boxes;

	//		// Step 5: Post-processing
	//		for (int i = 0; i < output_buffer.rows; i++) {
	//			cv::Mat classes_scores = output_buffer.row(i).colRange(4, dimensions);
	//			cv::Point class_id;
	//			double max_class_score;
	//			minMaxLoc(classes_scores, nullptr, &max_class_score, nullptr, &class_id);

	//			if (max_class_score > _modelConfig.detectionScoreThreshold) {
	//				class_scores.push_back(static_cast<float>(max_class_score));
	//				class_ids.push_back(class_id.x);

	//				float cx = output_buffer.at<float>(i, 0);
	//				float cy = output_buffer.at<float>(i, 1);
	//				float w = output_buffer.at<float>(i, 2);
	//				float h = output_buffer.at<float>(i, 3);

	//				int left = static_cast<int>((cx - 0.5f * w) * scale);
	//				int top = static_cast<int>((cy - 0.5f * h) * scale);
	//				int width = static_cast<int>(w * scale);
	//				int height = static_cast<int>(h * scale);

	//				left = std::max(0, left);
	//				top = std::max(0, top);
	//				width = std::min(input.cols - left, width);
	//				height = std::min(input.rows - top, height);

	//				boxes.emplace_back(left, top, width, height);
	//			}
	//		}
	//		// NMS
	//		int classNameSize = static_cast<int>(_classes.size());
	//		std::vector<int> indices;
	//		cv::dnn::NMSBoxes(boxes, class_scores, _modelConfig.detectionScoreThreshold, _modelConfig.modelMNSThreshold, indices);
	//		for (int id : indices) {
	//			if (class_scores[id] >= _modelConfig.detectionScoreThreshold) {
	//				Object result;
	//				int class_id = class_ids[id];
	//				result.classId = class_id;
	//				if (!_classes.empty()) {
	//					if (result.classId < classNameSize) {
	//						result.className = _classes[result.classId];
	//					}
	//					else {
	//						result.className = _classes[classNameSize - 1]; // Use last valid class name if out of range
	//					}
	//				}
	//				else {
	//					result.className = "Unknown"; // Fallback if _classes is empty
	//				}

	//				result.confidence = class_scores[id];
	//				result.box = boxes[id];
	//				result.polygon = ANSUtilityHelper::RectToNormalizedPolygon(result.box, input.cols, input.rows);
	//				result.cameraId = camera_id;
	//				outputs.push_back(result);
	//			}
	//		}
	//	}
	//	catch (const std::exception& e) {
	//		_logger.LogFatal("OPENVINOOD::RunInference", e.what(), __FILE__, __LINE__);
	//	}
	//	catch (...) {
	//		_logger.LogFatal("OPENVINOOD::RunInference", "Unknown error occurred", __FILE__, __LINE__);
	//	}
	//	return outputs; // 
	//}