// test.cpp : 此文件包含 "main" 函数。程序执行将在此处开始并结束。
//

#include <iostream>

#include <windows.h>
#include <io.h>
#include <ShellAPI.h>
#include <atltypes.h>

#include <direct.h>
#include "tlhelp32.h"

#include "YoloFeatureExtractor.h"
#include "ImageProcessor.h"
#include "SQLiteVecManager.h"

#include <filesystem>

#include <opencv2/core.hpp>
#include <opencv2/core/ocl.hpp>
#include <opencv2/highgui.hpp>
#include <opencv2/imgcodecs.hpp>
#include <opencv2/imgproc.hpp>





void checkBuild()
{
	//查询opencv编译时配置
	std::cout << cv::getBuildInformation() << std::endl;
}

void checkSimd()
{
	//查询opencv线程
	int numTh = cv::getNumThreads(); //默认值是cpu的逻辑线程数
	int numCore = cv::getNumberOfCPUs();
	std::cout << "getNumThreads=" << numTh << std::endl;
	std::cout << "getNumberOfCPUs=" << numCore << std::endl;

	//查询opencv当前是否开启了并行优化功能
	bool opt = cv::useOptimized(); //默认值是true
	std::cout << "useOptimized=" << opt << std::endl;

	//查询opencv当前是否支持具体的CPU指令集
	bool check1 = cv::checkHardwareSupport(CV_CPU_SSE4_1);
	bool check2 = cv::checkHardwareSupport(CV_CPU_SSE4_2);
	bool check3 = cv::checkHardwareSupport(CV_CPU_AVX2);
	std::cout << "CV_CPU_SSE4_1=" << check1 << std::endl;
	std::cout << "CV_CPU_SSE4_2=" << check2 << std::endl;
	std::cout << "CV_CPU_AVX2=" << check3 << std::endl;

	//查询完整的硬件支持清单
	std::cout << "HardwareSupport:" << std::endl;
	std::cout << "CV_CPU_MMX: " << cv::checkHardwareSupport(CV_CPU_MMX) << std::endl;
	std::cout << "CV_CPU_SSE: " << cv::checkHardwareSupport(CV_CPU_SSE) << std::endl;
	std::cout << "CV_CPU_SSE2: " << cv::checkHardwareSupport(CV_CPU_SSE2) << std::endl;
	std::cout << "CV_CPU_SSE3: " << cv::checkHardwareSupport(CV_CPU_SSE3) << std::endl;
	std::cout << "CV_CPU_SSSE3: " << cv::checkHardwareSupport(CV_CPU_SSSE3) << std::endl;
	std::cout << "CV_CPU_SSE4_1: " << cv::checkHardwareSupport(CV_CPU_SSE4_1) << std::endl;
	std::cout << "CV_CPU_SSE4_2: " << cv::checkHardwareSupport(CV_CPU_SSE4_2) << std::endl;
	std::cout << "CV_CPU_POPCNT: " << cv::checkHardwareSupport(CV_CPU_POPCNT) << std::endl;
	std::cout << "CV_CPU_FP16: " << cv::checkHardwareSupport(CV_CPU_FP16) << std::endl;
	std::cout << "CV_CPU_AVX: " << cv::checkHardwareSupport(CV_CPU_AVX) << std::endl;
	std::cout << "CV_CPU_AVX2: " << cv::checkHardwareSupport(CV_CPU_AVX2) << std::endl;
	std::cout << "CV_CPU_FMA3: " << cv::checkHardwareSupport(CV_CPU_FMA3) << std::endl;
	std::cout << "CV_CPU_AVX_512F: " << cv::checkHardwareSupport(CV_CPU_AVX_512F) << std::endl;
	std::cout << "CV_CPU_AVX_512BW: " << cv::checkHardwareSupport(CV_CPU_AVX_512BW) << std::endl;
	std::cout << "CV_CPU_AVX_512CD: " << cv::checkHardwareSupport(CV_CPU_AVX_512CD) << std::endl;
	std::cout << "CV_CPU_AVX_512DQ: " << cv::checkHardwareSupport(CV_CPU_AVX_512DQ) << std::endl;
	std::cout << "CV_CPU_AVX_512ER: " << cv::checkHardwareSupport(CV_CPU_AVX_512ER) << std::endl;
	std::cout << "CV_CPU_AVX_512IFMA512: " << cv::checkHardwareSupport(CV_CPU_AVX_512IFMA512) << std::endl;
	std::cout << "CV_CPU_AVX_512IFMA: " << cv::checkHardwareSupport(CV_CPU_AVX_512IFMA) << std::endl;
	std::cout << "CV_CPU_AVX_512PF: " << cv::checkHardwareSupport(CV_CPU_AVX_512PF) << std::endl;
	std::cout << "CV_CPU_AVX_512VBMI: " << cv::checkHardwareSupport(CV_CPU_AVX_512VBMI) << std::endl;
	std::cout << "CV_CPU_AVX_512VL: " << cv::checkHardwareSupport(CV_CPU_AVX_512VL) << std::endl;
	std::cout << "CV_CPU_NEON: " << cv::checkHardwareSupport(CV_CPU_NEON) << std::endl;
	std::cout << "CV_CPU_VSX: " << cv::checkHardwareSupport(CV_CPU_VSX) << std::endl;
	std::cout << "CV_CPU_AVX512_SKX: " << cv::checkHardwareSupport(CV_CPU_AVX512_SKX) << std::endl;
	std::cout << "CV_HARDWARE_MAX_FEATURE: " << cv::checkHardwareSupport(CV_HARDWARE_MAX_FEATURE) << std::endl;
	std::cout << std::endl;

	//cv::setUseOptimized(false);
	//cv::setNumThreads(1);
}

