#include "ANSODTraininingEngine.h"
#include "Utility.h"
#include <random>
#include <cstdio>
#include <memory>
#include <regex>
namespace ANSCENTER
{
	ANSODTRE::ANSODTRE() {
		_projectDirectory = "C:\\ProgramData\\Sh7O7nUe7vJ";
		_engineDirectory = "C:\\ProgramData\\mt5P9nTe5vG\\Engine\\CUDA11";
		_modelTemplateDirectory = "C:\\ProgramData\\mt5P9nTe5vG\\Engine\\CUDA11";
		_currentExperimentNumber = 1;
		_isLicenseValid = false;
		_licenseKey = "";
		_zipPassword = "";
		_changeProjectDirectory = false;
		_isLatestEngine = false; // No use for this case
	}
	ANSODTRE::~ANSODTRE() {
		_projects.clear();
		_zipPassword = "";
		_isLicenseValid = false;
		_changeProjectDirectory = false;
		_isLatestEngine = false;
	}
	void ANSODTRE::CheckLicense() {

		try {
			_isLicenseValid = ANSCENTER::ANSLicenseHelper::LicenseVerification(_licenseKey, 1002, "ODHUB-LV");//Default productId=1002 (ODHUB)
		}
		catch (std::exception& e) {
			this->_logger.LogFatal("ANSODTRE::CheckLicense()", e.what(), __FILE__, __LINE__);
		}
	}
	bool ANSODTRE::Init(std::string licenseKey,
						std::string projectDirectory, 
						std::string engineDirectory, 
						std::string modelTemplateDirectory,
						std::string modelZipPassword, 
						bool isLatestEngine) 
	{
		_licenseKey = licenseKey;
		_isLatestEngine = isLatestEngine;
		CheckLicense();
		if (!_isLicenseValid) {
			this->_logger.LogError("ANSODTRE::Init", "Invalid license key", __FILE__, __LINE__);
			return false;
		}
		if (projectDirectory.empty())_projectDirectory = "C:\\ProgramData\\Sh7O7nUe7vJ\\odhub";
		else _projectDirectory = projectDirectory +"\\odhub";

		_engineDirectory = engineDirectory;
		_modelTemplateDirectory = modelTemplateDirectory;
		if (_engineDirectory.empty())_engineDirectory = "C:\\ProgramData\\ANSCENTER\\Shared";
		if (_modelTemplateDirectory.empty())_modelTemplateDirectory = "C:\\ProgramData\\ANSCENTER\\Shared";	

		// Change the project directory to 1 level down if there is 1 folder in the project directory
		if(!FolderExist(_projectDirectory)) fs::create_directory(_projectDirectory);

		_modelZipPassword = modelZipPassword;
		_zipPassword = "Sh7O7nUe7vJ/417W0gWX+dSdfcP9hUqtf/fEqJGqxYL3PedvHubJag==";
		_modelZipPassword = _zipPassword;


		// Some logging in here
		std::string _ansdnnEngine = _engineDirectory + "\\ansdn.exe";
		if (!FileExist(_ansdnnEngine))this->_logger.LogError("ANSODTRE::Init", "ODHUB Engine is not installed. Please install ANS Runtime Engine", __FILE__, __LINE__);

		std::string fastCfgTemplateFile = _modelTemplateDirectory + "\\FAST.cfg";
		std::string accurateCfgTemplateFile = _modelTemplateDirectory + "\\ACCURATE.cfg";
		std::string veryAccurateCfgTemplateFile = _modelTemplateDirectory + "\\VERYACCURATE.cfg";
		std::string customCfgTemplateFile = _modelTemplateDirectory + "\\CUSTOM.cfg";

		if (!FileExist(fastCfgTemplateFile))this->_logger.LogError("ANSODTRE::Init", "FAST.cfg template file is missing", __FILE__, __LINE__);
		if (!FileExist(accurateCfgTemplateFile))this->_logger.LogError("ANSODTRE::Init", "ACCURATE.cfg template file is missing", __FILE__, __LINE__);
		if (!FileExist(veryAccurateCfgTemplateFile))this->_logger.LogError("ANSODTRE::Init", "VERYACCURATE.cfg template file is missing", __FILE__, __LINE__);
		if (!FileExist(customCfgTemplateFile))this->_logger.LogError("ANSODTRE::Init", "CUSTOM.cfg template file is missing", __FILE__, __LINE__);

		std::string fastWeight = _modelTemplateDirectory + "\\FAST.weights";
		std::string accurateWeight = _modelTemplateDirectory + "\\BASE.weights";

		if (!FileExist(fastWeight))this->_logger.LogError("ANSODTRE::Init", "FAST model file is missing", __FILE__, __LINE__);
		if (!FileExist(accurateWeight))this->_logger.LogError("ANSODTRE::Init", "ACCURATE model file is missing", __FILE__, __LINE__);

		if (FolderExist(_projectDirectory) &&
			FolderExist(_engineDirectory) &&
			FolderExist(_modelTemplateDirectory))return true;
		return false;
	}
	std::vector<std::string> ANSODTRE::GetProjects() {
		_projects.clear();
		std::vector<std::string> _projectNames;
		for (auto& entry : std::filesystem::directory_iterator(_projectDirectory))
		{
			if (entry.is_directory())
			{
				_projects.push_back(entry.path().string());
				_projectNames.push_back(ExtractFolderName(entry.path().string()));
			}
		}

		return _projectNames;
	}
	std::vector<std::string> ANSODTRE::GetProjectExperiments(std::string projectName) {
		_projectName = Space2Underscore(projectName);
		std::string projectPath = FindFilePathInFileList(_projects, _projectName);
		std::string experimentPath = projectPath + "\\experiments";
		if (FolderExist(experimentPath)) {
			std::vector<std::string> files;
			std::vector<std::string> experiments;
			for (const auto& entry : std::filesystem::directory_iterator(experimentPath))
			{
				if (entry.is_regular_file() && entry.path().extension() == ".zip")
				{
					experiments.push_back(GetFileNameWithoutExtension(entry.path().string()));
				}
			}
			return experiments;
		}
		else {
			std::vector<std::string> experiments;
			return experiments;
		}
		
	}
	std::string ANSODTRE::GetProjectExperimentStatus(std::string projectName, int experimentNumber) {
		_projectName = Space2Underscore(projectName);
		std::string projectPath = FindFilePathInFileList(_projects, _projectName);
		if (FolderExist(projectPath)) {
			std::string experimentFolder = projectPath + "\\experiments";
			if (FolderExist(experimentFolder)) {
				std::string experimentPath = projectPath + "\\experiments\\" + std::to_string(experimentNumber) + ".zip";
				if (FileExist(experimentPath)) {
					std::string evaluationPath = projectPath + "\\evaluations\\mAP_" + std::to_string(experimentNumber) + ".txt";
					if (FileExist(evaluationPath)) return "EXPORTED";
					else return "TRAINED";
				}
				else {
					std::string configurationPath = projectPath + "\\experiments\\" + std::to_string(experimentNumber) + "\\train.cfg";
					if (FileExist(configurationPath))
					{
						std::string mapFile = projectPath + "\\experiments\\" + std::to_string(experimentNumber) + "\\mAP.txt";
						if (FileExist(mapFile)) return "RUNNING";
						else return "CONFIGURED";
					}
					else {
						return "UPLOADED";
					}
				}
			}
			else {
				std::string dataPath = projectPath + "\\data";
				if (FileExist(dataPath))
				{
					return "CREATED";
				}
				else {
					return "UPLOADED";
				}
			}
		}
		else return "NOTFOUND";
	}
	bool ANSODTRE::CreateProject(std::string projectName) {
		_projectName = Space2Underscore(projectName);
		std::string projectPath = FindFilePathInFileList(_projects, _projectName);
		if (projectPath.empty()) {
			std::string newProjectPath = _projectDirectory + "\\" + _projectName;
			if (FolderExist(newProjectPath))return false;
			fs::create_directory(newProjectPath);
			_projects.push_back(newProjectPath);
			return true;
		}
		return false;
	}
	bool ANSODTRE::DeleteProject(std::string projectName) {
		_projectName = Space2Underscore(projectName);
		std::string projectPath = FindFilePathInFileList(_projects, _projectName);
		if (FolderExist(projectPath)) {
			return DeleteFolder(projectPath);
		}
		return false;
	}
	bool ANSODTRE::SetWorkingDirectory(std::string workingDirectory) {
		if (FolderExist(workingDirectory)) {
			size_t lastSlashPos = workingDirectory.find_last_of("\\/");
			if (lastSlashPos != std::string::npos && workingDirectory.substr(lastSlashPos + 1) == "data") {
				workingDirectory = workingDirectory.substr(0, lastSlashPos);
			}
			_workingDirectory = workingDirectory;
			std::string dataPath = _workingDirectory + "\\data";
			std::string pbLabelPath = dataPath + "\\label_map.pbtxt";
			if (FileExist(pbLabelPath)) return true;

			else false;
		}
		return false;
	}
	bool ANSODTRE::UploadTrainingData(std::string projectName) {
		_projectName = Space2Underscore(projectName);
		std::string projectPath = FindFilePathInFileList(_projects, _projectName);
		std::string dataPath = projectPath + "\\data";
		if (FolderExist(dataPath))DeleteFolder(dataPath);		
		fs::create_directory(dataPath);
		std::string sourceDataPath = _workingDirectory + "\\data";
		bool useANSClassFile = true;
		std::string sourceLabelPath = sourceDataPath + "\\ANS_Class.json";
		if (!FileExist(sourceLabelPath)) {
			sourceLabelPath = sourceDataPath + "\\label_map.pbtxt";
			useANSClassFile = false;
		}
		if (!FileExist(sourceLabelPath)) return false;
		std::string classesName = dataPath + "\\classes.names";
		if (FileExist(classesName))DeleteFile(classesName);
		int numberOfClass = GenerateClassNamesFile(sourceLabelPath, classesName, useANSClassFile);
		if (numberOfClass <= 0) return false;
		GenerateTrainDataSet(sourceDataPath, dataPath + "\\train", dataPath + "\\test", 0.8);
		std::string objectDataFile = dataPath + "\\object.data";
		GenerateObjectDataFile(numberOfClass, objectDataFile);
		return FileExist(objectDataFile);
	}
	

