//#undef UpdateResource

#include "OptiXContextManager.h"
#include "OptiXModule.h"
#include "OptiXBuffer.h"

#include "RenderCore.h"
#include "EngineUtils.h"

#include <Runtime/Engine/Classes/Engine/Engine.h>
#include "Runtime/Engine/Public/SceneView.h"
#include "Runtime/Engine/Classes/Materials/MaterialInstanceDynamic.h"

#include "Runtime/Engine/Classes/Engine/TextureCube.h"
#include "Runtime/Engine/Public/TextureResource.h"
#include "Runtime/Engine/Classes/Engine/TextureRenderTarget2D.h"
#include "Runtime/Engine/Classes/Engine/TextureRenderTargetCube.h"

// VR
#include "Runtime/HeadMountedDisplay/Public/IHeadMountedDisplay.h"
#include "Runtime/HeadMountedDisplay/Public/IXRTrackingSystem.h"
#include "Runtime/HeadMountedDisplay/Public/IXRCamera.h"


#include "Runtime/Engine/Classes/GameFramework/GameUserSettings.h"

#include "Async.h"

// Console variables todo

static TAutoConsoleVariable<int32> CVarDisableTrace(
	TEXT("optix.DisableTrace"),
	0,
	TEXT("Defines if Optix should perform a constant trace.\n"),
	ECVF_Scalability | ECVF_RenderThreadSafe);

static TAutoConsoleVariable<int32> CVarDisableLaserTrace(
	TEXT("optix.DisableLaserTrace"),
	0,
	TEXT("Defines if Optix should perform a constant trace.\n"),
	ECVF_Scalability | ECVF_RenderThreadSafe);



FOptiXContextManager::FOptiXContextManager(const FAutoRegister& AutoRegister)
	: FSceneViewExtensionBase(AutoRegister)
{
	UE_LOG(LogTemp, Display, TEXT("FOptiXContextManager, is in rendering thread: %i"), static_cast<int32>(IsInRenderingThread()));

	RTXOn = 0;

	LaserMaxDepth = 20;
	LaserEntryPoint = 1; // Default, will be overwritten anyway

	LaserBufferWidth = 50 * 50;
	LaserBufferHeight = LaserMaxDepth * 2;

	LaserBufferSize = LaserBufferHeight * LaserBufferWidth;

	bValidCubemap.AtomicSet(false);

	OnSceneChangedDelegate.AddRaw(this, &FOptiXContextManager::SceneChangedCallback);

}


void FOptiXContextManager::SetupViewFamily(FSceneViewFamily & InViewFamily)
{
	//UE_LOG(LogTemp, Warning, TEXT("SetupViewFamily"));
}

void FOptiXContextManager::SetupView(FSceneViewFamily & InViewFamily, FSceneView & InView)
{
	//UE_LOG(LogTemp, Warning, TEXT("SetupView"));

	// TODO Check Width/Height
}

void FOptiXContextManager::BeginRenderViewFamily(FSceneViewFamily & InViewFamily)
{
	//UE_LOG(LogTemp, Warning, TEXT("BeginRenderViewFamily"));
}

// Called on render thread at the start of rendering, for each view, after PreRenderViewFamily_RenderThread call.
void FOptiXContextManager::PreRenderView_RenderThread(FRHICommandListImmediate & RHICmdList, FSceneView & InView)
{
}

// Called on render thread at the start of rendering.
void FOptiXContextManager::PreRenderViewFamily_RenderThread(FRHICommandListImmediate & RHICmdList, FSceneViewFamily & InViewFamily)
{
	//UE_LOG(LogTemp, Warning, TEXT("PreRenderViewFamily_RenderThread"));
	if (!bIsInitialized && bStartTracing)
	{
		InitCUDADX();
	}
}

void FOptiXContextManager::PostRenderViewFamily_RenderThread(FRHICommandListImmediate & RHICmdList, FSceneViewFamily & InViewFamily)
{
	//UE_LOG(LogTemp, Warning, TEXT("PostRenderViewFamily_RenderThread"));
		// Laser Test part:

}

