nmWTAI-Platform/ML/Training/Training/DatasetIO.h

#pragma once
#ifndef DATASET_IO_H
#define DATASET_IO_H

#include <string>
#include <vector>
#include <fstream>
#include <iostream>

#include "pch.h"
#include "SceneIO.h"   // 为了用 PebiScene

// DatasetIO：把 (scene + gridOutput2) 打包成一个 dataset.bin，供后续直接加载
class DatasetIO
{
public:
	// 保存 dataset（包含 scene + gridOutput2 + header(sceneKey)）
	static bool saveDataset(const std::string& filename,
		const PebiScene& scene,
		const HX_NWTM_GRID_INPUT& gridInput,
		const HX_NWTM_GRID_OUTPUT2& grid);

	// 加载 dataset（读出 scene + gridOutput2，且校验 header）
	// strictVerify=true：会用读出来的 scene 重新计算 sceneKey 与 header 对比（防止文件损坏/写错）
	static bool loadDataset(const std::string& filename,
		PebiScene& outScene,
		HX_NWTM_GRID_OUTPUT2& outGrid,
		bool strictVerify = true);

	// 只读 header 并返回 sceneKey（用于与外部 scene.bin 比对串场景）
	static bool readDatasetSceneKey(const std::string& filename,
		unsigned long long& outSceneKey,
		unsigned int& outCells,
		unsigned int& outWells,
		int& outSolverType);

	static bool fileExists(const std::string& filename);

	// 计算 sceneKey：与 GridCacheIO 相同算法（FNV-1a 64）
	static unsigned long long computeSceneKey64(const HX_NWTM_GRID_INPUT& in);

private:
	// ---------------- 文件头 ----------------
	struct DatasetHeader
	{
		unsigned int magic;        // 'DTHX' 等
		unsigned int version;      // 格式版本
		unsigned int reserved;     // 预留
		unsigned long long sceneKey; // 串场景指纹
		unsigned int nCells;       // cell 数
		unsigned int nWells;       // 井数（从 Rate.t.size() 或 scene.wellName.size() 推）
		int solverType;            // 求解器类型
	};

	enum { MAGIC = 0x58485444 }; // 'DTHX' little-endian
	enum { VERSION = 2 };

	// 从 scene 生成 gridInput（用于 strictVerify 重新算 sceneKey）
	static HX_NWTM_GRID_INPUT buildGridInputFromScene(const PebiScene& scene);

	// ---------------- 基础 IO（安全尺寸） ----------------
	static bool writeUInt(std::ofstream& fs, unsigned int v);
	static bool writeULL(std::ofstream& fs, unsigned long long v);
	static bool writeInt(std::ofstream& fs, int v);
	static bool writeDouble(std::ofstream& fs, double v);

	static bool readUInt(std::ifstream& fs, unsigned int& v);
	static bool readULL(std::ifstream& fs, unsigned long long& v);
	static bool readInt(std::ifstream& fs, int& v);
	static bool readDouble(std::ifstream& fs, double& v);

	static bool writeSafeSize(std::ofstream& fs, size_t sz);
	static bool readSafeSize(std::ifstream& fs, size_t& sz);

	// ---------------- string/vector/scene IO ----------------
	static bool writeString(std::ofstream& fs, const std::string& s);
	static bool readString(std::ifstream& fs, std::string& s);

	static bool writeStdVecD(std::ofstream& fs, const std::vector<double>& v);
	static bool readStdVecD(std::ifstream& fs, std::vector<double>& v);

	static bool writeStdVecI(std::ofstream& fs, const std::vector<int>& v);
	static bool readStdVecI(std::ifstream& fs, std::vector<int>& v);

	static bool writeStdVec2D(std::ofstream& fs, const std::vector<std::vector<double> >& v);
	static bool readStdVec2D(std::ifstream& fs, std::vector<std::vector<double> >& v);

	static bool writeStdVec3D(std::ofstream& fs, const std::vector<std::vector<std::vector<double> > >& v);
	static bool readStdVec3D(std::ifstream& fs, std::vector<std::vector<std::vector<double> > >& v);

	static bool writeScene(std::ofstream& fs, const PebiScene& scene);
	static bool readScene(std::ifstream& fs, PebiScene& scene);

	// ---------------- gridOutput2 IO ----------------
	static bool writeVec1D(std::ofstream& fs, const dVec1& v);
	static bool writeIVec1D(std::ofstream& fs, const iVec1& v);
	static bool writeVec2D(std::ofstream& fs, const dVec2& v);
	static bool writeIVec2D(std::ofstream& fs, const iVec2& v);
	static bool writeVec3D(std::ofstream& fs, const dVec3& v);

	static bool readVec1D(std::ifstream& fs, dVec1& v);
	static bool readIVec1D(std::ifstream& fs, iVec1& v);
	static bool readVec2D(std::ifstream& fs, dVec2& v);
	static bool readIVec2D(std::ifstream& fs, iVec2& v);
	static bool readVec3D(std::ifstream& fs, dVec3& v);

	static bool writeGridOutput2(std::ofstream& fs, const HX_NWTM_GRID_OUTPUT2& g);
	static bool readGridOutput2(std::ifstream& fs, HX_NWTM_GRID_OUTPUT2& g);

	static bool writeHeader(std::ofstream& fs, const DatasetHeader& hdr);
	static bool readHeader(std::ifstream& fs, DatasetHeader& hdr);

	// ---------------- FNV-1a 64 hash ----------------
	static unsigned long long fnv1a64_init();
	static void fnv1a64_update(unsigned long long& h, const void* data, size_t len);

	static void hashInt(unsigned long long& h, int v);
	static void hashUInt(unsigned long long& h, unsigned int v);
	static void hashDouble(unsigned long long& h, double v);
	static void hashVec1D(unsigned long long& h, const dVec1& v);
	static void hashIVec1D(unsigned long long& h, const iVec1& v);
	static void hashVec2D(unsigned long long& h, const dVec2& v);
	static void hashIVec2D(unsigned long long& h, const iVec2& v);
	static void hashVec3D(unsigned long long& h, const dVec3& v);
};

#endif // DATASET_IO_H