	bool ANSODTRE::CreateTrainingEngine(std::string projectName,
										int experimentNumber,
										int extractorType, 
										long numberStep, 
										int batchSize, 
										double learningRate)// extractorType =0: FAST, 1: ACCURATE; 2: VERY ACCURATE
	{
		_projectName = Space2Underscore(projectName);
		const int _extractorType = extractorType;
		std::string projectPath = FindFilePathInFileList(_projects, _projectName);
		std::string experimentPath = projectPath + "\\experiments\\" + std::to_string(experimentNumber);
		std::string experimentParentFolder = GetParentFolder(experimentPath);
		if (!FolderExist(experimentParentFolder))fs::create_directory(experimentParentFolder);
		//0. Check if the experiement folder exist
		std::string zipExperimentPath = experimentPath + ".zip";	
		if (FileExist(zipExperimentPath)) { // If the experiment folder is already zipped, we need to unzip it
			bool unzipResult = ExtractProtectedZipFile(zipExperimentPath, _zipPassword, std::to_string(experimentNumber), experimentPath);
			if (!unzipResult) {
				unzipResult=ExtractProtectedZipFile(zipExperimentPath, _modelZipPassword, std::to_string(experimentNumber), experimentPath);
			}
			if (unzipResult) {
				DeleteFile(zipExperimentPath);
			}
			DeleteWeightsFiles(experimentPath);
		}

		//1. Create a brand new experiment folder
		if (!FolderExist(experimentPath))fs::create_directory(experimentPath);
		
		std::string trainCfgFile = experimentPath + "\\train.cfg";
		std::string testCfgFile = experimentPath + "\\test.cfg";
		std::string classesName = projectPath + "\\data\\classes.names";
		std::string objectDataFile = projectPath + "\\data\\object.data";

		//2. Delete the old configuration files
		if (FileExist(trainCfgFile))DeleteFile(trainCfgFile);
		if (FileExist(testCfgFile))DeleteFile(testCfgFile);

		//3. Copy classes.names to the experiment folder
		std::string experimentClassesName = experimentPath + "\\classes.names";
		if (FileExist(experimentClassesName)) DeleteFile(experimentClassesName);
		CopyFile(classesName, experimentClassesName);
		
		//4. Copy object.data to the experiment folder
		std::string experimentObjectDataFile = experimentPath + "\\object.data";
		if (FileExist(experimentObjectDataFile)) DeleteFile(experimentObjectDataFile);
		CopyFile(objectDataFile, experimentObjectDataFile);

		int numberOfClass = ReadClassesValue(experimentObjectDataFile);

		int step1 = static_cast<int>(numberStep * 0.8);
		int step2 = static_cast<int>(numberStep * 0.9);

		int subDivision = 64; // Hardcoded
		int filters = (numberOfClass + 5) * 3;

		//2. Create a training configuration file
		std::map<std::string, std::string> trainReplacements = {
				{"BATCH", std::to_string(batchSize)},
				{"SUBDIVISION", "64"},
				{"LEARNING_RATE",  std::to_string(learningRate)},
				{"MAX_BATCHES", std::to_string(numberStep)},
				{"STEP1", std::to_string(step1)},
				{"STEP2", std::to_string(step2)},
				{"FILTERS", std::to_string(filters)},
				{"CLASSES", std::to_string(numberOfClass)}};
		std::map<std::string, std::string> testReplacements = {
				{"BATCH", "1"},
				{"SUBDIVISION", "1"},
				{"LEARNING_RATE",  std::to_string(learningRate)},
				{"MAX_BATCHES", std::to_string(numberStep)},
				{"STEP1", std::to_string(step1)},
				{"STEP2", std::to_string(step2)},
				{"FILTERS", std::to_string(filters)},
				{"CLASSES", std::to_string(numberOfClass)}};
		
		std::string trainCfgTemplateFile;;
		if (_extractorType == 0) {
			trainCfgTemplateFile = _modelTemplateDirectory + "\\FAST.cfg";
		}
		else if (_extractorType == 1) {
			trainCfgTemplateFile = _modelTemplateDirectory + "\\ACCURATE.cfg";
		}
		else if (_extractorType == 2) {
			trainCfgTemplateFile = _modelTemplateDirectory + "\\VERYACCURATE.cfg";
		}
		else if (_extractorType == 3) {
			trainCfgTemplateFile = _modelTemplateDirectory + "\\CUSTOM.cfg";
		}
		else {
			trainCfgTemplateFile = _modelTemplateDirectory + "\\CUSTOM.cfg";
		}
		std::cout<< "trainCfgTemplateFile: " << trainCfgTemplateFile << std::endl;
		GenerateConfigFile(trainCfgTemplateFile, trainCfgFile, trainReplacements);
		GenerateConfigFile(trainCfgTemplateFile, testCfgFile, testReplacements);

		if (!FileExist(trainCfgFile))	this->_logger.LogError("ANSODTRE::CreateTrainingEngine", "Missing train.cfg file", __FILE__, __LINE__);
		if (!FileExist(testCfgFile))	this->_logger.LogError("ANSODTRE::CreateTrainingEngine", "Missing test.cfg file", __FILE__, __LINE__);
		if (!FileExist(experimentClassesName))	this->_logger.LogError("ANSODTRE::CreateTrainingEngine", "MMissing classes.names file", __FILE__, __LINE__);
		if (!FileExist(experimentObjectDataFile))	this->_logger.LogError("ANSODTRE::CreateTrainingEngine", "Missing object.data file", __FILE__, __LINE__);

		if ((FileExist(trainCfgFile)) && 
			(FileExist(testCfgFile)) &&
			(FileExist(experimentClassesName)) &&
			(FileExist(experimentObjectDataFile)))return true;
		return false;
	}
	std::string ANSODTRE::GenerateTrainingCommand(std::string projectName,int experimentNumber, int extractorType) {
		_projectName = Space2Underscore(projectName);
		std::string projectPath = FindFilePathInFileList(_projects, _projectName);
		if (FolderExist(projectPath)) {
			std::string experimentPath = projectPath + "\\experiments\\" + std::to_string(experimentNumber);
			std::string trainScript = _engineDirectory + "\\ansdn.exe";
			std::string experimentObjectDataFile = experimentPath + "\\object.data";
			std::string trainCfgFile = experimentPath + "\\train.cfg";
			std::string preTrainWeightFile = _modelTemplateDirectory + "\\FAST.weights";
			if (extractorType > 0)preTrainWeightFile = _modelTemplateDirectory + "\\BASE.weights";

			// Override the preTrainWeightFile if the experiment folder is already zipped
			std::string zipExperimentPath = experimentPath + ".zip";
			std::string experimentParentFolder = GetParentFolder(experimentPath);
			if (FileExist(zipExperimentPath)) { // If the experiment folder is already zipped, we need to unzip it
				bool unzipResult = ExtractProtectedZipFile(zipExperimentPath, _zipPassword, std::to_string(experimentNumber), experimentPath);
				if (!unzipResult) {
					unzipResult = ExtractProtectedZipFile(zipExperimentPath, _modelZipPassword, std::to_string(experimentNumber), experimentPath);
				}
				if (unzipResult) {
					DeleteFile(zipExperimentPath);
				}
			}
			std::string lastTrainWeightFile = experimentPath + "\\train_last.weights";
			if(FileExist(lastTrainWeightFile))preTrainWeightFile = experimentPath + "\\train_last.weights";

			if (!FileExist(trainScript))	this->_logger.LogError("ANSODTRE::GenerateTrainingCommand", "ANSODTRE Engine is not installed", __FILE__, __LINE__);
			if (!FileExist(trainCfgFile))	this->_logger.LogError("ANSODTRE::GenerateTrainingCommand", "Missing train.cfg file", __FILE__, __LINE__);
			if (!FileExist(preTrainWeightFile))	this->_logger.LogError("ANSODTRE::GenerateTrainingCommand", "Missing pre-trained weight file", __FILE__, __LINE__);
			if (!FileExist(experimentObjectDataFile))	this->_logger.LogError("ANSODTRE::GenerateTrainingCommand", "Missing object.data file", __FILE__, __LINE__);

			if ((FileExist(trainScript)) &&
				(FileExist(trainCfgFile)) &&
				(FileExist(preTrainWeightFile)) &&
				(FileExist(experimentObjectDataFile))) {
				std::string trainCommand = trainScript + " detector train " + "\""+ 
										   experimentObjectDataFile + "\""+ " " + +"\"" +
										   trainCfgFile + +"\"" + " " + "\"" +
										   preTrainWeightFile + "\"" +
										   " -dont_show " + " -map";
				return trainCommand;
			}
			this->_logger.LogError("ANSODTRE::GenerateTrainingCommand", "Missing configuration files", __FILE__, __LINE__);
			return "";
		}
		else {
			this->_logger.LogError("ANSODTRE::GenerateTrainingCommand", "project path is not exist", __FILE__, __LINE__);
			return "";

		}
	} 

