OptiXLensComponent.cpp

// Fill out your copyright notice in the Description page of Project Settings.

#include "OptiXLensComponent.h"

#include "Runtime/Engine/Classes/Engine/TextureRenderTargetCube.h"


#include "OptiXModule.h"

UOptiXLensComponent::UOptiXLensComponent(const FObjectInitializer& ObjectInitializer)
	: Super(ObjectInitializer)
{
	Radius1 = 0.8 * 100; // todo default values
	Radius2 = 1.0 * 100;
	LensRadius = 0.1 * 100;

	LensType1 = ELensSideType::CONVEX;
	LensType2 = ELensSideType::CONVEX;
	LensThickness = 0.025 * 100;
	CurrentWavelength = 450.0f;

	// A little hacky but w/e
	for (const TPair<FString, FGlassDefinition>& Pair : FOptiXModule::Get().GetGlassDefinitions())
	{
		GlassType = Pair.Key;
		break;
	}
	
}


void UOptiXLensComponent::BeginPlay()
{
	Super::BeginPlay();


	// do this on the game thread
	// hook into WL update:
	//UE_LOG(LogTemp, Display, TEXT("Begin Play on LensComponent, GameThread"));

	FOptiXModule::Get().GetOptiXContextManager()->WavelengthChangedEvent.AddUFunction(this, "OnWavelengthChangedEvent");
}


void UOptiXLensComponent::UpdateOptiXComponentVariables()
{
	check(IsInRenderingThread());


	OptiXGeometryInstance->SetFloat("radius", Radius1);
	//UE_LOG(LogTemp, Display, TEXT("radius1 : %f"), Radius1);

	OptiXGeometryInstance->SetFloat("radius2", Radius2);
	//UE_LOG(LogTemp, Display, TEXT("radius2 : %f"), Radius2);

	OptiXGeometryInstance->SetFloat("lensRadius", LensRadius);
	//UE_LOG(LogTemp, Display, TEXT("lensradius : %f"), LensRadius);

	OptiXGeometryInstance->SetFloat("halfCylinderLength", GetCylinderLength(LensThickness) / 2.0f);
	//UE_LOG(LogTemp, Display, TEXT("halfCylinderLength : %f"), GetCylinderLength(LensThickness) / 2.0f);

	OptiXGeometryInstance->SetInt("side1Type", static_cast<int32>(LensType1));
	//UE_LOG(LogTemp, Display, TEXT("type1 : %i"), static_cast<int32>(LensType1));

	OptiXGeometryInstance->SetInt("side2Type", static_cast<int32>(LensType2));
	//UE_LOG(LogTemp, Display, TEXT("type2 : %i"), static_cast<int32>(LensType2));


	double WL2 = FMath::Pow(CurrentWavelength / 1000.0, 2.0f);
	FGlassDefinition Def = FOptiXModule::Get().GetGlassDefinitions()[GlassType];
	//UE_LOG(LogTemp, Display, TEXT("Glass Def: %f, %f, %F"), Def.B.X, Def.B.Y, Def.B.Z);

	float Index = FMath::Sqrt(1 +
		Def.B.X * WL2 / (WL2 - Def.C.X) +
		Def.B.Y * WL2 / (WL2 - Def.C.Y) +
		Def.B.Z * WL2 / (WL2 - Def.C.Z));

	OptiXMaterial->SetFloat("refraction_index", Index);

	MarkDirty();
}


void UOptiXLensComponent::UpdateOptiXComponent()
{
	if (OptiXGeometry != nullptr && OptiXTransform != nullptr && OptiXAcceleration != nullptr)
	{
		FMatrix T = GetComponentToWorld().ToMatrixNoScale();
		OptiXTransform->SetMatrix(T);
		//OptiXAcceleration->MarkDirty();
		OptiXContext->GetGroup("top_object")->GetAcceleration()->MarkDirty();
		FOptiXModule::Get().GetOptiXContextManager()->OnSceneChangedDelegate.Broadcast();
	}
}