void FOptiXContextManager::PostRenderView_RenderThread(FRHICommandListImmediate & RHICmdList, FSceneView & InView)
{
	//UE_LOG(LogTemp, Warning, TEXT("VPM: PostRenderViewFamily_RenderThread %s"), *InView.ViewMatrices.GetViewProjectionMatrix().ToString());

	if (!bIsInitialized && !bClearToLaunch && !OptiXContext.IsValid() && !bStartTracing)
	{
		return;
	}

	// Init the yet uninitialized optix components - this queue should be empty and do nothing if no new components are registered.
	InitOptiXComponents(RHICmdList);

	// Update the remaining variables TODO this needs to be only done once not once per eye!
	OptiXContext->UpdateVariables();
	UpdateOptiXComponentVariables();
	UpdateRequestedCubemaps(RHICmdList);

	RemovePendingChildrenFromGroups();
	// Clean up any dangling optix objects here to not interfere with launch
	DestroyOptiXObjects();


	OptiXContext->SetMatrix("invViewProjection", InView.ViewMatrices.GetInvViewProjectionMatrix());
	OptiXContext->SetMatrix("viewProjection", InView.ViewMatrices.GetViewProjectionMatrix());

	FIntPoint Size = OptiXContext->GetBuffer("result_color")->GetSize2D();

	//TextureRegion.Height = Size.Y;
	//TextureRegion.Width = Size.X;
	//TextureRegion.SrcX = 0;
	//TextureRegion.SrcY = 0;
	//TextureRegion.DestX = 0;
	//TextureRegion.DestY = 0;

	double start = FPlatformTime::Seconds();

	// Update texture refs?
	//OutputTextureColorRightRef = ((FTexture2DResource*)OutputTexture->Resource)->GetTexture2DRHI();
	//OutputTextureDepthRightRef = ((FTexture2DResource*)DepthTexture->Resource)->GetTexture2DRHI();
	//OutputTextureDepthLeftRef = ((FTexture2DResource*)DepthTexture2->Resource)->GetTexture2DRHI();
	//OutputTextureColorLeftRef = ((FTexture2DResource*)OutputTexture2->Resource)->GetTexture2DRHI();

	//UE_LOG(LogTemp, Warning, TEXT("Res: %s"), *Size.ToString());

	bIsTracing.AtomicSet(true);
	OptiXContext->Launch(0, Size.X, Size.Y);
	bIsTracing.AtomicSet(false);

	double end = FPlatformTime::Seconds();
	//UE_LOG(LogTemp, Warning, TEXT("Launch took %f seconds"), end - start);

	//UE_LOG(LogTemp, Warning, TEXT("Res : %i %i"), Width, Height);

	start = FPlatformTime::Seconds();

	if (InView.StereoPass == EStereoscopicPass::eSSP_LEFT_EYE) // check validity
	{		
		//float* Data2 = static_cast<float*>(OptiXContext->GetBuffer("result_depth")->MapNative());
		//RHICmdList.UpdateTexture2D(OutputTextureDepthLeftRef, 0, TextureRegion, Size.X * 4, (uint8*)Data2);
		//OptiXContext->GetBuffer("result_depth")->Unmap();

		if (Resources[0] == NULL && Resources[1] == NULL)
		{
			return;
		}

		cudaGraphicsMapResources(2, Resources, 0);
		PrintLastCudaError("cudaGraphicsMapResources");

		if (CudaResourceDepthLeft == NULL)
		{
			cudaGraphicsUnmapResources(2, Resources, 0);
			return;
		}

		// Copy Depth
		cudaArray *CuArrayDepth;
		cudaGraphicsSubResourceGetMappedArray(&CuArrayDepth, CudaResourceDepthLeft, 0, 0);
		PrintLastCudaError("cudaGraphicsSubResourceGetMappedArray");

		cudaMemcpy2DToArray(
			CuArrayDepth, // dst array
			0, 0,    // offset
			CudaLinearMemoryDepth, Width * sizeof(float),       // src
			Width * sizeof(float), Height, // extent
			cudaMemcpyDeviceToDevice); // kind
		PrintLastCudaError("cudaMemcpy2DToArray");

		// Copy Color

		cudaArray *CuArrayColor;
		cudaGraphicsSubResourceGetMappedArray(&CuArrayColor, CudaResourceColorLeft, 0, 0);
		PrintLastCudaError("cudaGraphicsSubResourceGetMappedArray");

		cudaMemcpy2DToArray(
			CuArrayColor, // dst array
			0, 0,    // offset
			CudaLinearMemoryColor, Width * 4 * sizeof(float),       // src
			Width * 4 * sizeof(float), Height, // extent
			cudaMemcpyDeviceToDevice); // kind
		PrintLastCudaError("cudaMemcpy2DToArray");


		cudaGraphicsUnmapResources(2, Resources, 0);
		PrintLastCudaError("cudaGraphicsUnmapResources");

		//D3DDeviceContext->Flush();
		LaunchLaser();
		UpdateCubemapBuffer(RHICmdList);

	}
	else if(InView.StereoPass == EStereoscopicPass::eSSP_RIGHT_EYE)
	{
		/*optix::uchar4* Data = static_cast<optix::uchar4*>(OptiXContext->GetBuffer("result_color")->MapNative());
		RHICmdList.UpdateTexture2D(OutputTextureColorRightRef, 0, TextureRegion, Size.X * 4, (uint8*)Data);
		OptiXContext->GetBuffer("result_color")->Unmap();*/

		////float* Data2 = static_cast<float*>(OptiXContext->GetBuffer("result_depth")->MapNative());
		////RHICmdList.UpdateTexture2D(OutputTextureDepthRightRef, 0, TextureRegion, Size.X * 4, (uint8*)Data2);
		////OptiXContext->GetBuffer("result_depth")->Unmap();

		if (Resources[2] == NULL && Resources[3] == NULL)
		{
			return;
		}
		cudaGraphicsMapResources(2, Resources + 2, 0);
		PrintLastCudaError("cudaGraphicsMapResources");

		if (CudaResourceDepthRight == NULL)
		{
			cudaGraphicsUnmapResources(2, Resources + 2, 0);
			return;
		}
		// Depth
		cudaArray *CuArrayDepth;
		cudaGraphicsSubResourceGetMappedArray(&CuArrayDepth, CudaResourceDepthRight, 0, 0);
		PrintLastCudaError("cudaGraphicsSubResourceGetMappedArray");

		cudaMemcpy2DToArray(
			CuArrayDepth, // dst array
			0, 0,    // offset
			CudaLinearMemoryDepth, Width * sizeof(float),       // src
			Width * sizeof(float), Height, // extent
			cudaMemcpyDeviceToDevice); // kind
		//PrintLastCudaError("cudaMemcpy2DToArray");


		// Color
		cudaArray *CuArrayColor;
		cudaGraphicsSubResourceGetMappedArray(&CuArrayColor, CudaResourceColorRight, 0, 0);
		PrintLastCudaError("cudaGraphicsSubResourceGetMappedArray");

		cudaMemcpy2DToArray(
			CuArrayColor, // dst array
			0, 0,    // offset
			CudaLinearMemoryColor, Width * 4 * sizeof(float),       // src
			Width * 4 * sizeof(float), Height, // extent
			cudaMemcpyDeviceToDevice); // kind
		PrintLastCudaError("cudaMemcpy2DToArray");

		cudaGraphicsUnmapResources(2, Resources + 2, 0);
		PrintLastCudaError("cudaGraphicsUnmapResources");
	}
	//D3DDeviceContext->Flush();
	//end = FPlatformTime::Seconds();
	//UE_LOG(LogTemp, Warning, TEXT("Update took %f seconds"), end - start);

	else if (InView.StereoPass == EStereoscopicPass::eSSP_FULL)
	{
		UE_LOG(LogTemp, Display, TEXT("Full Pass"));
	}

	if (bCleanup)
	{
		CleanupOptiXOnEnd();
	}

}


