#include "VoxelForgeLivingPlaneActor.h"

#include "Engine/TextureRenderTarget2D.h"
#include "Engine/TextureRenderTarget.h"
#include "Engine/StaticMesh.h"
#include "Engine/World.h"
#include "Materials/MaterialInstanceDynamic.h"
#include "Materials/MaterialInterface.h"
#include "RenderGraphBuilder.h"
#include "RenderGraphUtils.h"
#include "RHICommandList.h"
#include "RHIResources.h"
#include "DataDrivenShaderPlatformInfo.h"
#include "GlobalShader.h"
#include "ShaderParameterStruct.h"
#include "UObject/ConstructorHelpers.h"
#include "Components/StaticMeshComponent.h"

namespace
{
	constexpr int32 LivingPlaneThreadGroupSize = 8;

	struct FLivingPlaneStepInput
	{
		FVector4f ObjectPositionRadius[2] = {FVector4f(0, 0, 0, -1), FVector4f(0, 0, 0, -1)};
		FVector4f ObjectVelocityPush[2] = {FVector4f(0, 0, 0, 0), FVector4f(0, 0, 0, 0)};
		float TimeSeconds = 0.0f;
	};

	class FVFLivingPlaneStepCS final : public FGlobalShader
	{
		DECLARE_GLOBAL_SHADER(FVFLivingPlaneStepCS);
		SHADER_USE_PARAMETER_STRUCT(FVFLivingPlaneStepCS, FGlobalShader);

		BEGIN_SHADER_PARAMETER_STRUCT(FParameters, )
			SHADER_PARAMETER(FIntPoint, GridSize)
			SHADER_PARAMETER(float, Dt)
			SHADER_PARAMETER(float, Dx)
			SHADER_PARAMETER(float, PlaneLength)
			SHADER_PARAMETER(float, WaveSpeed)
			SHADER_PARAMETER(float, Damping)
			SHADER_PARAMETER(float, FoundationStiffness)
			SHADER_PARAMETER(float, TimeSeconds)
			SHADER_PARAMETER(FVector2f, BreathingSettings)
			SHADER_PARAMETER(FVector4f, ImprintSettings)
			SHADER_PARAMETER_ARRAY(FVector4f, ObjectPositionRadius, [2])
			SHADER_PARAMETER_ARRAY(FVector4f, ObjectVelocityPush, [2])
			SHADER_PARAMETER_RDG_TEXTURE_SRV(Texture2D<float4>, SourceState)
			SHADER_PARAMETER_RDG_TEXTURE_UAV(RWTexture2D<float4>, TargetState)
		END_SHADER_PARAMETER_STRUCT()

	public:
		static bool ShouldCompilePermutation(const FGlobalShaderPermutationParameters& Parameters)
		{
			return IsFeatureLevelSupported(Parameters.Platform, ERHIFeatureLevel::SM5);
		}
	};

	IMPLEMENT_GLOBAL_SHADER(FVFLivingPlaneStepCS,
		"/Plugin/VoxelForgeLivingPlane/LivingPlane.usf", "MainCS", SF_Compute)

	UTextureRenderTarget2D* CreateFieldTarget(UObject* Outer, int32 Size, FName Name)
	{
		UTextureRenderTarget2D* Target = NewObject<UTextureRenderTarget2D>(Outer, Name, RF_Transient);
		Target->RenderTargetFormat = RTF_RGBA32f;
		Target->bCanCreateUAV = true;
		Target->bAutoGenerateMips = false;
		Target->ClearColor = FLinearColor::Transparent;
		Target->Filter = TF_Bilinear;
		Target->AddressX = TA_Clamp;
		Target->AddressY = TA_Clamp;
		Target->InitAutoFormat(Size, Size);
		Target->UpdateResourceImmediate(true);
		return Target;
	}
}

