#include "ANSCustomFallDetection.h"

ANSCustomFallDetector::ANSCustomFallDetector()
{
}
bool ANSCustomFallDetector::OptimizeModel(bool fp16)
{
	if (!_detector) return false;
	if (_detector->Optimize(fp16) != 0) return false;
	return true;
}
std::vector<CustomObject> ANSCustomFallDetector::RunInference(const cv::Mat& input)
{
	return RunInference(input, "CustomCam");
}
bool ANSCustomFallDetector::Destroy()
{
	_detector.reset();
	_motiondetector.reset();
	return true;
}
bool ANSCustomFallDetector::Initialize(const std::string& modelDirectory, float detectionScoreThreshold, std::string& labelMap)
{
	try {
		_modelDirectory = modelDirectory;
		_detectionScoreThreshold = detectionScoreThreshold;

		// Create detector and motion detector through ANSLIB factory
		_detector = ANSLIBPtr(ANSCENTER::ANSLIB::Create(), &ANSCENTER::ANSLIB::Destroy);
		_motiondetector = ANSLIBPtr(ANSCENTER::ANSLIB::Create(), &ANSCENTER::ANSLIB::Destroy);
		if (!_detector || !_motiondetector) {
			return false;
		}

		// Check hardware type
		engineType = _detector->GetEngineType();
		if (engineType == 0) {
			_detectorModelType = 4;     // TENSORRT
			_detectorDetectionType = 1; // DETECTION
		}
		else {
			_detectorModelType = 5;     // OPENVINO
			_detectorDetectionType = 1; // DETECTION
		}
		// Motion detector is always CPU-based
		_motionModelType = 5;       // OPENVINO
		_motionDetectionType = 1;   // DETECTION

		if (_detectionScoreThreshold < 0.25f) _detectionScoreThreshold = 0.25f;

		std::string licenseKey = "";
		float detConfThreshold = 0.5f;
		float detNMSThreshold = 0.5f;
		bool autoEngineDetection = true;
		bool loadEngineOnCreation = false;

		// Load detector model
		std::string detLabelMap;
		int detResult = _detector->LoadModelFromFolder(
			licenseKey.c_str(),
			"", "",
			_detectionScoreThreshold,
			detConfThreshold,
			detNMSThreshold,
			autoEngineDetection,
			_detectorModelType,
			_detectorDetectionType,
			loadEngineOnCreation,
			modelDirectory.c_str(),
			detLabelMap);

		if (detResult != 1) {
			return false;
		}

		// Load motion detector (algorithmic, empty model directory)
		std::string motionLabelMap;
		float motionScoreThreshold = 0.5f;
		int motionResult = _motiondetector->LoadModelFromFolder(
			licenseKey.c_str(),
			"motion", "motion.names",
			motionScoreThreshold,
			0.5f, 0.5f,
			autoEngineDetection,
			_motionModelType,
			_motionDetectionType,
			loadEngineOnCreation,
			"",
			motionLabelMap);

		labelMap = detLabelMap;
		return (detResult == 1);
	}
	catch (const std::exception& e) {
		return false;
	}
	catch (...) {
		return false;
	}
}
ANSCustomFallDetector::~ANSCustomFallDetector()
{
	Destroy();
}
std::vector<CustomObject> ANSCustomFallDetector::RunInference(const cv::Mat& input, const std::string& camera_id) {
	std::lock_guard<std::recursive_mutex> lock(_mutex);
	std::vector<CustomObject> results;
	std::vector<ANSCENTER::Object> output;
	if (input.empty())
	{
		return results;
	}
	if ((input.cols < 10) ||
		(input.rows < 10))
	{
		return results;
	}
	try {

		if (_detector && _motiondetector) {
			// Detect movement using the motion detector
			std::vector<ANSCENTER::Object> movementObjects;
			_motiondetector->RunInference(input, camera_id.c_str(), movementObjects);
			if (!movementObjects.empty()) {
				std::vector<ANSCENTER::Object> weaponRects;
				for (ANSCENTER::Object& obj : movementObjects) {
					weaponRects.push_back(obj);
				}
				int detectedArea = _detectedArea.width * _detectedArea.height;
				if ((!weaponRects.empty()) ||
					(detectedArea > 0)) {
					if ((_detectedArea.width < 10) || (_detectedArea.height < 10))
					{
						_detectedArea = GenerateMinimumSquareBoundingBox(weaponRects, 640);
						int x1_new = _detectedArea.x;
						int y1_new = _detectedArea.y;
						_detectedArea.x = std::max(x1_new, 0);
						_detectedArea.y = std::max(y1_new, 0);
						_detectedArea.width = std::min(640, input.cols - _detectedArea.x);
						_detectedArea.height = std::min(640, input.rows - _detectedArea.y);
					}

					cv::Mat frame = input.clone();
					std::vector<ANSCENTER::Object> detectedObjects;
					if (_detectedArea.width * _detectedArea.height > 0) {
						cv::Mat detectedObj = frame(_detectedArea).clone();
						_detector->RunInference(detectedObj, camera_id.c_str(), detectedObjects);
						detectedObj.release();
					}
					frame.release();
					// Do detections
					float maxScore = 0.0;
					if (!detectedObjects.empty()) {
						for (auto& detectedObj : detectedObjects) {
							detectedObj.box.x += _detectedArea.x;
							detectedObj.box.y += _detectedArea.y;
							detectedObj.cameraId = camera_id;
							if (detectedObj.confidence > this->_detectionScoreThreshold)
							{
								output.push_back(detectedObj);
								if (maxScore < detectedObj.confidence)
								{
									maxScore = detectedObj.confidence;
									int cropSize = 640;
									int imagegSize = std::max(input.cols, input.rows);
									if (imagegSize > 1920) cropSize = 640;
									else if (imagegSize > 1280) cropSize = 480;
									else if (imagegSize > 640) cropSize = 320;
									else cropSize = 224;
									maxScore = detectedObj.confidence;
									int x1 = detectedObj.box.x;
									int y1 = detectedObj.box.y;
									int xc = x1 + detectedObj.box.width / 2;
									int yc = y1 + detectedObj.box.height / 2;
									int  x1_new = std::max(xc - cropSize / 2, 0);
									int  y1_new = std::max(yc - cropSize / 2, 0);
									x1_new = std::min(x1_new, input.cols - cropSize);
									y1_new = std::min(y1_new, input.rows - cropSize);
									_detectedArea.x = std::max(x1_new, 0);
									_detectedArea.y = std::max(y1_new, 0);
									_detectedArea.width = std::min(cropSize, input.cols - _detectedArea.x);
									_detectedArea.height = std::min(cropSize, input.rows - _detectedArea.y);
								}
								_isFallDetected = true;
								_retainDetectedArea = 0;

							}
							else
							{
								if (_isFallDetected) {
									_retainDetectedArea++;
									if (_retainDetectedArea >= 40) {
										_detectedArea.width = 0;
										_detectedArea.height = 0;
										_retainDetectedArea = 0;
										_isFallDetected = false;
									}
								}
								else {
									_detectedArea.width = 0;
									_detectedArea.height = 0;
									_retainDetectedArea = 0;
								}
							}
						}
					}
					else {
						if (_isFallDetected) {
							_retainDetectedArea++;
							if (_retainDetectedArea >= 40) {
								_detectedArea.width = 0;
								_detectedArea.height = 0;
								_retainDetectedArea = 0;
								_isFallDetected = false;
							}
						}
						else {
							_detectedArea.width = 0;
							_detectedArea.height = 0;
							_retainDetectedArea = 0;
						}
					}
				}
			}
			for (auto obj : output) {
				if (obj.confidence >= this->_detectionScoreThreshold) {
					CustomObject customObj;
					customObj.box = obj.box;
					customObj.classId = obj.classId;
					customObj.confidence = obj.confidence;
					customObj.cameraId = obj.cameraId;
					customObj.className = obj.className;
					customObj.extraInfo = "Detection Score Threshold:" + std::to_string(_detectionScoreThreshold);
					results.push_back(customObj);
				}
			}

		}
		return results;
	}
	catch (std::exception& e) {
		return results;
	}
}
cv::Rect ANSCustomFallDetector::GenerateMinimumSquareBoundingBox(const std::vector<ANSCENTER::Object>& detectedObjects, int minSize) {
	int adMinSize = minSize - 20;
	if (adMinSize < 0) adMinSize = 0;
	if (detectedObjects.empty()) return cv::Rect(0, 0, minSize, minSize);

	int minX = detectedObjects[0].box.x;
	int minY = detectedObjects[0].box.y;
	int maxX = detectedObjects[0].box.x + detectedObjects[0].box.width;
	int maxY = detectedObjects[0].box.y + detectedObjects[0].box.height;

	for (const auto& rect : detectedObjects) {
		minX = std::min(minX, rect.box.x);
		minY = std::min(minY, rect.box.y);
		maxX = std::max(maxX, rect.box.x + rect.box.width);
		maxY = std::max(maxY, rect.box.y + rect.box.height);
	}

	int width = maxX - minX;
	int height = maxY - minY;
	int squareSize = std::max({ width, height, adMinSize });

	int centerX = minX + width / 2;
	int centerY = minY + height / 2;
	int squareX = centerX - squareSize / 2;
	int squareY = centerY - squareSize / 2;

	return cv::Rect(squareX - 10, squareY - 10, squareSize + 20, squareSize + 20);
}
