nmWTAI-Platform/ML/Training/Training/Training.cpp

#include <Windows.h>
#include <iostream>
#include <fstream>
#include <string>

#include "pch.h"
#include "SceneIO.h"
#include "GridCacheIO.h"
#include "DatasetIO.h"

typedef void (*HX_NWTM_GRID_Func)(HX_NWTM_GRID_OUTPUT1&, HX_NWTM_GRID_OUTPUT2&,
	const HX_NWTM_GRID_INPUT&, std::string);
typedef void (*HX_NWTM_MODEL_Func)(HX_NWTM_MODEL_OUTPUT&, const HX_NWTM_MODEL_INPUT&, std::string);

// 获取目录
static std::string getExeDir()
{
	char buf[MAX_PATH] = {0};
	DWORD len = GetModuleFileNameA(NULL, buf, MAX_PATH);
	if (len == 0) {
		return ".";
	}
	std::string path(buf, len);
	size_t pos = path.find_last_of("\\/");
	if (pos != std::string::npos) {
		path.resize(pos);
	}
	return path;
}

static std::string getParentDir(const std::string& path)
{
	size_t pos = path.find_last_of("\\/");
	return (pos == std::string::npos ? std::string(".") : path.substr(0, pos));
}

static std::string getFileStem(const std::string& path)
{
	size_t slash = path.find_last_of("\\/");
	size_t begin = (slash == std::string::npos ? 0 : slash + 1);
	size_t dot = path.find_last_of('.');
	if (dot == std::string::npos || dot < begin) {
		dot = path.size();
	}
	return path.substr(begin, dot - begin);
}

static std::string joinPath(const std::string& dir, const std::string& filename)
{
	if (dir.empty() || dir == ".") {
		return filename;
	}
	char last = dir[dir.size() - 1];
	return dir + (last == '\\' || last == '/' ? "" : "\\") + filename;
}

static std::string deriveSceneSiblingPath(const std::string& scenePath, const std::string& prefix)
{
	std::string stem = getFileStem(scenePath);
	std::string suffix;
	if (stem.size() > 6 && stem.compare(0, 6, "scene_") == 0) {
		suffix = stem.substr(6);
	}
	std::string filename = prefix + (suffix.empty() ? "" : "_" + suffix) + ".bin";
	return joinPath(getParentDir(scenePath), filename);
}

// 从 scene 重建 GRID 输入（与之前一致）
static HX_NWTM_GRID_INPUT rebuildGridInput(const PebiScene& scene)
{
	HX_NWTM_GRID_INPUT input;
	input.D = scene.D;
	input.GridControl = scene.GridControl;
	input.Boundary = scene.Boundary;
	input.VerticalWell = scene.VerticalWell;
	input.HorizontalWell = scene.HorizontalWell;
	input.FractureVerticalWell = scene.FractureVerticalWell;
	input.MultistageFracturedHorizontalWell = scene.MultistageFracturedHorizontalWell;
	input.InclinedWell = scene.InclinedWell;
	input.Fault = scene.Fault;
	return input;
}

