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

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

Texture2D FlowFieldMap;
SamplerState FlowFieldMapSampler;

Texture2D ColorMap;
SamplerState ColorMapSampler;

float4 SpeedRange;

void DrawParticles
(
	in float2 Particle : TEXCOORD0,
	out float4 OutColor : SV_Target0
)
{
    float2 pos = Particle.xy;
    
    //UE Transform
    pos.y = 1.0 - pos.y;
    
    float4 value = Texture2DSample(FlowFieldMap, FlowFieldMapSampler, pos);
    if (value.a == 0.0)
    {
        OutColor = float4(0,0,0,0);
    }
    else
    {
        float2 velocity = lerp(SpeedRange.xy, SpeedRange.ba, value.xy);
        float speed = length(velocity) / length(SpeedRange.zw);
        OutColor = Texture2DSample(ColorMap, ColorMapSampler, float2(speed, 0.5));
    }
}

Texture2D<float4> Texture;
float FadeOpacity;
float FadeOutOpacity;
void DrawScreen
(
	in FScreenVertexOutput Input,
	out float4 OutColor : SV_Target0
)
{
    float4 Color = Texture.Load(uint3(Input.Position.xy, 0));
    
    if (Color.a < FadeOutOpacity)
        OutColor = float4(0,0,0,0);
    else
    {
        OutColor = Color;
        OutColor.a = floor(255.0 * Color.a * FadeOpacity) / 255.0;
    }
    //OutColor = floor(255.0 * Color * FadeOpacity) / 255.0;
}

Texture2D<float4> Particles;
uint2 FlowFieldMapResolution;
float RandomSeed;
uint2 ParticlesResolution;
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);
}

void UpdateParticles
(
	in FScreenVertexOutput Input,
	out float4 OutColor : SV_Target0
)
{
    float4 particle = Particles.Load(uint3(Input.Position.xy, 0));
    float2 pos = float2(particle.r / 255.0 + particle.b, particle.g / 255.0 + particle.a);
    float4 color = GetFlowFieldValue(pos);
    if (color.a == 0.0)
    {
        pos = frac(1.0 + pos + (RandomSeed * pos));
        float2 seed = (pos + Input.UV) * RandomSeed;
        pos = float2(rand(seed + 1.3), rand(seed + 2.1));
    }
    else
    {
        float2 velocity = lerp(SpeedRange.xy, SpeedRange.zw, lookup(pos));
        float speed = length(velocity) / length(SpeedRange.ba);
		
        float2 offset = 0.f;
        if (Spherical > 0.f)
        {
            float distortion = cos(radians(pos.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;
        }

        pos = frac(1.0 + pos + offset);

        float2 seed = (pos + Input.UV) * 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));
        pos = lerp(pos, random_pos, drop);
    }
		
    OutColor = float4(frac(pos * 255.0), floor(pos * 255.0) / 255.0);
}