void FOptiXContextManager::LaunchLaser()
{
	if (/*bSceneChanged && */ bLaserIsInitialized && !CVarDisableLaserTrace.GetValueOnRenderThread())
	{
		if (LaserActor.IsValid())
		{
			//bool bQueueTransformUpdate = LaserActor->OptiXLaserComponent->bPatternChanged;
			LaserActor->OptiXLaserComponent->UpdateOptiXContextVariables();
			//LaserActor->LineInstancedStaticMeshComponent->UpdateLines();
			//if (bQueueTransformUpdate)
			//{
			//	AsyncTask(ENamedThreads::GameThread, [Laser = LaserActor.Get()]() {
			//		Laser->LineInstancedStaticMeshComponent->UpdateLines();
			//	}
			//	);
			//}
		}

		// uuuuuuuuh
		static uint32 RandomSeed = 0;
		OptiXContext->SetUint("random_frame_seed", RandomSeed++);
		//UE_LOG(LogTemp, Warning, TEXT("Launching Laser Trace at Entry Point: %i"), LaserEntryPoint);
		bIsTracing.AtomicSet(true);
		OptiXContext->Launch(1, 50, 50, 20);
		bIsTracing.AtomicSet(false);


		if (Resources[4] == NULL)
		{
			return;
		}

		cudaGraphicsMapResources(1, Resources + 4, 0);
		PrintLastCudaError("cudaGraphicsMapResources");

		if (CudaResourceIntersections == NULL)
		{
			cudaGraphicsUnmapResources(1, Resources + 4, 0);
			return;
		}

		cudaArray *CuArrayIntersections;
		cudaGraphicsSubResourceGetMappedArray(&CuArrayIntersections, CudaResourceIntersections, 0, 0);
		PrintLastCudaError("cudaGraphicsSubResourceGetMappedArray");

		cudaMemcpy2DToArray(
			CuArrayIntersections, // dst array
			0, 0,    // offset
			CudaLinearMemoryIntersections, LaserBufferWidth * 4 * sizeof(float),       // src
			LaserBufferWidth * 4 * sizeof(float), LaserBufferHeight, // extent
			cudaMemcpyDeviceToDevice); // kind
		PrintLastCudaError("cudaMemcpy2DToArray");


		cudaGraphicsUnmapResources(1, Resources + 4, 0);
		PrintLastCudaError("cudaGraphicsUnmapResources");

		//optix::float4* DataLaser = static_cast<optix::float4*>(LaserOutputBuffer->MapNative(0, RT_BUFFER_MAP_READ));
		////FMemory::Memcpy(IntersectionData.GetData(), DataLaser, LaserOutputBuffer->GetSize1D() * sizeof(FVector4));

		//if (DataLaser == nullptr)
		//{
		//	UE_LOG(LogTemp, Error, TEXT("Error when trying to map laser output buffer: Got NULL"));
		//	return;
		//}

		//optix::float4 invData = optix::make_float4(0, -1, 0, 1);

		//uint32 N = 0;
		//// Loop over indices
		//for (uint32 i = 0; i < 50 * 50; ++i)
		//{
		//	if (DataLaser[i * 20 * 2].x == invData.x && DataLaser[i * 20 * 2].y == invData.y && DataLaser[i * 20 * 2].z == invData.z)
		//	{
		//		continue;
		//	}

		//	//if (!PreviousLaserResults.IsValidIndex(N))
		//	//{
		//	//	PreviousLaserResults.AddDefaulted(1);
		//	//	PreviousLaserResults[N].AddZeroed(20 * 2);
		//	//}


		//	TPair<uint32, TArray<FVector>> QueueItem;
		//	QueueItem.Key = N;
		//	N++;
		//	bool bEnqueue = false;
		//	for (uint32 Intersection = i * 20 * 2; Intersection < i * 20 * 2 + 20 * 2; ++Intersection)
		//	{
		//		FVector Pos(DataLaser[Intersection].x, DataLaser[Intersection].y, DataLaser[Intersection].z);
		//		//if (Pos != PreviousLaserResults[N][Intersection - i * 20 * 2])
		//		//{
		//		//	PreviousLaserResults[N][Intersection - i * 20 * 2] = Pos;
		//		//	bEnqueue = true;
		//		//}
		//		QueueItem.Value.Add(Pos);
		//	}
		//	//N++;
		//	//if(bEnqueue)
		//	LaserIntersectionQueue.Enqueue(QueueItem);
		//}

		//LaserOutputBuffer->Unmap();

		bSceneChanged.AtomicSet(false);
		LaserTraceFinishedEvent.Broadcast();
	}
}


bool FOptiXContextManager::IsActiveThisFrame(FViewport * InViewport) const
{
	//UE_LOG(LogTemp, Warning, TEXT("IsActiveThisFrame"));

	bool bDisableTrace = CVarDisableTrace.GetValueOnGameThread(); // Bad naming fix me
	return OptiXContext.IsValid() && !bDisableTrace && bStartTracing/* && TrackingSystem->IsHeadTrackingAllowed()*/;
}