	std::string ANSODTRE::GenerateCustomTrainingCommand(std::string projectName, int experimentNumber, int extractorType, std::string pretrainedModel) {
		_projectName = Space2Underscore(projectName);
		std::string projectPath = FindFilePathInFileList(_projects, _projectName);
		if (FolderExist(projectPath)) {
			std::string experimentPath = projectPath + "\\experiments\\" + std::to_string(experimentNumber);
			std::string trainScript = _engineDirectory + "\\ansdn.exe";
			std::string experimentObjectDataFile = experimentPath + "\\object.data";
			std::string trainCfgFile = experimentPath + "\\train.cfg";
			
			std::string preTrainWeightFile = pretrainedModel;
			if (!FileExist(preTrainWeightFile)) {
				preTrainWeightFile = _modelTemplateDirectory + "\\FAST.weights";
				if (extractorType > 0)preTrainWeightFile = _modelTemplateDirectory + "\\BASE.weights";
			}
			// Override the preTrainWeightFile if the experiment folder is already zipped
			std::string zipExperimentPath = experimentPath + ".zip";
			std::string experimentParentFolder = GetParentFolder(experimentPath);
			if (FileExist(zipExperimentPath)) { // If the experiment folder is already zipped, we need to unzip it
				bool unzipResult = ExtractProtectedZipFile(zipExperimentPath, _zipPassword, std::to_string(experimentNumber), experimentPath);
				if (!unzipResult) {
					unzipResult = ExtractProtectedZipFile(zipExperimentPath, _modelZipPassword, std::to_string(experimentNumber), experimentPath);
				}
				if (unzipResult) {
					DeleteFile(zipExperimentPath);
				}
			}
			std::string lastTrainWeightFile = experimentPath + "\\train_last.weights";
			if (FileExist(lastTrainWeightFile))preTrainWeightFile = experimentPath + "\\train_last.weights";

			if (!FileExist(trainScript))	this->_logger.LogError("ANSODTRE::GenerateCustomTrainingCommand", "ANSODTRE Engine is not installed", __FILE__, __LINE__);
			if (!FileExist(trainCfgFile))	this->_logger.LogError("ANSODTRE::GenerateCustomTrainingCommand", "Missing train.cfg file", __FILE__, __LINE__);
			if (!FileExist(preTrainWeightFile))	this->_logger.LogError("ANSODTRE::GenerateCustomTrainingCommand", "Missing pre-trained weight file", __FILE__, __LINE__);
			if (!FileExist(experimentObjectDataFile))	this->_logger.LogError("ANSODTRE::GenerateCustomTrainingCommand", "Missing object.data file", __FILE__, __LINE__);

			if ((FileExist(trainScript)) &&
				(FileExist(trainCfgFile)) &&
				(FileExist(preTrainWeightFile)) &&
				(FileExist(experimentObjectDataFile))) {
				std::string trainCommand = trainScript + " detector train " + "\"" +
					experimentObjectDataFile + "\"" + " " + +"\"" +
					trainCfgFile + +"\"" + " " + "\"" +
					preTrainWeightFile + "\"" +
					" -dont_show " + " -map";
				return trainCommand;
			}
			this->_logger.LogError("ANSODTRE::GenerateCustomTrainingCommand", "Missing configuration files", __FILE__, __LINE__);
			return "";
		}
		else {
			this->_logger.LogError("ANSODTRE::GenerateCustomTrainingCommand", "project path is not exist", __FILE__, __LINE__);
			return "";

		}
	}