// 从 scene + gridOutput2 重建 MODEL 输入（与之前一致）
static HX_NWTM_MODEL_INPUT rebuildModelInput(const PebiScene& scene, const HX_NWTM_GRID_OUTPUT2& gridOutput)
{
	HX_NWTM_MODEL_INPUT input(gridOutput);

	input.T = scene.solverType;

	input.Rate.t = scene.Rate.t;
	input.Rate.qo = scene.Rate.qo;
	input.Rate.qg = scene.Rate.qg;
	input.Rate.qw = scene.Rate.qw;

	input.CS.C = scene.CS.C;
	input.CS.S = scene.CS.S;

	input.PVT.p = scene.PVT.p;
	input.PVT.pb = scene.PVT.pb;
	input.PVT.Rso = scene.PVT.Rso;
	input.PVT.Bo = scene.PVT.Bo;
	input.PVT.Co = scene.PVT.Co;
	input.PVT.miuo = scene.PVT.miuo;
	input.PVT.rouo = scene.PVT.rouo;
	input.PVT.Rv = scene.PVT.Rv;
	input.PVT.Bg = scene.PVT.Bg;
	input.PVT.Cg = scene.PVT.Cg;
	input.PVT.miug = scene.PVT.miug;
	input.PVT.roug = scene.PVT.roug;
	input.PVT.Z = scene.PVT.Z;
	input.PVT.Rsw = scene.PVT.Rsw;
	input.PVT.Bw = scene.PVT.Bw;
	input.PVT.Cw = scene.PVT.Cw;
	input.PVT.miuw = scene.PVT.miuw;
	input.PVT.rouw = scene.PVT.rouw;
	input.PVT.V = scene.PVT.V;
	input.PVT.k_kinitial = scene.PVT.k_kinitial;
	input.PVT.Cf_Cfinitial = scene.PVT.Cf_Cfinitial;
	input.PVT.So = scene.PVT.So;
	input.PVT.Kro = scene.PVT.Kro;
	input.PVT.Sg = scene.PVT.Sg;
	input.PVT.Krg = scene.PVT.Krg;
	input.PVT.Sw = scene.PVT.Sw;
	input.PVT.Krw = scene.PVT.Krw;

	input.Base.Pi = scene.Base.Pi;
	input.Base.Cti = scene.Base.Cti;
	input.Base.Cf = scene.Base.Cf;
	input.Base.Soi = scene.Base.Soi;
	input.Base.Sgi = scene.Base.Sgi;
	input.Base.Swi = scene.Base.Swi;
	input.Base.d = scene.Base.d;
	input.Base.dt_Min = scene.Base.dt_Min;
	input.Base.dt_Max = scene.Base.dt_Max;

	// 用参考值填满每个网格单元（最小可用）
	size_t nCells = gridOutput.Trinodexy.size();
	input.Base.k = dVec1(nCells, scene.Base.k_ref);
	input.Base.phi = dVec1(nCells, scene.Base.phi_ref);
	input.Base.h = dVec1(nCells, scene.Base.h_ref);

	return input;
}

// 导出验证样本（CSV）
static void exportValidationSample(const HX_NWTM_MODEL_OUTPUT& output, const std::string& filename)
{
	std::ofstream f(filename.c_str());
	f << "t,pw\n";
	if (!output.pw.empty()) {
		for (size_t i = 0; i < output.t.size() && i < output.pw[0].size(); ++i) {
			f << output.t[i] << "," << output.pw[0][i] << "\n";
		}
	}
	f.close();
	std::cout << "已导出: " << filename << std::endl;
}

// 打印输入摘要：用于快速确认“为什么 steps 变了”
static void printInputSummary(const PebiScene& scene)
{
	size_t wells = scene.Rate.t.size();
	size_t tlen0 = (wells > 0 ? scene.Rate.t[0].size() : 0);

	std::cout << "输入摘要：wells=" << (unsigned int)wells
		<< ", solver=" << scene.solverType
		<< ", Rate.t[0].len=" << (unsigned int)tlen0
		<< ", PVT点=" << (unsigned int)scene.PVT.p.size()
		<< std::endl;
}