void FOptiXContextManager::Init()
{

	// TODO Fix me there's still an optix error in there somewhere
	//if (CubemapSampler.IsValid())
	//{
	//	CubemapSampler->RemoveFromRoot();
	//	CubemapSampler->GetNativeTextureSampler()->destroy();
	//	CubemapSampler->MarkPendingKill();
	//	CubemapSampler.Reset();
	//}
	//if (CubemapBuffer.IsValid())
	//{
	//	CubemapBuffer->RemoveFromRoot();
	//	CubemapBuffer->GetNativeBuffer()->destroy();
	//	CubemapBuffer->MarkPendingKill();
	//	CubemapBuffer.Reset();
	//}

	// Probably don't need this at all
	if (GEngine)
	{
		GEngine->ForceGarbageCollection();
	}

	// Shouldn't be anything in the queues but clean up anyway just to be sure.
	DestroyOptiXObjects();

	//TODO: Shut this thing down correctly - for now just clean up anything when restarting
	CleanupOptiXOnEnd();

	InitContext();
	InitRendering();
	InitBuffers();
	InitPrograms();
	InitLaser();
	InitCubemap();
	//InitCUDADX();
	bIsInitialized = false;
	bStartTracing = true;
}

void FOptiXContextManager::SceneChangedCallback()
{
	bSceneChanged.AtomicSet(true);
}

void FOptiXContextManager::InitContext()
{

	UE_LOG(LogTemp, Display, TEXT("Initializing Context in ContextManager"));

	// Needs to be called BEFORE the context is created!
	//rtGlobalSetAttribute(RT_GLOBAL_ATTRIBUTE_ENABLE_RTX, sizeof(RTXOn), &RTXOn);

	OptiXContext = NewObject<UOptiXContext>(GetTransientPackage(), UOptiXContext::StaticClass());
	OptiXContext->AddToRoot();
	NativeContext = OptiXContext->Init();


	//OptiXContext->SetExceptionEnabled(RTexception::RT_EXCEPTION_PAYLOAD_ACCESS_OUT_OF_BOUNDS, false);
	//OptiXContext->SetExceptionEnabled(RTexception::RT_EXCEPTION_USER_EXCEPTION_CODE_OUT_OF_BOUNDS, true);
	//OptiXContext->SetExceptionEnabled(RTexception::RT_EXCEPTION_TRACE_DEPTH_EXCEEDED, true);
	//OptiXContext->SetExceptionEnabled(RTexception::RT_EXCEPTION_PROGRAM_ID_INVALID, true);
	//OptiXContext->SetExceptionEnabled(RTexception::RT_EXCEPTION_TEXTURE_ID_INVALID, true);
	//OptiXContext->SetExceptionEnabled(RTexception::RT_EXCEPTION_BUFFER_ID_INVALID, true);
	
	//OptiXContext->SetExceptionEnabled(RTexception::RT_EXCEPTION_INDEX_OUT_OF_BOUNDS, true);
	//OptiXContext->SetExceptionEnabled(RTexception::RT_EXCEPTION_STACK_OVERFLOW, true);
	//OptiXContext->SetExceptionEnabled(RTexception::RT_EXCEPTION_INVALID_RAY, true);
	//OptiXContext->SetExceptionEnabled(RTexception::RT_EXCEPTION_INTERNAL_ERROR, true);
	//OptiXContext->SetExceptionEnabled(RTexception::RT_EXCEPTION_BUFFER_INDEX_OUT_OF_BOUNDS, true);
	//OptiXContext->SetExceptionEnabled(RTexception::RT_EXCEPTION_USER, true);
	//OptiXContext->SetExceptionEnabled(RTexception::RT_EXCEPTION_USER_MAX, true);

	//OptiXContext->SetExceptionEnabled(RTexception::RT_EXCEPTION_ALL, true);



	//NativeContext->setPrintEnabled(true);
	//NativeContext->setPrintLaunchIndex(100, 100);
	// Set some default values, they can (and should) be overwritten in the game mode as they're scene specific
	OptiXContext->SetRayTypeCount(2);
	OptiXContext->SetEntryPointCount(1);
	//OptiXContext->SetStackSize(4000);
	OptiXContext->SetMaxTraceDepth(31);

	OptiXContext->SetInt("max_depth", 10);
	OptiXContext->SetFloat("scene_epsilon", 1.e-2f);

	TopObject = OptiXContext->CreateGroup();
	TopAcceleration = OptiXContext->CreateAcceleration("Trbvh"); // Here the accel structure seems to be actually needed
	//TopAcceleration->AddToRoot();
	TopAcceleration->SetProperty("refit", "1");

	TopObject->SetAcceleration(TopAcceleration.Get());

	OptiXContext->SetGroup("top_object", TopObject.Get());

	// Keep buffers and programs with the camera manager for now, there's no real reason yet to force a refacturing there

}