	std::string ANSODTRE::GenerateEvaluateModelCommand(std::string projectName, int experimentNumber, bool &zipModelExist) {
		_projectName = Space2Underscore(projectName);
		 zipModelExist = false;

		std::string projectPath = FindFilePathInFileList(_projects, _projectName);
		if (FolderExist(projectPath)) {
			std::string experimentPath = projectPath + "\\experiments\\" + std::to_string(experimentNumber);
			std::string trainScript = _engineDirectory + "\\ansdn.exe";
			std::string experimentObjectDataFile = experimentPath + "\\object.data";
			std::string testCfgFile = experimentPath + "\\test.cfg";
			std::string preTrainWeightFile = experimentPath + "\\train_last.weights";
			// Override the preTrainWeightFile if the experiment folder is already zipped
			std::string zipExperimentPath = experimentPath + ".zip";
			std::string experimentParentFolder = GetParentFolder(experimentPath);
			if (FileExist(zipExperimentPath)) { // If the experiment folder is already zipped, we need to unzip it
				bool unzipResult = ExtractProtectedZipFile(zipExperimentPath, _zipPassword, std::to_string(experimentNumber), experimentPath);
				if (!unzipResult) {
					unzipResult = ExtractProtectedZipFile(zipExperimentPath, _modelZipPassword, std::to_string(experimentNumber), experimentPath);
				}
				if (unzipResult) {
					zipModelExist = true;
				}
			}
			if (!FileExist(trainScript))	this->_logger.LogError("ANSODTRE::GenerateEvaluateModelCommand", "ANSODTRE Engine is not installed", __FILE__, __LINE__);
			if (!FileExist(testCfgFile))	this->_logger.LogError("ANSODTRE::GenerateEvaluateModelCommand", "Missing train.cfg file", __FILE__, __LINE__);
			if (!FileExist(preTrainWeightFile))	this->_logger.LogError("ANSODTRE::GenerateEvaluateModelCommand", "Missing pre-trained weight file", __FILE__, __LINE__);
			if (!FileExist(experimentObjectDataFile))	this->_logger.LogError("ANSODTRE::GenerateEvaluateModelCommand", "Missing object.data file", __FILE__, __LINE__);

			if ((FileExist(trainScript)) &&
				(FileExist(testCfgFile)) &&
				(FileExist(preTrainWeightFile)) &&
				(FileExist(experimentObjectDataFile))) {
				std::string evaluateCommand = trainScript + " detector map " + "\"" +
										   experimentObjectDataFile + "\"" + " " + "\"" +  
										   testCfgFile + "\"" + " " + "\"" +
										   preTrainWeightFile + "\"" +
										   " -dont_show " + " -map";
				return evaluateCommand;
			}
			this->_logger.LogError("ANSODTRE::GenerateEvaluateModelCommand", "Missing configuration files", __FILE__, __LINE__);
			return "";
		}
		else {
			this->_logger.LogError("ANSODTRE::GenerateEvaluateModelCommand", "project path is not exist", __FILE__, __LINE__);
			return "";
		}
	}		