int main(int argc, char* argv[])
{
	std::cout << "=== PEBI Scene Processor (Dataset First) ===" << std::endl;
	if (argc > 4) {
		std::cerr << "Usage: training.exe [scene.bin] [dataset.bin] [grid_cache.bin]" << std::endl;
		return 1;
	}

	// 按需开关：发布版保持 true（因为 dataset 是后续批量采样的核心）
	const bool ENABLE_DATASET_FIRST = true;
	const bool ENABLE_DATASET_WRITE = true; // 当走 scene+grid 生成时，是否写 dataset.bin
	 
	// 路径
	std::string exeDir = getExeDir();

	// license 仍然在工程根目录下的 Bin\Res\license
	std::string lic = exeDir + "\\..\\..\\..\\Bin\\Res\\license\\HXNWTM_license.dat";

	// 所有数据统一放到 ML/nmWTAI-ML/data/temp
	//std::string dataDir      = exeDir + "\\..\\..\\Data";
	//std::string scenePath    = dataDir + "\\scene.bin";
	//std::string datasetPath  = dataDir + "\\dataset.bin";
	//std::string gridCachePath= dataDir + "\\grid_cache.bin";
	std::string dataDir      = exeDir + "\\..\\..\\nmWTAI-ML\\data\\temp";
	std::string scenePath = (argc >= 2 ? argv[1] : dataDir + "\\scene.bin");
	std::string datasetPath = (argc >= 3 ? argv[2] : deriveSceneSiblingPath(scenePath, "dataset"));
	std::string gridCachePath = (argc >= 4 ? argv[3] : deriveSceneSiblingPath(scenePath, "grid_cache"));

	std::cout << "Scene: " << scenePath << std::endl;
	std::cout << "Dataset: " << datasetPath << std::endl;
	std::cout << "Grid cache: " << gridCachePath << std::endl;

	// 1) 加载 DLL
	HMODULE hx = LoadLibraryW(L"HX_NWTM.dll");
	if (!hx) {
		std::cerr << "LoadLibrary failed, err=" << GetLastError() << std::endl;
		return 1;
	}

	HX_NWTM_GRID_Func  HX_NWTM_GRID  = (HX_NWTM_GRID_Func)GetProcAddress(hx, "HX_NWTM_GRID");
	HX_NWTM_MODEL_Func HX_NWTM_MODEL = (HX_NWTM_MODEL_Func)GetProcAddress(hx, "HX_NWTM_MODEL");
	if (!HX_NWTM_GRID || !HX_NWTM_MODEL) {
		std::cerr << "GetProcAddress failed, err=" << GetLastError() << std::endl;
		FreeLibrary(hx);
		return 2;
	}

	// 2) 检查 license
	DWORD attr = GetFileAttributesA(lic.c_str());
	if (attr == INVALID_FILE_ATTRIBUTES) {
		std::cerr << "License not found: " << lic << std::endl;
		FreeLibrary(hx);
		return 3;
	}
	std::cout << "License OK: " << lic << std::endl;

	// 3) dataset 优先尝试（但如果对应 scene 存在，会做 sceneKey 串场景校验）
	PebiScene scene;
	HX_NWTM_GRID_OUTPUT2 gridOutput2;
	bool gotFromDataset = false;

	if (ENABLE_DATASET_FIRST && DatasetIO::fileExists(datasetPath)) {
		std::cout << "\n发现 dataset: " << datasetPath << "，尝试直接加载..." << std::endl;

		// 先读 dataset header
		unsigned long long dsKey = 0;
		unsigned int dsCells = 0;
		unsigned int dsWells = 0;
		int dsSolver = 0;

		bool headerOK = DatasetIO::readDatasetSceneKey(datasetPath, dsKey, dsCells, dsWells, dsSolver);

		// 如果对应 scene 存在，我们就顺便做一次“串场景保护”：sceneKey 不匹配就不用 dataset
		bool sceneKeyMatch = true;

		if (headerOK && DatasetIO::fileExists(scenePath)) {
			std::cout << "\n加载 scene（仅用于串场景校验）: " << scenePath << std::endl;
			if (SceneIO::loadScene(scenePath, scene)) {
				HX_NWTM_GRID_INPUT gi = rebuildGridInput(scene);
				unsigned long long expected = DatasetIO::computeSceneKey64(gi);
				if (expected != dsKey) {
					sceneKeyMatch = false;
					std::cout << "dataset sceneKey 不匹配（串场景保护触发），忽略 dataset，准备重建..." << std::endl;
				}
			} else {
				std::cout << "scene 加载失败，跳过串场景校验，直接尝试 dataset 自检..." << std::endl;
				sceneKeyMatch = true;
			}
		}

		if (sceneKeyMatch) {
			// 真正加载 dataset（strictVerify=true 会用 dataset 内部 scene 重新算 key）
			PebiScene dsScene;
			HX_NWTM_GRID_OUTPUT2 dsGrid;
			if (DatasetIO::loadDataset(datasetPath, dsScene, dsGrid, true)) {
				scene = dsScene;
				gridOutput2 = dsGrid;
				gotFromDataset = true;

				std::cout << "dataset 加载成功：cells=" << (unsigned int)gridOutput2.Trinodexy.size()
					<< ", wells=" << (unsigned int)(scene.Rate.t.size())
					<< ", solver=" << scene.solverType << std::endl;
			} else {
				std::cout << "dataset 加载失败，准备走 scene+grid_cache..." << std::endl;
			}
		}
	}

	// 4) 如果没有拿到 dataset，就走 scene + grid_cache（带串场景保护）
	if (!gotFromDataset) {
		std::cout << "\n加载 scene: " << scenePath << std::endl;
		if (!SceneIO::loadScene(scenePath, scene)) {
			std::cerr << "scene 读取失败: " << scenePath << std::endl;
			FreeLibrary(hx);
			return 4;
		}

		// 输入摘要：帮助你定位 “为什么 steps 变了”
		std::cout << "scene 读取成功：Wells=" << (unsigned int)scene.wellName.size()
			<< ", PVT点=" << (unsigned int)scene.PVT.p.size()
			<< ", Solver=" << scene.solverType << std::endl;
		printInputSummary(scene);

		HX_NWTM_GRID_INPUT gridInput = rebuildGridInput(scene);

		HX_NWTM_GRID_OUTPUT1 gridOutput1;
		bool loadedFromCache = false;

		if (GridCacheIO::fileExists(gridCachePath)) {
			std::cout << "\n读取网格缓存: " << gridCachePath << std::endl;
			if (GridCacheIO::loadGrid(gridCachePath, gridOutput2, gridInput)) {
				loadedFromCache = true;
			} else {
				std::cout << "缓存无效或串场景不匹配，准备重新生成..." << std::endl;
			}
		}

		if (!loadedFromCache) {
			std::cout << "\n生成网格..." << std::endl;
			try {
				HX_NWTM_GRID(gridOutput1, gridOutput2, gridInput, lic);
			} catch (...) {
				std::cerr << "HX_NWTM_GRID 异常" << std::endl;
				FreeLibrary(hx);
				return 5;
			}

			if (gridOutput2.Trinodexy.empty()) {
				std::cerr << "HX_NWTM_GRID 返回空网格，停止写入缓存和 dataset" << std::endl;
				FreeLibrary(hx);
				return 5;
			}

			std::cout << "保存网格缓存: " << gridCachePath << std::endl;
			if (!GridCacheIO::saveGrid(gridCachePath, gridOutput2, gridInput)) {
				std::cout << "警告：保存网格缓存失败（继续执行）" << std::endl;
			}
		}

		std::cout << "网格就绪: " << (unsigned int)gridOutput2.Trinodexy.size()
			<< " cells" << (loadedFromCache ? " (from cache)" : " (generated)") << std::endl;

		// 生成/更新 dataset.bin（可开关）
		if (ENABLE_DATASET_WRITE) {
			std::cout << "\n写入 dataset: " << datasetPath << std::endl;
			if (!DatasetIO::saveDataset(datasetPath, scene, gridInput, gridOutput2)) {
				std::cout << "警告：dataset 写入失败（继续执行）" << std::endl;
			} else {
				std::cout << "dataset 写入成功" << std::endl;
			}
		}
	}

	//// 5) 求解模型（这一步不是“多余验证”，而是 dataset 生成/采样最终都要用到的核心步骤）
	//std::cout << "\n开始求解模型..." << std::endl;

	//HX_NWTM_MODEL_INPUT modelInput = rebuildModelInput(scene, gridOutput2);
	//HX_NWTM_MODEL_OUTPUT modelOutput;

	//try {
	//	HX_NWTM_MODEL(modelOutput, modelInput, lic);
	//} catch (...) {
	//	std::cerr << "HX_NWTM_MODEL 异常" << std::endl;
	//	FreeLibrary(hx);
	//	return 6;
	//}

	//std::cout << "===================================\n";
	//std::cout << "求解完成: t=" << (unsigned int)modelOutput.t.size()
	//	<< " steps, wells=" << (unsigned int)modelOutput.pw.size() << std::endl;

	//exportValidationSample(modelOutput, dataDir + "\\validation.csv");

	FreeLibrary(hx);
	std::cout << "\nDone!" << std::endl;
	return 0;
}