void FOptiXContextManager::InitRendering()
{

	UE_LOG(LogTemp, Display, TEXT("Initializing Rendering in ContextManager"));


	// Are we using an HMD?
	if (GEngine->XRSystem.IsValid() && GEngine->XRSystem->GetHMDDevice() != nullptr)
	{
		UE_LOG(LogTemp, Display, TEXT("Got HMD in ContextManager"));

		bWithHMD = GEngine->XRSystem->GetHMDDevice()->IsHMDEnabled();
	}
	else
	{
		UE_LOG(LogTemp, Display, TEXT("Running without HMD in ContextManager"));

		bWithHMD = false;
	}

	// Viewport size:
	FViewport* CurrentViewport = GEngine->GameViewport->Viewport;

	Width = CurrentViewport->GetSizeXY().X / 2.0;
	Height = CurrentViewport->GetSizeXY().Y;

	UE_LOG(LogTemp, Display, TEXT("Got viewport sizes: %i, %i"), Width, Height);
	UE_LOG(LogTemp, Warning, TEXT("Full Res: %i %i"), Width * 2, Height);


	// Apparently those can be 0 in a packaged build? 
	// Catch that case:
	if (Width == 0 || Height == 0)
	{
		UGameUserSettings* GameSettings = GEngine->GetGameUserSettings();
		Width = GameSettings->GetScreenResolution().X;
		Height = GameSettings->GetScreenResolution().Y;
		UE_LOG(LogTemp, Display, TEXT("Fallback to viewport size in settings: %i, %i"), Width, Height);

	}


	// Create the textures:

	OutputTexture = UTexture2D::CreateTransient(Width, Height, PF_A32B32G32R32F);
	OutputTexture->AddToRoot();
	//// Allocate the texture HRI
	OutputTexture->UpdateResource();

	DepthTexture = UTexture2D::CreateTransient(Width, Height, PF_R32_FLOAT);
	DepthTexture->AddToRoot();
	//// Allocate the texture HRI
	DepthTexture->UpdateResource();

	OutputTextureColorRightRef = ((FTexture2DResource*)OutputTexture->Resource)->GetTexture2DRHI();
	OutputTextureDepthRightRef = ((FTexture2DResource*)DepthTexture->Resource)->GetTexture2DRHI();

	if (bWithHMD)
	{
		OutputTexture2 = UTexture2D::CreateTransient(Width, Height, PF_A32B32G32R32F);
		OutputTexture2->AddToRoot();
		//// Allocate the texture HRI
		OutputTexture2->UpdateResource();

		DepthTexture2 = UTexture2D::CreateTransient(Width, Height, PF_R32_FLOAT);
		DepthTexture2->AddToRoot();
		//// Allocate the texture HRI
		DepthTexture2->UpdateResource();

		OutputTextureDepthLeftRef = ((FTexture2DResource*)DepthTexture2->Resource)->GetTexture2DRHI();
		OutputTextureColorLeftRef = ((FTexture2DResource*)OutputTexture2->Resource)->GetTexture2DRHI();

		// TODO Maybe we need to do this after setting the parameter?
	}

	UE_LOG(LogTemp, Display, TEXT("Created the Textures"));

	// Laser Texture
	LaserIntersectionTexture = UTexture2D::CreateTransient(LaserBufferWidth, LaserBufferHeight, PF_A32B32G32R32F); // TODO Hardcoded values
	LaserIntersectionTexture->AddToRoot();
	//// Allocate the texture HRI
	LaserIntersectionTexture->UpdateResource();

	LaserIntersectionTextureRef = ((FTexture2DResource*)LaserIntersectionTexture->Resource)->GetTexture2DRHI();

	// Set up the material

	// Load the materials
	RegularMaterial = LoadObject<UMaterial>(GetTransientPackage(), TEXT("Material'/OptiX/PPMaterials/TextureMaterial.TextureMaterial'"));
	VRMaterial = LoadObject<UMaterial>(GetTransientPackage(), TEXT("Material'/OptiX/PPMaterials/TextureMaterialVR.TextureMaterialVR'"));
	LaserMaterial = LoadObject<UMaterial>(GetTransientPackage(), TEXT("Material'/OptiX/Laser/LaserMaterial.LaserMaterial'"));
	LaserMaterialDynamic = UMaterialInstanceDynamic::Create(LaserMaterial.Get(), OptiXContext.Get(), "DynamicLaserMaterial");

	LaserMaterialDynamic->SetTextureParameterValue("IntersectionTexture", LaserIntersectionTexture.Get());
	LaserMaterialDynamic->SetScalarParameterValue("Lines", 50);
	LaserMaterialDynamic->SetScalarParameterValue("Segments", 20);


	if(RegularMaterial == nullptr || VRMaterial == nullptr)
	{
		UE_LOG(LogTemp, Error, TEXT("Couldn't load dummy Material!"));
	}

	if (bWithHMD)
	{
		DynamicMaterial = UMaterialInstanceDynamic::Create(VRMaterial.Get(), OptiXContext.Get(), "DynamicVRMaterial");
		DynamicMaterial->SetTextureParameterValue("TextureRight", OutputTexture.Get());
		DynamicMaterial->SetTextureParameterValue("DepthRight", DepthTexture.Get());
		DynamicMaterial->SetTextureParameterValue("TextureLeft", OutputTexture2.Get());
		DynamicMaterial->SetTextureParameterValue("DepthLeft", DepthTexture2.Get());


		DynamicMaterial2D = UMaterialInstanceDynamic::Create(RegularMaterial.Get(), OptiXContext.Get(), "DynamicNonVRMaterial");
		DynamicMaterial2D->SetTextureParameterValue("Texture", OutputTexture.Get());
		DynamicMaterial2D->SetTextureParameterValue("Depth", DepthTexture.Get());

	}
	else
	{
		DynamicMaterial = UMaterialInstanceDynamic::Create(RegularMaterial.Get(), OptiXContext.Get(), "DynamicNonVRMaterial");
		DynamicMaterial->SetTextureParameterValue("Texture", OutputTexture.Get());
		DynamicMaterial->SetTextureParameterValue("Depth", DepthTexture.Get());
	}

	UE_LOG(LogTemp, Display, TEXT("Finished Initializing Rendering in ContextManager"));


	//OutputTextureColorRightRef = ((FTexture2DResource*)OutputTexture->Resource)->GetTexture2DRHI();
	//OutputTextureDepthRightRef = ((FTexture2DResource*)DepthTexture->Resource)->GetTexture2DRHI();
	//OutputTextureDepthLeftRef = ((FTexture2DResource*)DepthTexture2->Resource)->GetTexture2DRHI();
	//OutputTextureColorLeftRef = ((FTexture2DResource*)OutputTexture2->Resource)->GetTexture2DRHI();

}

void FOptiXContextManager::InitBuffers()
{
	OutputBuffer = OptiXContext->CreateOutputBufferColor(Width, Height);
	OutputDepthBuffer = OptiXContext->CreateOutputBufferDepth(Width, Height);

	OptiXContext->SetBuffer("result_color", OutputBuffer.Get());
	OptiXContext->SetBuffer("result_depth", OutputDepthBuffer.Get());
}