void UOptiXLensComponent::InitOptiXGeometry()
{
	OptiXGeometry = OptiXContext->CreateGeometry();
	OptiXGeometry->SetPrimitiveCount(1u);
	UOptiXProgram* BB = OptiXContext->CreateProgramFromPTXFile(OptiXPTXDir + "generated/lens_parametric.ptx", "bounds");
	UOptiXProgram* IP = OptiXContext->CreateProgramFromPTXFile(OptiXPTXDir + "generated/lens_parametric.ptx", "intersect");

	OptiXGeometry->SetBoundingBoxProgram(BB);
	OptiXGeometry->SetIntersectionProgram(IP);
}

void UOptiXLensComponent::InitOptiXMaterial()
{
	UOptiXProgram* CHPerspective = OptiXContext->CreateProgramFromPTXFile(OptiXPTXDir + "generated/glass_perspective_camera.ptx", "closest_hit_radiance");
	UOptiXProgram* CHIterative = OptiXContext->CreateProgramFromPTXFile(OptiXPTXDir + "generated/glass_iterative_camera.ptx", "closest_hit_radiance");

	OptiXMaterial = OptiXContext->CreateMaterial();
	OptiXMaterial->SetClosestHitProgram(0, CHPerspective);
	OptiXMaterial->SetClosestHitProgram(1, CHIterative);

	OptiXMaterial->SetFloat("importance_cutoff", 1e-2f);
	OptiXMaterial->SetFloat3D("cutoff_color", 0.035f, 0.102f, 0.169f);
	OptiXMaterial->SetFloat("fresnel_exponent", 3.0f);
	OptiXMaterial->SetFloat("fresnel_minimum", 0.1f);
	OptiXMaterial->SetFloat("fresnel_maximum", 1.0f);
	OptiXMaterial->SetFloat("refraction_index", 1.4f);

	OptiXMaterial->SetFloat3D("refraction_color", 1.0f, 1.0f, 1.0f);
	OptiXMaterial->SetFloat3D("reflection_color", 1.0f, 1.0f, 1.0f);

	OptiXMaterial->SetInt("refraction_maxdepth", 10);
	OptiXMaterial->SetInt("reflection_maxdepth", 5);

	OptiXMaterial->SetFloat3D("extinction_constant", FMath::Loge(0.83f), FMath::Loge(0.83f), FMath::Loge(0.83f));

	OptiXMaterial->SetInt("lens_id", OptiXCubemapId);

}

void UOptiXLensComponent::InitOptiXGroups()
{
	OptiXGeometryInstance = OptiXContext->CreateGeometryInstance(OptiXGeometry, OptiXMaterial);

	OptiXGeometryInstance->SetFloat3D("center", 0.0f, 0.0f, 0.0f);

	OptiXGeometryInstance->SetFloat3DVector("orientation", {0, 0, -1}); // TODO Why this vector?!?!


	SetRadius1(Radius1);
	SetRadius2(Radius2);
	SetLensRadius(LensRadius);
	SetThickness(LensThickness);
	SetLensType1(LensType1);
	SetLensType2(LensType2);
	SetWavelength(CurrentWavelength); // min value

	OptiXTransform = OptiXContext->CreateTransform();

	FMatrix WorldTransform = GetComponentToWorld().ToMatrixNoScale();
	OptiXTransform->SetMatrix(WorldTransform);

	OptiXGeometryGroup = OptiXContext->CreateGeometryGroup();
	OptiXGeometryGroup->AddChild(OptiXGeometryInstance);

	OptiXAcceleration = OptiXContext->CreateAcceleration("NoAccel"); // Seems to be faster for now
	//OptiXAcceleration->SetProperty("refit", "1");

	OptiXGeometryGroup->SetAcceleration(OptiXAcceleration);

	OptiXTransform->SetChild(OptiXGeometryGroup);


	OptiXContext->GetGroup("top_object")->AddChild(OptiXTransform);
	MarkDirty();
}