	std::string ANSODTRE::GenerateTestCommand(std::string projectName, int experimentNumber, bool& zipModelExist, const std::string& testDataFolder) {
		_projectName = Space2Underscore(projectName);
		zipModelExist = false;
		std::string projectPath = FindFilePathInFileList(_projects, _projectName);
		if (FolderExist(projectPath)) {
			std::string experimentPath = projectPath + "\\experiments\\" + std::to_string(experimentNumber);
			std::string trainScript = _engineDirectory + "\\ansdn.exe";
			std::string experimentObjectDataFile = experimentPath + "\\object.data";
			std::string testCfgFile = experimentPath + "\\test.cfg";
			std::string testClassesName = experimentPath + "\\classes.names";
			std::string preTrainWeightFile = experimentPath + "\\train_last.weights";
			// Override the preTrainWeightFile if the experiment folder is already zipped
			std::string zipExperimentPath = experimentPath + ".zip";
			std::string experimentParentFolder = GetParentFolder(experimentPath);
			if (FileExist(zipExperimentPath)) { // If the experiment folder is already zipped, we need to unzip it
				bool unzipResult = ExtractProtectedZipFile(zipExperimentPath, _zipPassword, std::to_string(experimentNumber), experimentPath);
				if (!unzipResult) {
					unzipResult = ExtractProtectedZipFile(zipExperimentPath, _modelZipPassword, std::to_string(experimentNumber), experimentPath);
				}
				if (unzipResult) {
					zipModelExist = true;
				}
			}
			if (!FileExist(trainScript))	this->_logger.LogError("ANSODTRE::GenerateTestCommand", "ANSODTRE Engine is not installed", __FILE__, __LINE__);
			if (!FileExist(testCfgFile))	this->_logger.LogError("ANSODTRE::GenerateTestCommand", "Missing train.cfg file", __FILE__, __LINE__);
			if (!FileExist(preTrainWeightFile))	this->_logger.LogError("ANSODTRE::GenerateTestCommand", "Missing pre-trained weight file", __FILE__, __LINE__);
			if (!FileExist(experimentObjectDataFile))	this->_logger.LogError("ANSODTRE::GenerateTestCommand", "Missing object.data file", __FILE__, __LINE__);

			// Create test folder.
			// testDataFolder
			// - data: image files
			// - test.cfg
			// - classes.names
			// - object.data
			// - test.txt
			// Assuming that the data folder is already created and contains the image files
			std::string testFolder_ClassesNames = testDataFolder + "\\classes.names";
			std::string testFolder_ObjectData = testDataFolder + "\\object.data";
			std::string testFolder_TestTxt = testDataFolder + "\\test.txt";

			// Copy the classes.names file to testFolder_ClassesNames
			if (FileExist(testFolder_ClassesNames))DeleteFile(testFolder_ClassesNames);
			CopyFile(testClassesName, testFolder_ClassesNames);

			// List all the image files in the data folder and write to test.txt
			if (FileExist(testFolder_TestTxt))DeleteFile(testFolder_TestTxt);

			std::ofstream testTxtFile(testFolder_TestTxt);
			if (testTxtFile.is_open()) {
				for (const auto& entry : std::filesystem::directory_iterator(testDataFolder)) {
					if ((entry.is_regular_file() && entry.path().extension() == ".jpg")||
						(entry.is_regular_file() && entry.path().extension() == ".jpeg")|| 
						(entry.is_regular_file() && entry.path().extension() == ".png")||
						(entry.is_regular_file() && entry.path().extension() == ".bmp"))
					{  
						testTxtFile << entry.path().string() << std::endl;
					}
				}
				testTxtFile.close();
			}
			else {
				this->_logger.LogError("ANSODTRE::GenerateTestCommand", "Cannot open test.txt file to write image paths", __FILE__, __LINE__);
				return "";
			}

			// Create the object.data file in the test folder with follow contents
			//classes = 4 (read from classes.names file)
			//test = testTxtFile path
			int classCount = 0;
			std::ifstream classesFile(testFolder_ClassesNames);
			if (classesFile.is_open()) {
				std::string line;
				while (std::getline(classesFile, line)) {
					classCount++;
				}
				classesFile.close();
			}
			else {
				this->_logger.LogError("ANSODTRE::GenerateTestCommand", "Cannot open classes.names file to read class count", __FILE__, __LINE__);
				return "";
			}
			std::ofstream testObjectFile(testFolder_ObjectData);
			if (testObjectFile.is_open()) {
				testObjectFile << "classes = " << classCount << std::endl;
				testObjectFile << "test = " << testFolder_TestTxt << std::endl;
				testObjectFile.close();
			}
			else {
				this->_logger.LogError("ANSODTRE::GenerateTestCommand", "Cannot open object.data file to write image paths", __FILE__, __LINE__);
				return "";
			}
			
			if ((FileExist(trainScript)) &&
				(FileExist(testCfgFile)) &&
				(FileExist(preTrainWeightFile)) &&
				(FileExist(experimentObjectDataFile))) {
				std::string evaluateCommand = trainScript + " detector map " + "\"" +
					testFolder_ObjectData + "\"" + " " + "\"" +
					testCfgFile + "\"" + " " + "\"" +
					preTrainWeightFile + "\"" +
					" -dont_show " + " -map";
				return evaluateCommand;
			}
			this->_logger.LogError("ANSODTRE::GenerateTestCommand", "Missing configuration files", __FILE__, __LINE__);
			return "";
		}
		else {
			this->_logger.LogError("ANSODTRE::GenerateTestCommand", "project path is not exist", __FILE__, __LINE__);
			return "";
		}
	
	}

	
	bool ANSODTRE::DownloadModel(std::string projectName, int experimentNumber, std::string destZipDirectory, int modelModel) {
		_projectName = Space2Underscore(projectName);
		std::string projectPath = FindFilePathInFileList(_projects, _projectName);
		if (FolderExist(projectPath)) {
			//0. Get all files in the experiment folder
			std::string experimentPath = projectPath + "\\experiments\\" + std::to_string(experimentNumber);	
			std::string zipExperimentPath = experimentPath + ".zip";
			std::string experimentParentFolder = GetParentFolder(experimentPath);
			if (FileExist(zipExperimentPath)) { // If the experiment folder is already zipped, we need to unzip it
				if (ExtractProtectedZipFile(zipExperimentPath, _zipPassword, std::to_string(experimentNumber), experimentPath))
				{
					DeleteFile(zipExperimentPath);
				}
				else {
					this->_logger.LogError("ANSODTRE::DownloadModel", "Can not unzip model zip file to extract information. Invalid password.", __FILE__, __LINE__);
					return false;
				}
			}
			//1. Remove previuos weights files
			DeleteWeightsFiles(experimentPath);

			// 2. Get the names of classes.names
			std::string classesNameFile = experimentPath + "\\classes.names";
			std::ifstream inFile(classesNameFile);

			if (!inFile.is_open()) {
				this->_logger.LogError("ANSODTRE::DownloadModel", "cannot open file to save model information.", __FILE__, __LINE__);
				return false;
			}

			std::vector<std::string> classesNames;
			std::string line;
			while (std::getline(inFile, line)) {
				classesNames.push_back(line);
			}

			inFile.close();

			// 3. Create Categories folder
			std::string categoriesFolder = experimentPath + "\\Categories";
			if (FolderExist(categoriesFolder))DeleteFolder(categoriesFolder);
			fs::create_directory(categoriesFolder);

			for (size_t i = 0; i < classesNames.size(); ++i) {
				std::ofstream outFile(categoriesFolder + "/" + classesNames[i] + "." + std::to_string(i+1));
				outFile.close();
			}

			//4. Create a text file name yolo.txt
			std::string yoloTxtFile = experimentPath + "\\yolo.txt";
			std::ofstream outFile(yoloTxtFile);

			outFile.close();

			//5. Zip the experiment folder
			std::string zipFile = experimentParentFolder + "\\" + std::to_string(experimentNumber) + ".zip";
			if (ZipFolderWithPassword(experimentPath.c_str(), zipFile.c_str(), _zipPassword.c_str())) {
				DeleteFolder(experimentPath);
				std::string destZipFilePath = destZipDirectory + "\\" + projectName + "_" + std::to_string(experimentNumber) + ".zip";
				CopyFile(zipFile, destZipFilePath);
				if (FileExist(destZipFilePath))return true;
				return false;
			}
			this->_logger.LogError("ANSODTRE::DownloadModel", "Can not unzip model zip file to extract information. Invalid password.", __FILE__, __LINE__);
			return false;
		}
		else {
			this->_logger.LogError("ANSODTRE::DownloadModel", "Invalid project path.", __FILE__, __LINE__);
			return false;
		}
	}

