#pragma kernel CSMain

#include "UnityCG.cginc"
#include "../../ParticleWorks/Shader/Inc/Defines.cginc"
#include "../../ParticleWorks/Shader/Inc/Math.cginc"
#include "../../ParticleWorks/Shader/Noise/noise3Dgrad.cginc"
#include "../../ParticleWorks/Shader/ParticleUniforms.cginc"

CBUFFER_START(Variables)
	float timeStep;
	float timeValue;
	float bufferCount;
	//int home_count;
	float3 coeffsA;
	float3 coeffsB;
	float3 coeffsC; 
	float3 coeffsD;
	float4 FluidGridDim;
	float4 FluidGridRoot;
	float4 FluidGridSize;
CBUFFER_END

RWStructuredBuffer<Particle> ssbo;
RWStructuredBuffer<ParticleUniforms> uniform_buffer;
StructuredBuffer<int> home_count_buffer;
StructuredBuffer<float4> home_position_buffer;
StructuredBuffer<float4> home_velocity_buffer;
StructuredBuffer<float3> fluid_velocity_buffer;

[numthreads(WORK_GROUP_SIZE, 1, 1)]
void CSMain(uint3 id : SV_DispatchThreadID)
{
	int home_count = home_count_buffer[0];
	uint idx = id.x;
	Particle p = ssbo[idx];
	float4 home_position = (home_count > 0) ? home_position_buffer[p.seed * home_count] : float4(0, 0, 0, 1);
	float4 home_velocity = (home_count > 0) ? home_velocity_buffer[p.seed * home_count] : float4(0, 0, 0, 1);
	ParticleUniforms uni = uniform_buffer[0];

	float factor = p.color.a;
	float3 acc = float3(0, 0, 0);
	float3 F = float3(0, 0, 0);

	if (home_count > 0)
	{
		float duration = uni.lifeDuration * (1.0f + factor * uni.lifeVariation);
		float current_life = p.life * duration + timeStep;
		//bool is_reset = false;
		//if (current_life > duration)
		//{
		//	current_life -= duration;
		//	is_reset = true;
		//}
		//if (is_reset || p.life < 0.0f)
		if (current_life > duration && uni.lifeDuration > 0.1f)
		{
			current_life -= duration;
			//if (isEmit == 0 && current_life > 0.0) current_life -= duration;
			p.position = home_position.xyz;
			p.color.rgb = float3(1, 1, 1);
			p.velocity = home_velocity.xyz * 20.0f * pow(length(home_velocity.xyz), 2.0f);
			//p.velocity = forceDirection * 3.0f;
		}
		p.life = current_life / duration;
	}

	// noise
	{
		//float n_seed = p.seed * 0.12f;
		//float time_seed = timeValue * noiseTimeScale;
		float3 np = p.position * uni.noiseFrequency;
		float3 n1, n2;
		snoise(np, n1);
		snoise(np + float3(uni.noiseSeed, p.seed * 0.12f, timeValue * uni.noiseTimeScale), n2);
		F = cross(n1, n2);
		acc += F * uni.noiseStrength;
	}

	// convergence
	{
		F = -p.position;
		acc += F * uni.convergence;
	}
	
	// home as tangential
	{
		F = cross(p.position, float3(0, 1, 0));
		float d = clamp(abs(p.position.y), 1, 100);
		F = normalize(F) * 10.0f / pow(d, 2 * 2);
		acc += F * uni.homeStrength; 
	}

	// force
	{
		F = float3(0, 0, 0);
		if (length(uni.forceDirection) > 0.1f)
			F = normalize(uni.forceDirection) * 10.0f;
		acc += F * uni.forceStrength;
	}

	// impulse
	{
		//F = uni.impulsePosition - p.position;
		//float mag = F.x*F.x + F.y*F.y + F.z*F.z;
		//float rad2 = uni.impulseRadius * uni.impulseRadius;
		//acc += F * exp(-mag / rad2) * uni.impulseStrength;
	}

	// fluid
	{
		float3 dim = ((p.position - FluidGridRoot) / FluidGridSize) * FluidGridDim;
		dim = clamp(dim, float3(0, 0, 0), FluidGridDim - 1);
		int idx = (int)dim.x + (int)dim.y * FluidGridDim.x + (int)dim.z * FluidGridDim.x * FluidGridDim.y;
		float3 F = fluid_velocity_buffer[idx].xyz * 5.0f;
		acc += F * uni.impulseStrength;
	}
	
	
	acc *= uni.movement;
	p.velocity += acc * timeStep;
	p.position += p.velocity * timeStep;

	//float damping = lerp(0.965f, 0.985f, p.seed) * velocityDamping;
	p.velocity *= lerp(0.965f, 0.985f, p.seed) * uni.velocityDamping;
	
	// calculate rotate quat
	{
		float vel = length(p.velocity);
		float a = (1 + factor) * vel * 0.05f;
		float4 q = quat_from_axis_angle(normalize(p.velocity + float3(0, FLT_EPSILON, 0)), a);
		q = quat_mult(p.quat, q);
		q = normalize(q);
		p.quat = lerp(p.quat, q, saturate(vel - 0.01f));
	}

	// calculate model matrix
	{
		float size = uni.geomSize * lerp(1, pow(1.0f - saturate(p.life), 1.0f / 2.2f), any(uni.lifeDuration));
		size *= lerp(0.5, 1.0, p.seed);
		if (idx > (uint)(uni.numRate * bufferCount) || (p.life < 0.0f && uni.lifeDuration > 0.1f))
		{
			size = 0.0f;
		}

		float4x4 mat = make_translation_matrix(p.position);
		mat = mul(mat, quat_to_rotation_matrix(p.quat));
		mat = mul(mat, make_scaling_matrix(size));
		p.model_matrix = mat;
	}

	// ground collision
	{
		float3 pos = mul(p.model_matrix, float4(p.position, 1.0f)).xyz;
		pos = p.position;
		if (pos.y < 0.0f)
		{
			p.velocity.y = abs(p.velocity.y) * 0.8f;
			p.position.y = 0.0f;
		}
	}

	// cosine gradient
	{
		float t = (uni.lifeDuration > 0.1f) ? clamp(p.life, 0.0f, 1.0f) : factor;
		half3 rgb = cosine_gradient(coeffsA, coeffsB, coeffsC, coeffsD, t);
	#if !defined(UNITY_COLORSPACE_GAMMA)
		rgb = GammaToLinearSpace(rgb);
	#endif
		p.color.rgb = rgb;
	}
	
	ssbo[id.x] = p;
}