void checkOpenCL() //Open Computing Language:开放计算语言,可以附加在主机处理器的CPU或GPU上执行
{
	std::vector<cv::ocl::PlatformInfo> info;
	cv::ocl::getPlatfomsInfo(info);
	cv::ocl::PlatformInfo sdk = info.at(0);

	int number = sdk.deviceNumber();
	if (number < 1)
	{
		std::cout << "Number of devices:" << number << std::endl;
		return;
	}

	std::cout << "***********SDK************" << std::endl;
	std::cout << "Name:" << sdk.name() << std::endl;
	std::cout << "Vendor:" << sdk.vendor() << std::endl;
	std::cout << "Version:" << sdk.version() << std::endl;
	std::cout << "Version:" << sdk.version() << std::endl;
	std::cout << "Number of devices:" << number << std::endl;

	for (int i = 0; i < number; i++)
	{
		std::cout << std::endl;

		cv::ocl::Device device;
		sdk.getDevice(device, i);

		std::cout << "***********Device " << i + 1 << "***********" << std::endl;
		std::cout << "Vendor Id:" << device.vendorID() << std::endl;
		std::cout << "Vendor name:" << device.vendorName() << std::endl;
		std::cout << "Name:" << device.name() << std::endl;
		std::cout << "Driver version:" << device.vendorID() << std::endl;

		if (device.isAMD())
			std::cout << "Is AMD device" << std::endl;

		if (device.isIntel())
			std::cout << "Is Intel device" << std::endl;

		if (device.isNVidia())
			std::cout << "Is NVidia device" << std::endl;

		std::cout << "Global Memory size:" << device.globalMemSize() << std::endl;
		std::cout << "Memory cache size:" << device.globalMemCacheSize() << std::endl;
		std::cout << "Memory cache type:" << device.globalMemCacheType() << std::endl;
		std::cout << "Local Memory size:" << device.localMemSize() << std::endl;
		std::cout << "Local Memory type:" << device.localMemType() << std::endl;
		std::cout << "Max Clock frequency:" << device.maxClockFrequency() << std::endl;
	}
}

std::wstring getExePath()
{
	wchar_t exeFullPath[MAX_PATH]; // Full path
	std::wstring strPath = L"";

	GetModuleFileName(NULL, exeFullPath, MAX_PATH);
	strPath = (std::wstring)exeFullPath;    // Get full path of the file

	return strPath;
}

std::wstring GetProgramDir()
{
	wchar_t exeFullPath[MAX_PATH]; // Full path
	std::wstring strPath = L"";

	GetModuleFileName(NULL, exeFullPath, MAX_PATH);
	strPath = (std::wstring)exeFullPath;    // Get full path of the file

	int pos = strPath.find_last_of('\\', strPath.length());
	return strPath.substr(0, pos);  // Return the directory without the file name
}

