#include <iostream>
#include <vector>
#include <string>
#include <chrono>
#include <random>
#include <opencv2/opencv.hpp>
#include <onnxruntime_cxx_api.h>

using namespace std;
using namespace cv;
using namespace Ort;

// 定义检测结果结构体，用于在函数间传递数据
struct Detection {
	Rect box;       // 边界框 (x, y, width, height)
	float conf;     // 置信度
	int classId;    // 类别ID
};

class YOLO26 {
private:
	// 模型相关路径和参数
	string onnx_model_path;
	float conf_thres;
	float iou_thres;

	// ONNX Runtime 资源
	Env env;
	Session session;
	vector<const char*> input_node_names;
	vector<const char*> output_node_names;
	vector<string> input_node_names_alloc;
	vector<string> output_node_names_alloc;

	// 模型输入尺寸
	int64_t input_width;
	int64_t input_height;

	// 类别名称和颜色
	vector<string> class_names;
	vector<Scalar> color_palette;

public:
	YOLO26(const string& model_path, float conf, float iou)
		: onnx_model_path(model_path), conf_thres(conf), iou_thres(iou),
		env(ORT_LOGGING_LEVEL_WARNING, "YOLO26"), session(nullptr)
	{
		init_resources();
		init_classes_and_colors();
	}

private:
	// 初始化 ONNX Runtime 资源
	void init_resources() {
		SessionOptions session_options;
		session_options.SetIntraOpNumThreads(1);
		session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);
		// Windows 系统下，ONNX Runtime 的路径必须是宽字符 (wstring)
#ifdef _WIN32
		size_t newsize = onnx_model_path.length() + 1;
		std::wstring model_path_w(newsize, L'\0');
		size_t convertedChars = 0;
		// 将 string 转换为 wstring
		mbstowcs_s(&convertedChars, &model_path_w[0], newsize, onnx_model_path.c_str(), _TRUNCATE);
		session = Session(env, model_path_w.c_str(), session_options);
#else
	// Linux/Mac 系统下，直接使用 char*
		session = Session(env, onnx_model_path.c_str(), session_options);
#endif

		AllocatorWithDefaultOptions allocator;

		// 获取输入节点
		size_t num_input_nodes = session.GetInputCount();
		for (size_t i = 0; i < num_input_nodes; i++) {
			auto input_name = session.GetInputNameAllocated(i, allocator);
			input_node_names_alloc.push_back(input_name.get());
			input_node_names.push_back(input_node_names_alloc.back().c_str());

			auto type_info = session.GetInputTypeInfo(i);
			auto shape = type_info.GetTensorTypeAndShapeInfo().GetShape();
			input_height = shape[2];
			input_width = shape[3];
		}

		// 获取输出节点
		size_t num_output_nodes = session.GetOutputCount();
		for (size_t i = 0; i < num_output_nodes; i++) {
			auto output_name = session.GetOutputNameAllocated(i, allocator);
			output_node_names_alloc.push_back(output_name.get());
			output_node_names.push_back(output_node_names_alloc.back().c_str());
		}
	}

	// 初始化类别和颜色
	void init_classes_and_colors() {
		class_names = {
			"person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck",
			"boat", "traffic light", "fire hydrant", "stop sign", "parking meter", "bench",
			"bird", "cat", "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra",
			"giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
			"skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove",
			"skateboard", "surfboard", "tennis racket", "bottle", "wine glass", "cup",
			"fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", "orange",
			"broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
			"potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse",
			"remote", "keyboard", "cell phone", "microwave", "oven", "toaster", "sink",
			"refrigerator", "book", "clock", "vase", "scissors", "teddy bear", "hair drier",
			"toothbrush"
		};

		mt19937 rng(42);
		uniform_int_distribution<int> uni(0, 255);
		for (size_t i = 0; i < class_names.size(); ++i) {
			color_palette.emplace_back(Scalar(uni(rng), uni(rng), uni(rng)));
		}
	}

	// 辅助函数：Letterbox 调整大小
	void letterbox(const Mat& img, Mat& out_img, vector<int>& pad, float& ratio) {
		int shape_h = img.rows;
		int shape_w = img.cols;
		int new_h = input_height;
		int new_w = input_width;

		float r = min((float)new_h / shape_h, (float)new_w / shape_w);
		ratio = r;

		int new_unpad_w = (int)round(shape_w * r);
		int new_unpad_h = (int)round(shape_h * r);
		int dw = (new_w - new_unpad_w) / 2;
		int dh = (new_h - new_unpad_h) / 2;

		if (shape_w != new_unpad_w || shape_h != new_unpad_h) {
			resize(img, out_img, Size(new_unpad_w, new_unpad_h));
		}
		else {
			out_img = img.clone();
		}

		int top = dh, bottom = new_h - new_unpad_h - dh;
		int left = dw, right = new_w - new_unpad_w - dw;
		pad = { top, left };

		copyMakeBorder(out_img, out_img, top, bottom, left, right, BORDER_CONSTANT, Scalar(114, 114, 114));
	}