void FOptiXContextManager::InitPrograms()
{
	FString OptiXPTXDir = FOptiXModule::Get().OptiXPTXDir;

	// Generation Program
	RayGenerationProgram = OptiXContext->CreateProgramFromPTXFile
	(
		OptiXPTXDir + "generated/perspective_camera.ptx",
		"pinhole_camera"
	);
	OptiXContext->SetRayGenerationProgram(0, RayGenerationProgram.Get());

	// Exception program
	ExceptionProgram = OptiXContext->CreateProgramFromPTXFile
	(
		OptiXPTXDir + "generated/exception.ptx",
		"exception"
	);
	OptiXContext->SetExceptionProgram(0, ExceptionProgram.Get());

	// Miss Program
	MissProgram = OptiXContext->CreateProgramFromPTXFile
	(
		OptiXPTXDir + "generated/skybox.ptx",
		"skyboxLookup"
	);
	OptiXContext->SetMissProgram(0, MissProgram.Get());
	OptiXContext->SetFloat3DVector("bg_color", FVector(1.0, 1.0, 1.0));
}



void FOptiXContextManager::InitLaser()
{

	FString OptiXPTXDir = FOptiXModule::Get().OptiXPTXDir;


	LaserEntryPoint = OptiXContext->GetEntryPointCount();

	int32 RayTypeCount = OptiXContext->GetRayTypeCount();
	OptiXContext->SetRayTypeCount(RayTypeCount + 1);

	UE_LOG(LogTemp, Display, TEXT("Setting Laser Entry Point to %i"), LaserEntryPoint);
	UE_LOG(LogTemp, Display, TEXT("Setting Ray Type Index to %i"), RayTypeCount);


	// Increase EntryPointCount by 1
	OptiXContext->SetEntryPointCount(LaserEntryPoint + 1);

	// TODO maybe do this explicitely - loads the same program twice, but at least it's clear which one is used then.

	LaserExceptionProgram = OptiXContext->CreateProgramFromPTXFile
	(
		OptiXPTXDir + "generated/exception.ptx",
		"exception"
	);

	OptiXContext->SetExceptionProgram(1 /* todo- diff between raytypeindex and entrypointcount, this is 1 in the original app*/, LaserExceptionProgram.Get());

	LaserRayGenerationProgram = OptiXContext->CreateProgramFromPTXFile
	(
		OptiXPTXDir + "generated/laser_caster.ptx",
		"laser_caster"
	);
	OptiXContext->SetRayGenerationProgram(LaserEntryPoint, LaserRayGenerationProgram.Get());

	LaserMissProgram = OptiXContext->CreateProgramFromPTXFile
	(
		OptiXPTXDir + "generated/miss.ptx",
		"miss_iterative"
	);

	OptiXContext->SetMissProgram(1 /*LaserEntryPoint /* this is 1 in the original application, why? TODO*/, LaserMissProgram.Get());

	//LaserOutputBuffer = OptiXContext->CreateBuffer(RT_BUFFER_OUTPUT, RT_FORMAT_FLOAT4, LaserBufferSize);
	LaserOutputBuffer = OptiXContext->CreateOutputBufferIntersections(LaserBufferWidth, LaserBufferHeight);
	LaserOutputBuffer->AddToRoot();
	OptiXContext->SetBuffer("result_laser", LaserOutputBuffer.Get());

	OptiXContext->SetInt("max_depth_laser", LaserMaxDepth);

	UOptiXBuffer* LaserIndexBuffer = OptiXContext->CreateBuffer(RT_BUFFER_INPUT, RT_FORMAT_INT, 50, 50);
	//LaserIndexBuffer->AddToRoot();
	UOptiXBuffer* LaserDirectionBuffer = OptiXContext->CreateBuffer(RT_BUFFER_INPUT, RT_FORMAT_FLOAT3, 50, 50);
	//LaserDirectionBuffer->AddToRoot();
	OptiXContext->SetBuffer("laserIndex", LaserIndexBuffer);
	OptiXContext->SetBuffer("laserDir", LaserDirectionBuffer);


}

void FOptiXContextManager::InitCubemap()
{
	// todo max # cubemaps
	for (int32 i = 1; i < 10; i++) // 0 is reserved for this (player camera)
	{
		UnallocatedCubemapIds.Enqueue(i);
	}

	// TODO: Try and see if destroying/creating the whole thing and doing a memcpy on the GPU only is 
	// quicker than updating the cubemap each frame.

	CubemapsInputBuffer = OptiXContext->CreateBuffer(RT_BUFFER_INPUT, RTformat::RT_FORMAT_INT, 10);
	OptiXContext->SetBuffer("skyboxBuffer", CubemapsInputBuffer.Get());

	CubemapSampler = OptiXContext->CreateTextureSampler();
	//CubemapSampler->AddToRoot();
	CubemapSampler->SetWrapMode(0, RT_WRAP_CLAMP_TO_EDGE);
	CubemapSampler->SetWrapMode(1, RT_WRAP_CLAMP_TO_EDGE);
	CubemapSampler->SetWrapMode(2, RT_WRAP_CLAMP_TO_EDGE);
	CubemapSampler->SetIndexingMode(RT_TEXTURE_INDEX_NORMALIZED_COORDINATES);
	CubemapSampler->SetReadMode(RT_TEXTURE_READ_NORMALIZED_FLOAT);
	CubemapSampler->SetMaxAnisotropy(1.0f);
	CubemapSampler->SetMipLevelCount(1u);
	CubemapSampler->SetArraySize(1u);


	CubemapBuffer = OptiXContext->CreateCubemapBuffer(1024, 1024);
	//CubemapBuffer->AddToRoot();

	CubemapSampler->SetBufferWithTextureIndexAndMiplevel(0u, 0u, CubemapBuffer.Get());
	CubemapSampler->SetFilteringModes(RT_FILTER_LINEAR, RT_FILTER_LINEAR, RT_FILTER_NONE);

	OptiXContext->SetSkybox("skybox0", CubemapSampler.Get());

	//RequestCubemapId();
	AddCubemapToBuffer(0, CubemapSampler->GetId());
	
	//OptiXContext->SetTextureSampler("skybox", CubemapSampler.Get());

	UE_LOG(LogTemp, Display, TEXT("Successfully initialized cubemap."));
}