	std::string ANSODTRE::EvaluateModel(std::string projectName, int experimentNumber) {
		bool zipModelExist = false;
		std::string _projectName = Space2Underscore(projectName);
		std::string projectPath = FindFilePathInFileList(_projects, _projectName);
		if (FolderExist(projectPath)) {
			std::string experimentPath = projectPath + "\\experiments\\" + std::to_string(experimentNumber);
			std::string zipExperimentPath = experimentPath + ".zip";
			std::string evaluateCommand = GenerateEvaluateModelCommand(_projectName, experimentNumber, zipModelExist);
			if (evaluateCommand.empty()) {
				this->_logger.LogError("ANSODTRE::EvaluateModel", "Invalid command to evaluate model", __FILE__, __LINE__);
				return "";
			}
			std::string result = ExecuteCommand(evaluateCommand.c_str());

			if (result.empty()) {
				this->_logger.LogError("ANSODTRE::EvaluateModel", "Failed to evaluate model", __FILE__, __LINE__);
				return "";
			}
			//Remove previuos weights files
			DeleteWeightsFiles(experimentPath);
			if (FileExist(zipExperimentPath))DeleteFile(zipExperimentPath);
			//Zip the experiment folder
			if (ZipFolderWithPassword(experimentPath.c_str(), zipExperimentPath.c_str(), _zipPassword.c_str())) {
				DeleteFolder(experimentPath);
			}
			return ParseMAP(result);
		}
		else {
			this->_logger.LogError("ANSODTRE::EvaluateModel", "Invalid project path.", __FILE__, __LINE__);
			return "";
		}
	}
		