int main()
{
	try
	{
		// 检查 OpenCL 是否可用
		checkBuild();
		checkSimd();
		//checkOpenCL();

		// 检查OpenCL是否可用
		if (!cv::ocl::haveOpenCL())
		{
			std::cout << "opencl不可用" << std::endl;
		}
		else
		{
			std::cout << "opencl可用" << std::endl;

			// 启用OpenCL
			//cv::ocl::setUseOpenCL(true);
		}

		// 设置路径（可根据实际情况修改）
		std::wstring wsExePath = getExePath();
		std::wstring wsProgramDir = GetProgramDir();

		std::filesystem::path mainDir = wsProgramDir;

#ifdef _WIN64
		std::string sMainDir = mainDir.parent_path().parent_path().string();
#else
		std::string sMainDir = mainDir.parent_path().string();
#endif // WIN32		

		std::string modelPath = sMainDir + "/best.onnx";           // YOLO2026模型路径
		std::string classesPath = sMainDir + "/cls.names";             // 类别文件路径
		std::string galleryDir = sMainDir + "/images";       // 图库目录路径
		std::string searchImagePath = sMainDir + "/3.jpg"; // 搜索图片路径
		std::string databasePath = sMainDir + "/image_features.db";     // SQLite数据库路径

		std::filesystem::remove(databasePath);

		std::cout << "=== YOLO2026图像检索系统 (SQLite-Vec版本) ===" << std::endl;
		std::cout << "模型路径: " << modelPath << std::endl;
		std::cout << "图库路径: " << galleryDir << std::endl;
		std::cout << "查询图片: " << searchImagePath << std::endl;
		std::cout << "数据库路径: " << databasePath << std::endl;
		std::cout << "=========================================" << std::endl;

		// 初始化特征提取器
		std::cout << "正在初始化YOLO2026特征提取器..." << std::endl;
		YoloFeatureExtractor extractor(modelPath, classesPath);

		extractor.initOpenCL(); // 启用OpenCL加速（如果可用）

		// 获取图库中的所有图片
		std::cout << "正在扫描图库中的图片..." << std::endl;
		std::vector<std::string> galleryImages = ImageProcessor::getImagesInDirectory(galleryDir);
		std::cout << "找到 " << galleryImages.size() << " 张图片" << std::endl;

		if (galleryImages.empty())
		{
			std::cerr << "图库中没有找到有效的图片文件！" << std::endl;
			return -1;
		}

		// 初始化SQLite-Vec数据库
		std::cout << "正在初始化SQLite-Vec数据库..." << std::endl;
		SQLiteVecManager vecManager(databasePath);

		auto start_time = std::chrono::high_resolution_clock::now();

		// 检查数据库是否已存在数据
		bool databaseExists = vecManager.loadDatabase();
		std::cout << "数据库加载结果：" << std::boolalpha << databaseExists << std::endl;
		std::cout << "数据集数量：" << vecManager.getFeatureCount() << std::endl;
		if (databaseExists && vecManager.getFeatureCount() > 0)
		{
			std::cout << "数据库已存在 " << vecManager.getFeatureCount() << " 条特征记录" << std::endl;
		}
		else
		{
			std::cout << "重新构建数据库..." << std::endl;

			// 提取图库图片特征并向量
			std::cout << "开始提取图库图片特征..." << std::endl;
			int featureDimension = 0;
			int processedCount = 0;

			// 初始化数据库表结构
			vecManager.initializeDatabase(45); // 假设特征维度为1000，实际会在第一次处理时确定

			for (size_t i = 0; i < galleryImages.size(); ++i)
			{
				try
				{
					std::vector<float> features = extractor.extractFeatures(galleryImages[i]);
					if (!features.empty())
					{
						if (featureDimension == 0)
						{
							featureDimension = static_cast<int>(features.size());
							std::cout << "特征维度: " << featureDimension << std::endl;
							// 重新初始化数据库以匹配实际维度
							vecManager.initializeDatabase(featureDimension);
						}
						vecManager.addFeatureVector(features, galleryImages[i]);
						processedCount++;
						std::cout << "[" << (i + 1) << "/" << galleryImages.size() << "] 已处理: "
							<< galleryImages[i] << std::endl;
					}
				}
				catch (const std::exception & e)
				{
					std::cerr << "处理失败 [" << (i + 1) << "]: " << galleryImages[i]
						<< " (" << e.what() << ")" << std::endl;
				}
			}

			auto end_time = std::chrono::high_resolution_clock::now();
			auto duration = std::chrono::duration_cast<std::chrono::seconds>(end_time - start_time);

			std::cout << "特征提取完成，耗时: " << duration.count() << " 秒" << std::endl;
			std::cout << "成功处理图片数量: " << processedCount << std::endl;
			std::cout << "数据库记录数量: " << vecManager.getFeatureCount() << std::endl;

			if (processedCount == 0)
			{
				std::cerr << "没有成功提取任何图片特征！" << std::endl;
				return -1;
			}

			// 保存数据库
			vecManager.saveDatabase();
			std::cout << "数据库保存完成" << std::endl;
		}

		// 提取查询图片特征
		std::cout << "正在提取查询图片特征: " << searchImagePath << std::endl;
		std::vector<float> queryFeatures = extractor.extractFeatures(searchImagePath);
		if (queryFeatures.empty())
		{
			std::cerr << "无法提取查询图片特征！" << std::endl;
			return -1;
		}
		std::cout << "查询图片特征提取完成，维度: " << queryFeatures.size() << std::endl;

		// 进行相似性搜索
		std::cout << "正在进行相似性搜索..." << std::endl;
		std::vector<std::pair<std::string, float>> searchResults = vecManager.searchSimilarVectors(queryFeatures, 5);

		// 显示搜索结果
		std::cout << "\n=== 搜索结果 ===" << std::endl;
		std::cout << "查询图片: " << searchImagePath << std::endl;
		std::cout << "最相似的图片:" << std::endl;

		for (size_t i = 0; i < searchResults.size(); ++i)
		{
			std::cout << (i + 1) << ". " << searchResults[i].first
				<< " (相似度: " << searchResults[i].second << ")" << std::endl;
		}

		std::cout << "\n图像检索完成！" << std::endl;

	}
	catch (const std::exception & e)
	{
		std::string a = e.what();
		std::cerr << "程序执行出错: " << e.what() << std::endl;
		return -1;
	}

	sqlite3_sleep(100000);

	return 0;
}