int32 FOptiXContextManager::RequestCubemapId()
{
	if (UnallocatedCubemapIds.IsEmpty())
	{
		return 0;
	}
	int32 Id;
	UnallocatedCubemapIds.Dequeue(Id);
	return Id;
}

void FOptiXContextManager::DeleteCubemapId(int32 Id)
{
	if (Id <= 10)
	{
		UE_LOG(LogTemp, Warning, TEXT("Trying to free a cubemap that isn't there."));
		return;
	}
	// The Component itself should handle deletion of the sampler.
	UnallocatedCubemapIds.Enqueue(Id);
}

void FOptiXContextManager::UpdateCubemapBuffer(FRHICommandListImmediate & RHICmdList)
{
	if (!CameraActor.IsValid() || bValidCubemap)
	{
		return;
	}

	if (!CameraActor->bCubemapCaptured)
	{
		return;
	}


	int32 X = CameraActor->CubeRenderTarget->SizeX;
	int32 Y = X;

	SurfaceDataCube.Empty();
	SurfaceDataCube.SetNumZeroed(6);

	//TArray<FLinearColor> SD;

	optix::uchar4* BufferData = static_cast<optix::uchar4*>(CubemapBuffer->MapNative());

	FTextureRenderTargetCubeResource* RenderTargetCube = static_cast<FTextureRenderTargetCubeResource*>(CameraActor->CubeRenderTarget->GetRenderTargetResource());

	FIntRect InRectCube = FIntRect(0, 0, RenderTargetCube->GetSizeXY().X, RenderTargetCube->GetSizeXY().Y);
	FReadSurfaceDataFlags FlagsCube0(RCM_UNorm, CubeFace_PosX);
	FReadSurfaceDataFlags FlagsCube1(RCM_UNorm, CubeFace_NegX);
	FReadSurfaceDataFlags FlagsCube2(RCM_UNorm, CubeFace_PosY);
	FReadSurfaceDataFlags FlagsCube3(RCM_UNorm, CubeFace_NegY);
	FReadSurfaceDataFlags FlagsCube4(RCM_UNorm, CubeFace_PosZ);
	FReadSurfaceDataFlags FlagsCube5(RCM_UNorm, CubeFace_NegZ);

	RHICmdList.ReadSurfaceData(RenderTargetCube->GetTextureRHI(), InRectCube, SurfaceDataCube[0], FlagsCube0);
	RHICmdList.ReadSurfaceData(RenderTargetCube->GetTextureRHI(), InRectCube, SurfaceDataCube[1], FlagsCube1);
	RHICmdList.ReadSurfaceData(RenderTargetCube->GetTextureRHI(), InRectCube, SurfaceDataCube[2], FlagsCube2);
	RHICmdList.ReadSurfaceData(RenderTargetCube->GetTextureRHI(), InRectCube, SurfaceDataCube[3], FlagsCube3);
	RHICmdList.ReadSurfaceData(RenderTargetCube->GetTextureRHI(), InRectCube, SurfaceDataCube[4], FlagsCube4);
	RHICmdList.ReadSurfaceData(RenderTargetCube->GetTextureRHI(), InRectCube, SurfaceDataCube[5], FlagsCube5);

	uint32 MemSize = (X * Y * sizeof(FColor));
	FMemory::Memcpy(BufferData, SurfaceDataCube[0].GetData(), MemSize); // front
	FMemory::Memcpy(BufferData + X * Y * 1, SurfaceDataCube[1].GetData(), MemSize); // back
	FMemory::Memcpy(BufferData + X * Y * 2, SurfaceDataCube[2].GetData(), MemSize); // 
	FMemory::Memcpy(BufferData + X * Y * 3, SurfaceDataCube[3].GetData(), MemSize); // 
	FMemory::Memcpy(BufferData + X * Y * 4, SurfaceDataCube[4].GetData(), MemSize); // 
	FMemory::Memcpy(BufferData + X * Y * 5, SurfaceDataCube[5].GetData(), MemSize); //

	CubemapBuffer->Unmap();
	bValidCubemap.AtomicSet(true);
}

void FOptiXContextManager::AddCubemapToBuffer(int32 CubemapId, int32 SamplerId)
{
	int32* Data = static_cast<int32*>(CubemapsInputBuffer->MapNative());
	Data[CubemapId] = SamplerId;
	CubemapsInputBuffer->Unmap();
}