void UOptiXLensComponent::InitCubemap(FRHICommandListImmediate & RHICmdList)
{
	UE_LOG(LogTemp, Display, TEXT("Init Cubemap"));

	//OptiXContext->SetBuffer("skyboxBuffer", CubemapBuffer);

	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);
	CubemapSampler->SetFilteringModes(RT_FILTER_LINEAR, RT_FILTER_LINEAR, RT_FILTER_NONE);
	

	UpdateCubemap(RHICmdList);


	// Writes the variable into the input buffer
	FOptiXModule::Get().GetOptiXContextManager()->AddCubemapToBuffer(OptiXCubemapId, CubemapSampler->GetId());
	UE_LOG(LogTemp, Display, TEXT("Finished Init Cubemap"));

}

void UOptiXLensComponent::UpdateCubemap(FRHICommandListImmediate & RHICmdList)
{

	UE_LOG(LogTemp, Display, TEXT("Updating Cubemap"));


	int32 X = 1024; // todo hardcoded
	int32 Y = X;

	TArray<TArray<FColor>> SurfaceDataCube;
	SurfaceDataCube.SetNumZeroed(6);

	//TArray<FLinearColor> SD;

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

	FTextureRenderTargetCubeResource* RenderTargetCube = static_cast<FTextureRenderTargetCubeResource*>(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();

	UE_LOG(LogTemp, Display, TEXT("Finished Updating Cubemap"));


}


void UOptiXLensComponent::CleanOptiXComponent()
{
	if(OptiXContext != NULL && OptiXContext->GetGroup("top_object") != NULL)
		OptiXContext->GetGroup("top_object")->RemoveChild(OptiXTransform);
	OptiXTransform = nullptr;

	Super::CleanOptiXComponent();
	OptiXGeometryGroup = nullptr;

}

void UOptiXLensComponent::InitFromData(const FLensData& Data)
{
	SetLensRadius(Data.LensRadius * 10);
	SetRadius1(Data.Radius1 * 10);
	SetRadius2(Data.Radius2 * 10);
	SetThickness(Data.Thickness);
	SetLensType1(Data.LensTypeSide1);
	SetLensType2(Data.LensTypeSide2);
	SetGlassType(Data.GlassType);
}

void UOptiXLensComponent::SetThickness(float Thickness)
{
	UE_LOG(LogTemp, Display, TEXT("Setting Thickness: %f"), Thickness);
	LensThickness = Thickness;
	QueueOptiXContextUpdate();
	if(IsInGameThread())
		OnLensThicknessChanged.Broadcast(LensThickness);
}

float UOptiXLensComponent::GetThickness() const
{
	// No silly conversions...
	return LensThickness;
}

void UOptiXLensComponent::SetRadius1(float Radius)
{
	UE_LOG(LogTemp, Display, TEXT("Setting Radius 1: %f"), Radius);
	Radius1 = Radius;
	QueueOptiXContextUpdate();
}

float UOptiXLensComponent::GetRadius1() const
{
	return Radius1;
}

void UOptiXLensComponent::SetRadius2(float Radius)
{

	UE_LOG(LogTemp, Display, TEXT("Setting Radius 2: %f"), Radius);
	Radius2 = Radius;
	QueueOptiXContextUpdate();
}

float UOptiXLensComponent::GetRadius2() const
{
	return Radius2;
}

void UOptiXLensComponent::SetLensRadius(float Radius)
{
	UE_LOG(LogTemp, Display, TEXT("Setting Lens Radius: %f"), Radius);
	LensRadius = Radius;
	QueueOptiXContextUpdate();
	if (IsInGameThread())
		OnLensRadiusChanged.Broadcast(Radius);
}

float UOptiXLensComponent::GetLensRadius() const
{
	return LensRadius;
}

void UOptiXLensComponent::SetLensType1(ELensSideType Type)
{
	UE_LOG(LogTemp, Display, TEXT("Setting Lens Side 1 Type: %i"), static_cast<int>(Type));

	LensType1 = Type;
	// Recalculate Thickness and set new half cylinder length
	//SetThickness(LensThickness);
	QueueOptiXContextUpdate();
}

ELensSideType UOptiXLensComponent::GetLensType1() const
{
	return LensType1;
}

void UOptiXLensComponent::SetLensType2(ELensSideType Type)
{
	UE_LOG(LogTemp, Display, TEXT("Setting Lens Side 2 Type: %i"), static_cast<int>(Type));

	LensType2 = Type;
	// Recalculate Thickness and set new half cylinder length
	//SetThickness(LensThickness);
	QueueOptiXContextUpdate();
}

ELensSideType UOptiXLensComponent::GetLensType2() const
{
	return LensType2;
}

void UOptiXLensComponent::SetGlassType(FString Type)
{
	UE_LOG(LogTemp, Display, TEXT("Setting Glass Type: %s"), *Type);

	GlassType = Type;
	SetWavelength(CurrentWavelength); // ???
	QueueOptiXContextUpdate();
}

FString UOptiXLensComponent::GetGlassType() const
{
	return GlassType;
}

void UOptiXLensComponent::SetWavelength(float WL)
{
	UE_LOG(LogTemp, Display, TEXT("Setting new WL in lens: %s"), *GetName());
	CurrentWavelength = WL;
	QueueOptiXContextUpdate();
}

float UOptiXLensComponent::GetWavelength() const
{
	return CurrentWavelength;
}

void UOptiXLensComponent::MarkDirty()
{

	UE_LOG(LogTemp, Display, TEXT("Marking Dirty"));


	if (OptiXGeometry != nullptr)
	{
		OptiXGeometry->MarkDirty();
	}
	if (OptiXAcceleration != nullptr)
	{
		OptiXAcceleration->MarkDirty();
	}

	//manager_->TopAccelerationMarkDirty();
	UOptiXAcceleration* TopAccel = OptiXContext->GetGroup("top_object")->GetAcceleration();
	if (TopAccel != nullptr)
	{
		TopAccel->MarkDirty(); // This should never be null, but check anyway
	}

	//RecalculateBoundingBox(); // TODO

}


float UOptiXLensComponent::GetCylinderLength(float Thickness) const
{
	// Halfsphere thickness
	float HalfThickness1;
	float HalfThickness2;

	// Side 1
	if (LensType1 == ELensSideType::PLANE)
	{
		HalfThickness1 = 0;
	}
	else
	{
		HalfThickness1 = Radius1 - FMath::Sqrt(-LensRadius * LensRadius + Radius1 * Radius1);
		if (LensType1 == ELensSideType::CONCAVE)
		{
			HalfThickness1 *= -1;
		}
	}

	// Side 2
	if (LensType2 == ELensSideType::PLANE)
	{
		HalfThickness2 = 0;
	}
	else
	{
		HalfThickness2 = Radius2 - FMath::Sqrt(-LensRadius * LensRadius + Radius2 * Radius2);
		if (LensType2 == ELensSideType::CONCAVE)
		{
			HalfThickness2 *= -1;
		}
	}
	return Thickness - HalfThickness1 - HalfThickness2;
}


void UOptiXLensComponent::RecalculateBoundingBox()
{

	// Do we even need this? Seems like it was just used for phoenix stuff TODO
	return;

	// TODO Shouldn't have to call the getter here - this should just be the actor rotation transform?
	/*FVector Orientation = OptiXGeometry->GetFloat3D("orientation");
	float HalfThickness1;	
	LensType1 == ELensSideType::CONVEX ? HalfThickness1 = Radius1 - FMath::Sqrt(-LensRadius * LensRadius + Radius1 * Radius1) : HalfThickness1 = 0;

	float HalfThickness2;
	LensType2 == ELensSideType::CONVEX ? HalfThickness2 = Radius2 - FMath::Sqrt(-LensRadius * LensRadius + Radius2 * Radius2) : HalfThickness2 = 0;
*/
	// TODO

}

TArray<FString> UOptiXLensComponent::GetGlassDefinitionNames()
{
	TArray<FString> Names;

	for (const TPair<FString, FGlassDefinition>& Pair : FOptiXModule::Get().GetGlassDefinitions())
	{
		Names.Add(Pair.Key);
	}

	return Names;
}

void UOptiXLensComponent::OnWavelengthChangedEvent(float WL)
{
	SetWavelength(WL);
}