AVoxelForgeLivingPlaneActor::AVoxelForgeLivingPlaneActor()
{
	PrimaryActorTick.bCanEverTick = true;
	PlaneMesh = CreateDefaultSubobject<UStaticMeshComponent>(TEXT("LivingPlaneMesh"));
	SetRootComponent(PlaneMesh);
	PlaneMesh->SetCollisionEnabled(ECollisionEnabled::NoCollision);
	PlaneMesh->SetMobility(EComponentMobility::Movable);
	static ConstructorHelpers::FObjectFinder<UStaticMesh> PlaneAsset(
		TEXT("/Engine/BasicShapes/Plane.Plane"));
	if (PlaneAsset.Succeeded())
	{
		PlaneMesh->SetStaticMesh(PlaneAsset.Object);
	}

	ObjectInputs.SetNum(2);
}

void AVoxelForgeLivingPlaneActor::OnConstruction(const FTransform& Transform)
{
	Super::OnConstruction(Transform);
	if (PlaneMesh)
	{
		// /Engine/BasicShapes/Plane is 100 cm square; one Unreal scale unit
		// therefore corresponds to one metre of authored plane length.
		PlaneMesh->SetRelativeScale3D(FVector(PlaneLengthMeters, PlaneLengthMeters, 1.0));
	}
}

float AVoxelForgeLivingPlaneActor::GetStableStepSeconds() const
{
	const float Dx = PlaneLengthMeters / static_cast<float>(FMath::Max(GridSize, 16));
	const float EigenFrequencySquared = 8.0f * WaveSpeedMetersPerSecond * WaveSpeedMetersPerSecond
		/ FMath::Max(Dx * Dx, 1.0e-8f) + FoundationStiffnessPerSecondSquared;
	return 1.30f / FMath::Sqrt(FMath::Max(EigenFrequencySquared, 1.0e-8f));
}

void AVoxelForgeLivingPlaneActor::BeginPlay()
{
	Super::BeginPlay();
	GridSize = FMath::Clamp(GridSize, 16, 2048);
	PlaneLengthMeters = FMath::Max(PlaneLengthMeters, 1.0f);
	FixedStepSeconds = GetStableStepSeconds();
	CreateFieldTargets();

	if (PlaneMaterial && PlaneMesh)
	{
		DynamicPlaneMaterial = PlaneMesh->CreateDynamicMaterialInstance(0, PlaneMaterial);
		if (DynamicPlaneMaterial)
		{
			DynamicPlaneMaterial->SetScalarParameterValue(TEXT("PlaneLengthMeters"), PlaneLengthMeters);
			DynamicPlaneMaterial->SetScalarParameterValue(TEXT("FieldSize"), static_cast<float>(GridSize));
		}
	}
	UpdateMaterialField();
}

void AVoxelForgeLivingPlaneActor::CreateFieldTargets()
{
	FieldTargets[0] = CreateFieldTarget(this, GridSize, TEXT("LivingPlaneStateA"));
	FieldTargets[1] = CreateFieldTarget(this, GridSize, TEXT("LivingPlaneStateB"));
	CurrentFieldIndex = 0;
	SimulationStepCount = 0;
	SimulationAccumulatorSeconds = 0.0;
}

void AVoxelForgeLivingPlaneActor::SetObjectState(int32 ObjectIndex,
	const FVFLivingPlaneObjectInput& State)
{
	if (!ObjectInputs.IsValidIndex(ObjectIndex) || ObjectIndex >= 2)
	{
		return;
	}
	ObjectInputs[ObjectIndex] = State;
}

void AVoxelForgeLivingPlaneActor::ResetSimulation()
{
	if (!IsInGameThread() || !FieldTargets[0] || !FieldTargets[1])
	{
		return;
	}
	CreateFieldTargets();
	FixedStepSeconds = GetStableStepSeconds();
	UpdateMaterialField();
}

