#include <wrl/client.h>
#include <winml.h>
#include <stdio.h>
using Microsoft::WRL::ComPtr;

#include"cnpy.h"

#define PRINTDBG

int main(int argc,char**argv)
{
	//load model

	ComPtr<IWinMLRuntime> spRuntime;
	HRESULT res= WinMLCreateRuntime(&spRuntime);
	printf("%d\n", res);
	ComPtr<IWinMLModel> spModel;
	res=spRuntime->LoadModel(L"C:\\src4\\winml\\mnist.onnx", &spModel);
	printf("%d\n", res);

	//view graph metadata

	WINML_MODEL_DESC *desc;
	res = spModel->GetDescription(&desc);
	printf("%d\n", res);
	int count = 0;
	LPCWSTR key,value;
	res=spModel->EnumerateMetadata(count, &key, &value);
	printf("%d\n", res);
	WINML_VARIABLE_DESC *desc3;
	res = spModel->EnumerateModelInputs(count, &desc3);
	printf("%d\n", res);
	printf("%S %d %d %d %d %d\n", desc3->Name,desc3->Tensor.NumDimensions, desc3->Tensor.pShape[0], desc3->Tensor.pShape[1], desc3->Tensor.pShape[2], desc3->Tensor.pShape[3]);
	WINML_VARIABLE_DESC *desc2;
	res = spModel->EnumerateModelOutputs(count, &desc2);
	printf("%S %d %d %d %d %d\n", desc2->Name,desc2->Tensor.NumDimensions, desc2->Tensor.pShape[0], desc2->Tensor.pShape[1], desc2->Tensor.pShape[2], desc2->Tensor.pShape[3]);
	printf("%d\n", res);

	//bind resources

	//setup device to use for inferencing
	ComPtr<IWinMLEvaluationContext> spContext;
	ComPtr<ID3D12Device> spDevice;
	res=spRuntime->CreateEvaluationContext(spDevice.Get(), &spContext);
	printf("%d\n", res);

	//connect I/O data
	WINML_BINDING_DESC bindDescriptor;
	bindDescriptor.BindType = WINML_BINDING_TYPE::WINML_BINDING_TENSOR;
	bindDescriptor.Tensor.DataType=WINML_TENSOR_DATA_TYPE::WINML_TENSOR_FLOAT;
	bindDescriptor.Tensor.NumDimensions = 4;
	INT64 shape[4] = { 1,1,28,28 };
	bindDescriptor.Tensor.pShape = reinterpret_cast<INT64*>(&shape);
	//void *data=(void *)malloc(28 * 28 * sizeof(float));
	bindDescriptor.Tensor.DataSize = 28 * 28 * sizeof(float);

	char datafile[100] = { 0 };
	int nt = 0;
	if (argc == 2)
		nt = atoi(argv[1]);
	sprintf_s(datafile,99,"C:\\src4\\winml\\test_data_%d.npz", nt);
	//"C:\\src4\\winml\\test_data_0.npz"
	cnpy::NpyArray arr2 = cnpy::npz_load(datafile, "inputs");
	int comp = 0;
	comp = arr2.word_size == sizeof(float);
	comp = arr2.shape.size() == 3;
	comp = arr2.shape[0] == 1;
	comp = arr2.shape[1] == 28;
	comp = arr2.shape[2] == 28;
	float* mv1 = arr2.data<float>();
	for (int i = 0; i < 28; i++)
	{
		for (int j = 0; j < 28; j++)
		{
			char c = ' ';
			if (mv1[i * 28 + j] > 0.10)
				c = '*';
			else
				c = ' ';
			printf("%c", c);
			//printf("%.1f ", mv1[i*28+j]);
		}
		printf("\n");
	}

	bindDescriptor.Tensor.pData = /*data*/mv1;
	bindDescriptor.Name = desc3->Name;// LPCWSTR("hola");
	res=spContext->BindValue(&bindDescriptor);
	printf("%d\n", res);

	WINML_BINDING_DESC bindDescriptor2;
	bindDescriptor2.BindType = WINML_BINDING_TYPE::WINML_BINDING_TENSOR;
	bindDescriptor2.Tensor.DataType = WINML_TENSOR_DATA_TYPE::WINML_TENSOR_FLOAT;
	bindDescriptor2.Tensor.NumDimensions = 2;
	INT64 shape2[2] = { 1,10 };
	bindDescriptor2.Tensor.pShape = reinterpret_cast<INT64*>(&shape2);
	float *data2 = (float *)calloc(1 * 10 , sizeof(float));

#ifdef PRINTDBG
	for (int i = 0; i < 10; i++)
	{
		printf("%f ", data2[i]);
	}
	printf("\n");
#endif
	bindDescriptor2.Tensor.DataSize = 1 * 10 * sizeof(float);
	bindDescriptor2.Tensor.pData = data2;
	bindDescriptor2.Name = desc2->Name;// LPCWSTR("hola");
	res = spContext->BindValue(&bindDescriptor2);
	printf("%d\n", res);



	cnpy::NpyArray arr3 = cnpy::npz_load(datafile, "outputs");
	comp = 0;
	comp = arr3.word_size == sizeof(float);
	comp = arr3.shape.size() == 2;
	comp = arr3.shape[0] == 1;
	comp = arr3.shape[1] == 10;

	float* mv2 = arr3.data<float>();
#ifdef PRINTDBG
	for (int i = 0; i < 10; i++)
	{
			printf("%.2f ", mv2[i]);
	}
	printf("\n");
#endif

	//Evaluate model (inference)
	res=spRuntime->EvaluateModel(spContext.Get());

	//process results
	for (int i = 0; i < 10; i++)
	{
		printf("%.2f ", data2[i]);
	}
	printf("\n");
	int maxIndex = 0;
	float maxProbability = 0.0;
	for (int i = 0; i < 10 ; i++)
	{
		if (data2[i] > maxProbability)
		{
			maxIndex = i;
			maxProbability = data2[i];
		}
	}
	printf("numero detectado %d\n", maxIndex);
}