# version 440

#define tau 0.91//0.6
#define omega (1.0/tau)		// viscosity, etc.
const int q = 15;
const int ex[q] = { 0, -1, 0, 0, -1, -1, -1, -1, 1, 0, 0, 1, 1, 1, 1 };
const int ey[q] = { 0, 0, -1, 0, -1, -1, 1, 1, 0, 1, 0, 1, 1, -1, -1 };
const int ez[q] = { 0, 0, 0, -1, -1, 1, -1, 1, 0, 0, 1, 1, -1, 1, -1 };
const int inv[q] = { 0, 8, 9, 10, 11, 12, 13, 14, 1, 2, 3, 4, 5, 6, 7 };
const float z = 2.0 / 9.0, l = 1.0 / 72.0, s = 1.0 / 9.0;
const float w[q] = { z, s, s, s, l, l, l, l, s, s, s, l, l, l, l };
#define C_FLD 0
#define C_BND 1


layout(binding = 5) buffer df0 { float f0 [  ]; };
layout(binding = 6) buffer df1 { float f1 [  ]; };
layout(binding = 7) buffer dcF { int F [  ]; };
layout(binding = 8) buffer dcU { float U [  ]; };
layout(binding = 9) buffer dcV { float V [  ]; };
layout(binding = 10) buffer dcW { float W [  ]; };
layout(local_size_x = 10, local_size_y = 10, local_size_z = 10) in;

int per(int x, int nx)      // periodic bnd's
{
	if (x < 0) x = nx;
	else if (x > nx) x = 0;
	return x;
}

void main()
{
	int NX = 200;
	int NY = 70;
	int NZ = 70;
	float devFx = 0.0002;
	float devFy = 0.0;
	float devFz = 0.0;

	int i = int(gl_GlobalInvocationID.x);
	int j = int(gl_GlobalInvocationID.y);
	int k = int(gl_GlobalInvocationID.z);
	int idx = i + j * NX + k * NX * NY;
	float feq;
	float rho = 0;
	float ux = 0;
	float uy = 0;
	float uz = 0;
	if (F[idx] == C_FLD)
	{
		for (int m = 0; m < q; m++)         // calculate density and velocity
		{
			rho = rho + f0[idx * q + m];
			ux = ux + f0[idx * q + m] * ex[m];
			uy = uy + f0[idx * q + m] * ey[m];
			uz = uz + f0[idx * q + m] * ez[m];
		}
		ux /= rho;
		uy /= rho;
		uz /= rho;
		U[idx] = ux;
		V[idx] = uy;
		W[idx] = uz;
		ux = ux + 0.5 * devFx;
		uy = uy + 0.5 * devFy;
		uz = uz + 0.5 * devFz;
		for (int m = 0; m < q; m++)     // collision + streaming (the main solver is here, really)
		{
			int ip = i + ex[m];
			int jp = j + ey[m];
			int kp = k + ez[m];

			ip = per(ip, NX - 1);
			jp = per(jp, NY - 1);
			kp = per(kp, NZ - 1);

			int idxp = ip + jp * NX + kp * NX * NY;

			if (F[idxp] == C_BND)
				f1[idx * q + inv[m]] = (1 - omega) * f0[idx * q + m] + omega * w[m] * rho * (1.0 - (3.0 / 2.0) * (ux * ux + uy * uy * uz * uz) + 3.0 * (ex[m] * ux + ey[m] * uy + ez[m] * uz) + (9.0 / 2.0) * (ex[m] * ux + ey[m] * uy + ez[m] * uz) * (ex[m] * ux + ey[m] * uy + ez[m] * uz));
			else
				f1[idxp * q + m] = (1 - omega) * f0[idx * q + m] + omega * w[m] * rho * (1.0 - (3.0 / 2.0) * (ux * ux + uy * uy * uz * uz) + 3.0 * (ex[m] * ux + ey[m] * uy + ez[m] * uz) + (9.0 / 2.0) * (ex[m] * ux + ey[m] * uy + ez[m] * uz) * (ex[m] * ux + ey[m] * uy + ez[m] * uz));
		}
	}
	/*else
	{
		// reset flags and set df to equilibrium if flag set to 0
		U[ idx ] = 0;
		V[ idx ] = 0;
		W[ idx ] = 0;
		for(int m=0; m<q; m++)		// collision + streaming (the main solver is here, really)
		{
			int ip = i+ex[m];
			int jp = j+ey[m];
			int kp = k+ez[m];

			ip=per(ip,NX-1);
			jp=per(jp,NY-1);
            		kp=per(kp,NZ-1);

			int idxp =  ip+jp*NX+kp*NX*NY;
			f1[ idxp*q + m ] = w[m] * rho * (1.0 - (3.0/2.0) * (ux*ux + uy*uy * uz*uz)+ 3.0 * (ex[m] * ux + ey[m]*uy + ez[m]*uz)+ (9.0/2.0) * (ex[m] * ux + ey[m]*uy + ez[m]*uz) * (ex[m] * ux + ey[m]*uy + ez[m]*uz));
		}
	}*/
}







