// nnbuildmat.cpp
//

#include "precompiled.h"
#include "../../../renderer/tr_local.h"

/*
=============
NeuralPOMProgressCallback
=============
*/
static void __cdecl NeuralPOMProgressCallback(
	const char* message,
	uint32_t currentSample,
	uint32_t totalSamples,
	float meanLoss,
	void* userData) {
	if (message && message[0]) {
		common->Printf("%s: %u / %u loss %f\n", message, currentSample, totalSamples, meanLoss);
	}
	else {
		common->Printf("Neural Material training: %u / %u loss %f\n", currentSample, totalSamples, meanLoss);
	}
}

/*
=============
ExportNatMaterial
=============
*/
void ExportNatMaterial(const idMaterial* mtr, const char* materialName) {
	globalImages->keepPixelsResident = true;

	idImage* albedoImage = mtr->GetDiffuseImage(NULL);
	idImage* normalImage = mtr->GetBumpImage();
	idImage* specImage = mtr->GetSpecImage();

	QD3D12NeuralPOMImageRGBA8 albedoImg = {};
	albedoImg.width = albedoImage->uploadWidth;
	albedoImg.height = albedoImage->uploadHeight;
	albedoImg.rowPitchBytes = albedoImage->uploadWidth * 4;
	albedoImg.pixelsRGBA8 = albedoImage->residentPixels;

	QD3D12NeuralPOMImageRGBA8 normalImg = {};
	normalImg.width = normalImage->uploadWidth;
	normalImg.height = normalImage->uploadHeight;
	normalImg.rowPitchBytes = normalImage->uploadWidth * 4;
	normalImg.pixelsRGBA8 = normalImage->residentPixels;

	QD3D12NeuralPOMImageRGBA8 specularImg = {};
	specularImg.width = specImage->uploadWidth;
	specularImg.height = specImage->uploadHeight;
	specularImg.rowPitchBytes = specImage->uploadWidth * 4;
	specularImg.pixelsRGBA8 = specImage->residentPixels;

	idStr materialNameFixed = materialName;
	materialNameFixed.Replace("/", "__");

	QD3D12NeuralPOMTrainDesc desc = {};
	desc.size = sizeof(desc);
	desc.albedoImage = &albedoImg;
	desc.normalImage = &normalImg;
	desc.specularImage = &specularImg;
	strcpy(desc.outputPrefix, va("base/neural/%s", materialNameFixed.c_str()));
	desc.samples = 1600000;
	desc.samplesPerGpuBatch = 4192;
	desc.latentResolution = 256;
	desc.latentChannels = 8;
	desc.hiddenCount = 48;
	desc.normalStrength = 2.0f;
	desc.normalMapYSign = 1.0f;
	desc.flags = QD3D12_NEURAL_POM_TRAIN_FLAG_VERBOSE;
	desc.progress = NeuralPOMProgressCallback;
	desc.progressUserData = NULL;

	QD3D12NeuralPOMTrainStats stats = {};
	stats.size = sizeof(stats);

	int ok = QD3D12_NeuralPOMTrainMaterialD3D12(&desc, &stats);
	if (ok) {
		common->Printf("Neural Material %s exported successfully\n", materialName);
	}
	else {
		common->Printf("There was a problem exporting %s\n", materialName);
	}

	globalImages->keepPixelsResident = false;
}

/*
============
nnbuildmat_f
============
*/
void nnbuildmat_f(const idCmdArgs& args) {
	if (args.Argc() < 1)
	{
		common->Warning("Usage: nnbuildmat <material>\n");
	}

	const idMaterial* mtr = declManager->FindMaterial(args.Argv(1), false);

	if (mtr == NULL) {
		common->Warning("Failed to find material %s\n", args.Argv(1));
		return;
	}

	ExportNatMaterial(mtr, args.Argv(1));
}