public:
	// ---------------------------------------------------------
	// 1. 预处理 (Preprocessing)
	// ---------------------------------------------------------
	vector<float> preprocess(const Mat& img, vector<int>& pad, float& ratio) {
		Mat img_rgb;
		cvtColor(img, img_rgb, COLOR_BGR2RGB);

		Mat resized_img;
		letterbox(img_rgb, resized_img, pad, ratio);

		// 准备数据容器：CHW 格式
		vector<float> input_tensor_values(input_width * input_height * 3);

		// HWC -> CHW 并归一化 (0-1)
		for (int c = 0; c < 3; c++) {
			for (int h = 0; h < input_height; h++) {
				for (int w = 0; w < input_width; w++) {
					input_tensor_values[c * input_height * input_width + h * input_width + w] =
						resized_img.at<Vec3b>(h, w)[c] / 255.0f;
				}
			}
		}
		return input_tensor_values;
	}

	// ---------------------------------------------------------
	// 2. 推理 (Inference)
	// ---------------------------------------------------------
	vector<Value> inference(vector<float>& input_tensor_values) {
		// 创建内存信息
		auto memory_info = MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);

		// 定义输入形状 [1, 3, H, W]
		vector<int64_t> input_shape = { 1, 3, input_height, input_width };

		// 创建输入 Tensor
		Value input_tensor = Value::CreateTensor<float>(
			memory_info,
			input_tensor_values.data(),
			input_tensor_values.size(),
			input_shape.data(),
			input_shape.size()
			);

		// 运行推理
		return session.Run(RunOptions{ nullptr }, input_node_names.data(), &input_tensor, 1, output_node_names.data(), 1);
	}

	// ---------------------------------------------------------
	// 3. 后处理 (Postprocessing)
	// ---------------------------------------------------------
	vector<Detection> postprocess(vector<Value>& output_tensors, const vector<int>& pad, float ratio) {
		vector<Detection> detections;

		// 获取输出数据指针
		float* floatarr = output_tensors[0].GetTensorMutableData<float>();
		auto output_shape = output_tensors[0].GetTensorTypeAndShapeInfo().GetShape();

		// 假设输出形状为 [1, N, 6] 或 [N, 6]
		// 格式: [x1, y1, x2, y2, conf, class_id]
		int num_detections = 0;
		int dimensions = 0;

		if (output_shape.size() == 3) {
			num_detections = output_shape[1];
			dimensions = output_shape[2];
		}
		else if (output_shape.size() == 2) {
			num_detections = output_shape[0];
			dimensions = output_shape[1];
		}

		for (int i = 0; i < num_detections; ++i) {
			float* det = floatarr + i * dimensions;
			float conf = det[4];

			if (conf >= conf_thres) {
				float x1 = det[0];
				float y1 = det[1];
				float x2 = det[2];
				float y2 = det[3];
				int cls_id = (int)det[5];

				// 坐标还原：减去 padding 并除以缩放比例
				int r_x1 = int((x1 - pad[1]) / ratio);
				int r_y1 = int((y1 - pad[0]) / ratio);
				int r_x2 = int((x2 - pad[1]) / ratio);
				int r_y2 = int((y2 - pad[0]) / ratio);

				detections.push_back({
					Rect(r_x1, r_y1, r_x2 - r_x1, r_y2 - r_y1), // 转为 x,y,w,h
					conf,
					cls_id
					});
			}
		}
		return detections;
	}

	// ---------------------------------------------------------
	// 4. 可视化 (Visualization)
	// ---------------------------------------------------------
	Mat visualize(const Mat& img, const vector<Detection>& detections) {
		Mat result_img = img.clone();

		for (const auto& det : detections) {
			Scalar color = color_palette[det.classId % color_palette.size()];

			// 绘制矩形框
			rectangle(result_img, det.box, color, 2);

			// 绘制标签
			string label = class_names[det.classId] + ": " + to_string(det.conf).substr(0, 4);

			int baseLine;
			Size labelSize = getTextSize(label, FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseLine);
			int top = max(det.box.y, labelSize.height + 10);

			rectangle(result_img, Point(det.box.x, top - labelSize.height - 10),
				Point(det.box.x + labelSize.width, top + baseLine - 10), color, FILLED);

			putText(result_img, label, Point(det.box.x, top - 5), FONT_HERSHEY_SIMPLEX, 0.5, Scalar(0, 0, 0), 1);
		}
		return result_img;
	}
};

int main(int argc, char** argv) {
	string model_path = "yolo26s.onnx";
	string img_path = "bus.jpg";

	if (argc > 1) model_path = argv[1];
	if (argc > 2) img_path = argv[2];

	try {
		// 0. 初始化
		YOLO26 detector(model_path, 0.25f, 0.45f);

		// 读取图片
		Mat img = imread(img_path);
		if (img.empty()) {
			cerr << "Error: Could not open image." << endl;
			return -1;
		}

		cout << "Processing " << img_path << "..." << endl;

		// 运行 10 次测试
		for (int i = 0; i < 9; ++i) {
			auto start = chrono::high_resolution_clock::now();

			// 1. 预处理
			vector<int> pad;
			float ratio;
			vector<float> input_data = detector.preprocess(img, pad, ratio);

			// 2. 推理
			vector<Value> outputs = detector.inference(input_data);

			// 3. 后处理
			vector<Detection> results = detector.postprocess(outputs, pad, ratio);

			// 4. 可视化 (仅最后一次保存)
			if (i == 9) {
				Mat out_img = detector.visualize(img, results);
				imwrite("output.jpg", out_img);
				cout << "Saved output.jpg" << endl;
			}

			auto end = chrono::high_resolution_clock::now();
			chrono::duration<double> elapsed = end - start;
			cout << "Step " << i << " time: " << elapsed.count() << "s" << endl;
		}

	}
	catch (const exception& e) {
		cerr << "Exception: " << e.what() << endl;
		return -1;
	}

	return 0;
}