void FOptiXContextManager::InitCUDADX()
{
	
	// Setup DX:

	D3DDevice = (ID3D11Device*)GDynamicRHI->RHIGetNativeDevice();
	D3DDevice->GetImmediateContext(&D3DDeviceContext);

	// Create texture for now:

	Width = Width;
	Height = Height;

	OutputTextureDepthLeftRef = ((FTexture2DResource*)DepthTexture2->Resource)->GetTexture2DRHI();

	// Depth Left
	D3D11_TEXTURE2D_DESC DescDepthLeft;
	ZeroMemory(&DescDepthLeft, sizeof(D3D11_TEXTURE2D_DESC));
	ID3D11Texture2D* D3D11DepthLeftTexture = static_cast<ID3D11Texture2D*>(OutputTextureDepthLeftRef->GetNativeResource());
	D3D11DepthLeftTexture->GetDesc(&DescDepthLeft);
	UE_LOG(LogTemp, Display, TEXT("ID3D11Texture2D Info Depth: Format is %i"), int(DescDepthLeft.Format));

	// Depth Right
	OutputTextureDepthRightRef = ((FTexture2DResource*)DepthTexture->Resource)->GetTexture2DRHI();

	D3D11_TEXTURE2D_DESC DescDepthRight;
	ZeroMemory(&DescDepthRight, sizeof(D3D11_TEXTURE2D_DESC));
	ID3D11Texture2D* D3D11DepthRightTexture = static_cast<ID3D11Texture2D*>(OutputTextureDepthRightRef->GetNativeResource());
	D3D11DepthLeftTexture->GetDesc(&DescDepthRight);

	// Color Left

	OutputTextureColorLeftRef = ((FTexture2DResource*)OutputTexture2->Resource)->GetTexture2DRHI();

	D3D11_TEXTURE2D_DESC DescColorLeft;
	ZeroMemory(&DescColorLeft, sizeof(D3D11_TEXTURE2D_DESC));
	ID3D11Texture2D* D3D11ColorLeftTexture = static_cast<ID3D11Texture2D*>(OutputTextureColorLeftRef->GetNativeResource());
	D3D11ColorLeftTexture->GetDesc(&DescColorLeft);
	UE_LOG(LogTemp, Display, TEXT("ID3D11Texture2D Info Color : Format is %i"), int(DescColorLeft.Format));


	//// Color Right
	OutputTextureColorRightRef = ((FTexture2DResource*)OutputTexture->Resource)->GetTexture2DRHI();

	D3D11_TEXTURE2D_DESC DescColorRight;
	ZeroMemory(&DescColorRight, sizeof(D3D11_TEXTURE2D_DESC));
	ID3D11Texture2D* D3D11ColorRightTexture = static_cast<ID3D11Texture2D*>(OutputTextureColorRightRef->GetNativeResource());
	D3D11ColorRightTexture->GetDesc(&DescColorRight);


	//// Intersections
	LaserIntersectionTextureRef = ((FTexture2DResource*)LaserIntersectionTexture->Resource)->GetTexture2DRHI();

	D3D11_TEXTURE2D_DESC DescIntersections;
	ZeroMemory(&DescIntersections, sizeof(D3D11_TEXTURE2D_DESC));
	ID3D11Texture2D* D3D11IntersectionTexture = static_cast<ID3D11Texture2D*>(LaserIntersectionTextureRef->GetNativeResource());
	D3D11IntersectionTexture->GetDesc(&DescIntersections);
	

	// Register the unreal textures with cuda
	cudaGraphicsD3D11RegisterResource(&CudaResourceDepthLeft, D3D11DepthLeftTexture, cudaGraphicsRegisterFlagsNone);
	PrintLastCudaError("cudaGraphicsD3D11RegisterResource");

	cudaGraphicsD3D11RegisterResource(&CudaResourceDepthRight, D3D11DepthRightTexture, cudaGraphicsRegisterFlagsNone);
	PrintLastCudaError("cudaGraphicsD3D11RegisterResource");

	cudaGraphicsD3D11RegisterResource(&CudaResourceColorLeft, D3D11ColorLeftTexture, cudaGraphicsRegisterFlagsNone);
	PrintLastCudaError("cudaGraphicsD3D11RegisterResource");

	cudaGraphicsD3D11RegisterResource(&CudaResourceColorRight, D3D11ColorRightTexture, cudaGraphicsRegisterFlagsNone);
	PrintLastCudaError("cudaGraphicsD3D11RegisterResource");

	cudaGraphicsD3D11RegisterResource(&CudaResourceIntersections, D3D11IntersectionTexture, cudaGraphicsRegisterFlagsNone);
	PrintLastCudaError("cudaGraphicsD3D11RegisterResource");

	// Allocate the buffer memory
	//cudaMallocPitch(&CudaLinearMemoryDepth, &Pitch, Width * sizeof(float), Height);
	cudaMalloc(&CudaLinearMemoryDepth, Width * Height * sizeof(float));
	PrintLastCudaError("cudaMalloc");

	cudaMalloc(&CudaLinearMemoryColor, Width * Height * 4 * sizeof(float));
	PrintLastCudaError("cudaMalloc");

	cudaMalloc(&CudaLinearMemoryIntersections, LaserBufferWidth * LaserBufferHeight * 4 * sizeof(float));
	PrintLastCudaError("cudaMalloc");

	//cudaMallocPitch(&CudaLinearMemoryColorRight, &Pitch, Width * sizeof(optix::uchar4), Height);
	//PrintLastCudaError("cudaMallocPitch");

	//cudaMemset(CudaLinearMemory, 1, Pitch * Height);
	//PrintLastCudaError("cudaMemset");

	OptiXContext->GetBuffer("result_depth")->SetDevicePointer(0, CudaLinearMemoryDepth);
	OptiXContext->GetBuffer("result_color")->SetDevicePointer(0, CudaLinearMemoryColor);
	OptiXContext->GetBuffer("result_laser")->SetDevicePointer(0, CudaLinearMemoryIntersections);


	UE_LOG(LogTemp, Display, TEXT("Device Count: %i"), OptiXContext->GetDeviceCount());
	UE_LOG(LogTemp, Display, TEXT("Device Name 0: %s"), *OptiXContext->GetDeviceName(0));


	Resources[0] = CudaResourceDepthLeft;
	Resources[1] = CudaResourceColorLeft;
	Resources[2] = CudaResourceDepthRight;
	Resources[3] = CudaResourceColorRight;
	Resources[4] = CudaResourceIntersections;

	bIsInitialized = true;

}