// Copyright 2023 ZSttt, Inc. All Rights Reserved.

#include "/Engine/Public/Platform.ush"
#include "/Engine/Private/Common.ush"

RWStructuredBuffer<float2> RWParticles;
Texture2D FlowFieldMap;
SamplerState FlowFieldMapSampler;
uint2 FlowFieldMapResolution;
uint NumParticles;
float RandomSeed;
uint NumParticlesSqrt;
float4 SpeedRange;
float SpeedFactor;
float RandomRate;
float DropRate;
float Spherical;

static const float3 rand_constants = float3(12.9898, 78.233, 4375.85453);
float rand(float2 uv) 
{
    float t = dot(rand_constants.xy, uv);
    return frac(sin(t) * (rand_constants.z + t));
}

float4 GetFlowFieldValue(float2 UV)
{
	//UE Transform
    float2 pos = float2(UV.x,1.0 - UV.y);

    return Texture2DSample(FlowFieldMap, FlowFieldMapSampler, pos);
}

float2 lookup(float2 uv)
{
    float2 px = 1.0 / (float) FlowFieldMapResolution;
    float2 vc = (floor(uv * (float) FlowFieldMapResolution)) * px;
    float2 f = frac(uv * (float) FlowFieldMapResolution);
    float2 tl = GetFlowFieldValue(vc).xy;
    float2 tr = GetFlowFieldValue(vc + float2(px.x, 0)).xy;
    float2 bl = GetFlowFieldValue(vc + float2(0, px.y)).xy;
    float2 br = GetFlowFieldValue(vc + px).xy;
    return lerp(lerp(tl, tr, f.x), lerp(bl, br, f.x), f.y);
}

[numthreads(256, 1, 1)]
void UpdateParticles
(
	uint3 GroupID : SV_GroupID,
	uint3 ThreadId : SV_DispatchThreadID,
	uint3 GroupThreadID : SV_GroupThreadID
)
{
    uint Index = ThreadId.x + ThreadId.y * 256;
    if (Index < NumParticles)
    {
        float2 particle = RWParticles[Index];
        float4 color = GetFlowFieldValue(particle);
        if (color.a == 0.0)
        {
            particle = frac(1.0 + particle + (RandomSeed * particle));
            float Row = Index / NumParticlesSqrt;
            float Column = Index % NumParticlesSqrt;
            float2 seed = (particle + (float2(Column, Row) / (float) NumParticlesSqrt)) * RandomSeed;
            particle = float2(rand(seed + 1.3), rand(seed + 2.1));
        }
        else
        {
            float2 velocity = lerp(SpeedRange.xy, SpeedRange.zw, lookup(particle));
            float speed = length(velocity) / length(SpeedRange.ba);
		
            float2 offset = 0.f;
            if (Spherical > 0.f)
            {
                float distortion = cos(radians(particle.y * 180.0 - 90.0));
                offset = float2(velocity.x / distortion, velocity.y) * 0.0001 * SpeedFactor;
            }
            else
            {
                offset = float2(velocity.x, velocity.y) * 0.0001 * SpeedFactor;
            }

            particle = frac(1.0 + particle + offset);
		
            float Row = Index / NumParticlesSqrt;
            float Column = Index % NumParticlesSqrt;
            float2 seed = (particle + (float2(Column, Row) / (float) NumParticlesSqrt)) * RandomSeed;
		
            float drop_rate = RandomRate + speed * DropRate;
            float drop = step(1.0 - drop_rate, rand(seed));
		
            float2 random_pos = float2(rand(seed + 1.3), rand(seed + 2.1));
            particle = lerp(particle, random_pos, drop);
        }
        RWParticles[Index].xy = particle;
    }
}