#pragma kernel ResetKernel
#pragma kernel StepKernel

static int nthreads = 1;

RWTexture2D<float4> readC;
RWTexture2D<float4> readD;


RWTexture2D<float4> writeC;
RWTexture2D<float4> writeD;

int resolution;

float4 S(RWTexture2D<float4> target, uint2 id)
{
	return target[(id.xy + resolution)%resolution];
}

[numthreads(nthreads, nthreads, 1)]
void ResetKernel(uint3 id : SV_DispatchThreadID)
{

	float4 values = float4(0, 0, 0, 0);

	if (distance(id.xy, float2(resolution / 2.0, resolution / 2.0)) < 10)
	{
		values.xy = float2(0,.25);
	}

	readC[id.xy] = values;
	readD[id.xy] = float4(0, 0, 0, 1);

	writeC[id.xy] = values;
	writeD[id.xy] = float4(0, 0, 0, 1);
}

void Fluid(RWTexture2D<float4> target, uint2 id, float2 offset, float4 values, inout float2 velocity, inout float pressure, inout float divergence, inout float neighbors)
{
	float2 v = target[id.xy+offset].xy;
	float4 s = S(target, id.xy + offset.xy - v);

	offset = normalize(offset);

	velocity += offset * (s.w - values.w);

	pressure += s.w;

	divergence += dot(offset.xy, s.xy);

	++neighbors;
}

void StepVelocity(uint3 id, RWTexture2D<float4> write, RWTexture2D<float4> read, bool addJet)
{
	float4 values = S(read, id.xy);
	values = S(read, id.xy - values.xy);

	float2 velocity = float2(0,0);
	float pressure = 0;
	float neighbors = 0;
	float divergence = 0;

	Fluid(read, id.xy, float2(0, 1), values, velocity, pressure, divergence, neighbors);

	Fluid(read, id.xy, float2(0, -1), values, velocity, pressure, divergence, neighbors);

	Fluid(read, id.xy, float2(1, 0), values, velocity, pressure, divergence, neighbors);
	Fluid(read, id.xy, float2(-1, 0), values, velocity, pressure, divergence, neighbors);

	velocity = velocity / neighbors;
	divergence = divergence / neighbors;
	pressure = pressure/neighbors;

	values.w = pressure - divergence;
	values.xy -= velocity;

	if (addJet && distance(id.xy, float2(resolution / 2.0, resolution / 2.0)) < 10)
		values = float4(0,.25,0,values.w);

	write[id.xy] = values;

}

void StepFluid(uint3 id)
{
	float2 p = id.xy;

	for (int i = 0; i < 4; i++)
		p -= readC[p].xy;


	float4 values = S(readD, p);

	if (distance(id.xy, float2(resolution / 2.0, resolution / 2.0)) < 10)
		values = float4(0, 0, 1, 1);

	writeD[id.xy] = S(writeC, id.xy);

}

[numthreads(nthreads, nthreads, 1)]
void StepKernel(uint3 id : SV_DispatchThreadID)
{
	StepVelocity(id, writeC, readC, true);

	StepFluid(id);
}