void AVoxelForgeLivingPlaneActor::UpdateMaterialField()
{
	if (DynamicPlaneMaterial && FieldTargets[CurrentFieldIndex])
	{
		DynamicPlaneMaterial->SetTextureParameterValue(TEXT("LivingField"), FieldTargets[CurrentFieldIndex]);
	}
}

void AVoxelForgeLivingPlaneActor::EnqueueSimulationSteps(int32 NumSteps)
{
	if (NumSteps <= 0 || !FieldTargets[0] || !FieldTargets[1])
	{
		return;
	}

	FTextureRenderTargetResource* ResourceA = FieldTargets[0]->GameThread_GetRenderTargetResource();
	FTextureRenderTargetResource* ResourceB = FieldTargets[1]->GameThread_GetRenderTargetResource();
	if (!ResourceA || !ResourceB)
	{
		return;
	}

	TArray<FLivingPlaneStepInput> Steps;
	Steps.Reserve(NumSteps);
	const FTransform PlaneTransform = GetActorTransform();
	for (int32 Step = 0; Step < NumSteps; ++Step)
	{
		FLivingPlaneStepInput& Input = Steps.AddDefaulted_GetRef();
		Input.TimeSeconds = static_cast<float>(SimulationStepCount + Step) * FixedStepSeconds;
		for (int32 ObjectIndex = 0; ObjectIndex < 2; ++ObjectIndex)
		{
			if (!ObjectInputs.IsValidIndex(ObjectIndex) || !ObjectInputs[ObjectIndex].bActive)
			{
				continue;
			}

			const FVFLivingPlaneObjectInput& Object = ObjectInputs[ObjectIndex];
			const FVector LocalPosition = PlaneTransform.InverseTransformPosition(Object.WorldPositionCm);
			const FVector LocalVelocity = PlaneTransform.InverseTransformVector(Object.WorldVelocityCmPerSecond);
			const float RadiusMeters = FMath::Max(Object.RadiusCm * 0.01f, 0.001f);
			Input.ObjectPositionRadius[ObjectIndex] = FVector4f(
				static_cast<float>(LocalPosition.X * 0.01),
				static_cast<float>(LocalPosition.Y * 0.01),
				static_cast<float>(LocalPosition.Z * 0.01),
				RadiusMeters);
			Input.ObjectVelocityPush[ObjectIndex] = FVector4f(
				static_cast<float>(LocalVelocity.X * 0.01),
				static_cast<float>(LocalVelocity.Y * 0.01),
				static_cast<float>(LocalVelocity.Z * 0.01),
				FMath::Max(Object.PushAccelerationMetersPerSecondSquared, 0.0f));
		}
	}

	const int32 Resolution = GridSize;
	const int32 InitialFieldIndex = CurrentFieldIndex;
	const float Dt = static_cast<float>(FixedStepSeconds);
	const float Dx = PlaneLengthMeters / static_cast<float>(Resolution);
	const float PlaneLength = PlaneLengthMeters;
	const float WaveSpeed = WaveSpeedMetersPerSecond;
	const float Damping = DampingPerSecond;
	const float Foundation = FoundationStiffnessPerSecondSquared;
	const float BreathingStrength = bEnableBreathing ? BreathingStrengthMetersPerSecondSquared : 0.0f;
	const FVector2f BreathingSettings(bEnableBreathing ? 1.0f : 0.0f, BreathingStrength);
	const FVector4f ImprintSettings(bEnableImprint ? 1.0f : 0.0f,
		ImprintCaptureMetersPerSecond, FMath::Max(ImprintRecoverySeconds, 1.0e-4f),
		FMath::Max(ImprintMaxDepthMeters, 0.0f));

	ENQUEUE_RENDER_COMMAND(VoxelForgeLivingPlaneStep)(
		[ResourceA, ResourceB, Steps = MoveTemp(Steps), InitialFieldIndex, Resolution, Dt, Dx, PlaneLength,
			WaveSpeed, Damping, Foundation, BreathingSettings, ImprintSettings]
		(FRHICommandListImmediate& RHICmdList) mutable
		{
			FRDGBuilder GraphBuilder(RHICmdList, RDG_EVENT_NAME("VoxelForgeLivingPlane"));
			FRDGTextureRef Fields[2] = {
			RegisterExternalTexture(GraphBuilder, ResourceA->GetRenderTargetTexture().GetReference(), TEXT("LivingPlaneStateA")),
			RegisterExternalTexture(GraphBuilder, ResourceB->GetRenderTargetTexture().GetReference(), TEXT("LivingPlaneStateB"))};
			TShaderMapRef<FVFLivingPlaneStepCS> ComputeShader(GetGlobalShaderMap(GMaxRHIFeatureLevel));
			int32 CurrentIndex = InitialFieldIndex;
			for (const FLivingPlaneStepInput& StepInput : Steps)
			{
				const int32 TargetIndex = 1 - CurrentIndex;
				FVFLivingPlaneStepCS::FParameters* PassParameters =
					GraphBuilder.AllocParameters<FVFLivingPlaneStepCS::FParameters>();
				PassParameters->GridSize = FIntPoint(Resolution, Resolution);
				PassParameters->Dt = Dt;
				PassParameters->Dx = Dx;
				PassParameters->PlaneLength = PlaneLength;
				PassParameters->WaveSpeed = WaveSpeed;
				PassParameters->Damping = Damping;
				PassParameters->FoundationStiffness = Foundation;
				PassParameters->TimeSeconds = StepInput.TimeSeconds;
				PassParameters->BreathingSettings = BreathingSettings;
				PassParameters->ImprintSettings = ImprintSettings;
				for (int32 ObjectIndex = 0; ObjectIndex < 2; ++ObjectIndex)
				{
					PassParameters->ObjectPositionRadius[ObjectIndex] = StepInput.ObjectPositionRadius[ObjectIndex];
					PassParameters->ObjectVelocityPush[ObjectIndex] = StepInput.ObjectVelocityPush[ObjectIndex];
				}
				PassParameters->SourceState = GraphBuilder.CreateSRV(Fields[CurrentIndex]);
				PassParameters->TargetState = GraphBuilder.CreateUAV(Fields[TargetIndex]);
				FComputeShaderUtils::AddPass(GraphBuilder,
					RDG_EVENT_NAME("LivingPlaneStep"),
					ComputeShader, PassParameters,
					FIntVector(FMath::DivideAndRoundUp(Resolution, LivingPlaneThreadGroupSize),
						FMath::DivideAndRoundUp(Resolution, LivingPlaneThreadGroupSize), 1));
				CurrentIndex = TargetIndex;
			}
			GraphBuilder.Execute();
		});

	CurrentFieldIndex = (InitialFieldIndex + NumSteps) & 1;
	SimulationStepCount += NumSteps;
	UpdateMaterialField();
}

void AVoxelForgeLivingPlaneActor::Tick(float DeltaSeconds)
{
	Super::Tick(DeltaSeconds);
	if (FixedStepSeconds <= 0.0 || !FieldTargets[0] || !FieldTargets[1])
	{
		return;
	}

	SimulationAccumulatorSeconds += FMath::Max(static_cast<double>(DeltaSeconds), 0.0);
	const int32 StepsToRun = FMath::Min(
		static_cast<int32>(SimulationAccumulatorSeconds / FixedStepSeconds), 32);
	if (StepsToRun > 0)
	{
		SimulationAccumulatorSeconds -= StepsToRun * FixedStepSeconds;
		EnqueueSimulationSteps(StepsToRun);
	}
}

void AVoxelForgeLivingPlaneActor::EndPlay(const EEndPlayReason::Type EndPlayReason)
{
	Super::EndPlay(EndPlayReason);
	DynamicPlaneMaterial = nullptr;
	FieldTargets[0] = nullptr;
	FieldTargets[1] = nullptr;
}