	std::string ANSODTRE::ParseTrainingResults(std::string trainingResults) {
		return trainingResults;
	}

	bool ANSODTRE::CheckEngineStatus() {
		const char* targetProcess = "ansdn.exe"; // Replace with the target process name

		if (IsProcessRunning(targetProcess)) {
			return true;
		}
		else {
			return false;
		}
	}

	bool ANSODTRE::CheckEngine() {
		std::string inferenceCommand = "C:\\ProgramData\\ANSCENTER\\Shared\\ansdn.exe detector";
		std::string result = ExecuteCommand(inferenceCommand.c_str());
		if (result.empty()) {
			return false;
		}
		else {
			std::string target = "ODHUB_TRAIN_COMPLETE";
			// Check if the target text is found in the input text
			if (result.find(target) != std::string::npos) {
				std::cout << "Engine is installed correctly." << std::endl;
				return true;  // Text is found
			}
			else {
				std::cout << "Engine is not installed correctly." << std::endl;
				return false;  // Text is not found
			}
		}
	}
	//Private functions
	std::string ANSODTRE::ParseMAP(const std::string& text) {
		std::regex mapPattern("mean average precision \\(mAP@0\\.50\\) = ([0-9\\.]+)");
		std::smatch matches;
		if (std::regex_search(text, matches, mapPattern) && matches.size() > 1) {
			return matches[1].str();
		}
		else {
			return "0.00";
		}
	}
	std::string ANSODTRE::AdjustRootDirectory(const std::string& root) {
		std::filesystem::path rootPath(root);
		std::vector<std::filesystem::path> subdirectories;
		for (const auto& entry : std::filesystem::directory_iterator(rootPath)) {
			if (entry.is_directory()) {
				subdirectories.push_back(entry.path());
			}
		}
		if (subdirectories.size() == 1) {
			return subdirectories[0].string();
		}
		else {
			return root;
		}
	}
	int ANSODTRE::GenerateClassNamesFile(std::string labelMapFile, std::string objectDataFile, bool useANSClassFile) {
		if (!FileExist(labelMapFile)) return 0;
		if (useANSClassFile) {
			// Use the ANS Class file
			std::vector<std::string> classVec = ParseANSClass(labelMapFile);
			if (classVec.empty()) return 0;
			std::ofstream outFile(objectDataFile);
			std::string line;
			int fileLine = 0;
			for (auto& item : classVec) {
				outFile << item << "\n";
				fileLine++;
			}
			outFile.close();
			return fileLine;
		}
		else {
			std::ifstream inFile(labelMapFile);
			std::ofstream outFile(objectDataFile);
			if (!inFile.is_open() || !outFile.is_open()) {
				return 0;
			}
			std::string line;
			// Updated pattern: allows optional spaces between name and the colon
			std::regex namePattern(R"(name\s*:\s*'(.*)')");
			int fileLine = 0;

			while (std::getline(inFile, line)) {
				std::smatch matches;
				std::cout << line << std::endl;
				if (std::regex_search(line, matches, namePattern) && matches.size() > 1) {
					outFile << matches[1].str() << "\n";
					fileLine++;
				}
			}
			inFile.close();
			outFile.close();
			return fileLine;
		}
	}
	bool ANSODTRE::GenerateObjectDataFile(int numberOfClass, std::string objectDataFile) {
		std::string projectName = GetParentFolder(objectDataFile);
		if (FileExist(objectDataFile))DeleteFile(objectDataFile);
		std::string classesName = projectName + "\\classes.names";	
		std::string trainTxt = projectName + "\\train.txt";
		std::string testTxt = projectName + "\\test.txt";
		if ((FileExist(classesName)) &&
			(FileExist(trainTxt)) &&
			(FileExist(testTxt))) 
		{
			std::ofstream outFile(objectDataFile);
			outFile << "classes = " << numberOfClass << "\n";
			outFile << "train = " << trainTxt << "\n";
			outFile << "valid = " << testTxt << "\n";
			outFile << "names = " << classesName << "\n";
			outFile.close();
			return true;
		}
		return false;
	}
	void ANSODTRE::GenerateTrainDataSet(const std::string& sourceFolder, const std::string& trainFolder, const std::string& testFolder, double trainRatio) {
		if (FolderExist(trainFolder))DeleteFolder(trainFolder);
		if (FolderExist(testFolder))DeleteFolder(testFolder);
		fs::create_directory(trainFolder);
		fs::create_directory(testFolder);
		std::unordered_map<std::string, std::vector<std::filesystem::path>> fileGroups;
		for (const auto& entry : std::filesystem::directory_iterator(sourceFolder)) {
			std::string extension = entry.path().extension().string();
			boost::algorithm::to_lower(extension);

			if (entry.is_regular_file()) {
				if (extension == ".jpg" ||
					extension == ".jpeg" ||
					extension == ".png" ||
					extension == ".bmp") fileGroups[entry.path().stem().string()].push_back(entry.path());
			}
		}

		std::vector<std::string> groupNames;
		for (const auto& group : fileGroups) {
			groupNames.push_back(group.first);
		}

		std::ofstream trainFile(std::filesystem::path(trainFolder).parent_path() / "train.txt");
		std::ofstream testFile(std::filesystem::path(testFolder).parent_path() / "test.txt");

		size_t trainSize = static_cast<size_t>(groupNames.size() * trainRatio);
		for (size_t i = 0; i < groupNames.size(); ++i) {
			const auto& groupName = groupNames[i];
			const auto& files = fileGroups[groupName];
			for (const auto& file : files) {
				std::string txtSourceFile = sourceFolder + "\\" + GetFileNameWithoutExtension(file.filename().string()) + ".txt";
				if (i < trainSize) {
					std::string txtTrainFile = trainFolder + "\\" + GetFileNameWithoutExtension(file.filename().string()) + ".txt";
					std::string fileName = file.filename().string();
					std::string fileExtension = GetFileExtension(fileName);
					if (fileExtension == ".txt" ||
						fileExtension == ".jpeg" ||
						fileExtension == ".jpg" ||
						fileExtension == ".png" ||
						fileExtension == ".bmp")
					{
						if (FileExist(txtSourceFile)) {
							std::filesystem::copy(file, trainFolder / file.filename());
							CopyFile(txtSourceFile, txtTrainFile);
							trainFile << (trainFolder / file.filename()).string() << "\n";
						}
					}
					
				}
				else {
					std::string fileName = file.filename().string();
					std::string fileExtension = GetFileExtension(fileName);
					if (fileExtension == ".txt" ||
						fileExtension == ".jpeg" ||
						fileExtension == ".jpg" ||
						fileExtension == ".png" ||
						fileExtension == ".bmp") 
					{
						if (FileExist(txtSourceFile)) {
							std::string txtTestFile = testFolder + "\\" + GetFileNameWithoutExtension(file.filename().string()) + ".txt";
							std::filesystem::copy(file, testFolder / file.filename());
							CopyFile(txtSourceFile, txtTestFile);
							testFile << (testFolder / file.filename()).string() << "\n";
						}
					}
				}
			}
		}
		trainFile.close();
		testFile.close();
	}
	bool ANSODTRE::GenerateConfigFile(const std::string& templateFile, 
									  const std::string& outputFile, 
									  const std::map<std::string,std::string>& replacements) {
		std::ifstream inFile(templateFile);
		std::ofstream outFile(outputFile);

		if (!inFile.is_open() || !outFile.is_open()) {
			return false;
		}

		std::string line;
		while (std::getline(inFile, line)) {
			for (const auto& replacement : replacements) {
				size_t pos = line.find("$" + replacement.first);
				while (pos != std::string::npos) {
					line.replace(pos, replacement.first.size() + 1, replacement.second);
					pos = line.find("$" + replacement.first, pos + replacement.second.size());
				}
			}
			outFile << line << "\n";
		}

		inFile.close();
		outFile.close();

		return FileExist(outputFile);
	}
	int ANSODTRE::ReadClassesValue(const std::string& filePath) {
		std::ifstream inFile(filePath);
		if (!inFile.is_open()) {
			return 0;
		}
		std::ofstream outFile(filePath, std::ios_base::app);
		if (!outFile.is_open()) {
			return 0;
		}
		std::string currentFolderName = GetParentFolder(filePath);
		outFile << "backup = " << currentFolderName << "\n";
		outFile.close();

		std::string line;
		while (std::getline(inFile, line)) {
			if (line.find("classes") == 0) {
				size_t pos = line.find("=");
				if (pos != std::string::npos) {
					std::string classesValue = line.substr(pos + 1);
					inFile.close();
					return std::stoi(classesValue);
				}
			}
		}
		inFile.close();
		return 0;
	}
	void ANSODTRE::DeleteWeightsFiles(const std::string& folderPath) {
		for (const auto& entry : std::filesystem::directory_iterator(folderPath)) {
			if (entry.is_regular_file() && entry.path().extension() == ".weights" && entry.path().filename() != "train_last.weights") {
				std::filesystem::remove(entry.path());
			}
		}
	}
}

