// CelerisLab/kernels/kernel.cu

#include <stdio.h>
#include <stdint.h>
#include <cuda.h>

#include "macros.h"
#include "const.h"
#include "D2Q9.cu"

extern "C"
{
	__global__ void OneStep(uint8_t *flag, LBtype *f, LBtype *f_temp, int32_t *indx, LBtype *delta, LBtype *action, LBtype *obs, uint32_t *error_flag)
	{
		__shared__ LBtype f_share[NT * NQ];
		__shared__ LBtype obs_share[(N_OBJS * DIM > 0) ? N_OBJS * DIM : 1];

		int x, y, k;
		LBtype g[NQ], m[NQ];
		Index_lattice(x, y, k); // Only for D2
		int totalCells = NX * NY;
		int id = indx[k];

		for (int i = 0; i < NQ; i++)
		{
			f_share[threadIdx.x + i * NT] = f[k + i * totalCells];
		}
		for (int i = threadIdx.x; i < N_OBJS * DIM; i += NT)
		{
			obs_share[i] = 0;
		}

		__syncthreads();

		for (int i = 0; i < NQ; i++)
		{
			g[i] = f_share[threadIdx.x + i * NT];
		}

		if (flag[k] & FLUID)
		{
			CollisionKernel(g, m);

			for (int i = 0; i < NQ; i++)
			{
				if (isnan((double)g[i]) || isinf((double)g[i]))
				{
					atomicOr(error_flag, (uint32_t)1);
				}
				f_share[threadIdx.x + i * NT] = g[i];
			}
		}
		else if (flag[k] & SOLID)
		{
			if (x == 0)
			{
				for (int i = 0; i < NQ; i++)
				{
					m[i] = f_share[threadIdx.x + i * NT + 1];
				}
				ParabolicInlet(g, m, y);
			}
			else if (x == NX - 1)
			{
				for (int i = 0; i < NQ; i++)
				{
					m[i] = f_share[threadIdx.x + i * NT - 1];
				}
				PressureOutlet(g, m, y);
			}

			for (int i = 0; i < NQ; i++)
			{
				if (isnan((double)g[i]) || isinf((double)g[i]))
				{
					atomicOr(error_flag, (uint32_t)1);
				}
				f_share[threadIdx.x + i * NT] = g[i];
			}
		}

		__syncthreads();

		for (int i = 0; i < NQ; i++)
		{
			int x_neb = x + e[i][0];
			int y_neb = y + e[i][1];

			if (y != 0 && y != NY - 1)
			{
				if ((y == 1 && y_neb == 0) || (y == NY - 2 && y_neb == NY - 1))
				{
					f_temp[k + opp[i] * totalCells] = f_share[threadIdx.x + i * NT];
				}
				else
				{
					int k_neb = ((y_neb * NX + x_neb) + totalCells) % totalCells;
					f_temp[k_neb + i * totalCells] = f_share[threadIdx.x + i * NT];
				}
			}
		}

		__syncthreads();

		if (flag[k] & SOLID && flag[k] & INTERFACE)
		{
			LBtype Uw, Vw;
			int id_obj = *reinterpret_cast<int *>(&delta[id]);
			Uw = action[id_obj] * delta[id + 9];
			Vw = action[id_obj] * delta[id + 10];

			int x_neb, y_neb, k_neb;
			for (int i = 1; i < 9; i++)
			{
				x_neb = x + e[i][0];
				y_neb = y + e[i][1];
				k_neb = x_neb + y_neb * NX;
				if (flag[k_neb] & FLUID)
				{
					LBtype q = delta[id + i];
					int k_neb2 = (y + 2 * e[i][1]) * NX + (x + 2 * e[i][0]);
					LBtype temp = 6 * w[i] * (e[i][0] * Uw + e[i][1] * Vw);
					f_temp[k_neb + i * totalCells] = (q * f_temp[k + opp[i] * totalCells] \
						+ (1 - q) * f_temp[k_neb + opp[i] * totalCells] \
						+ q * f_temp[k_neb2 + i * totalCells] + temp) / (1 + q);
					f_temp[k + i * totalCells] = temp * Uw;
					k_neb2 = (y - e[i][1]) * NX + (x - e[i][0]);
					f_temp[k_neb2 + i * totalCells] = temp * Vw;

					temp = f_temp[k_neb + i * totalCells] + f_temp[k + opp[i] * totalCells];
					k_neb2 = (y - e[i][1]) * NX + (x - e[i][0]);
					atomicAdd(&obs_share[DIM * id_obj], -temp * e[i][0] + f_temp[k + i * totalCells]);
					atomicAdd(&obs_share[DIM * id_obj + 1], -temp * e[i][1] + f_temp[k_neb2 + i * totalCells]);
				}
			}
		}
		if (flag[k] & SENSOR)
		{
			LBtype u, v;
			u = (g[1] + g[5] + g[8] - g[3] - g[6] - g[7]) / RHO;
			v = (g[2] + g[5] + g[6] - g[4] - g[7] - g[8]) / RHO;
			if (isnan((double)u) || isinf((double)u) || isnan((double)v) || isinf((double)v))
			{
				atomicOr(error_flag, (uint32_t)1);
			}
			atomicAdd(&obs_share[DIM * id], u);
			atomicAdd(&obs_share[DIM * id + 1], v);
		}

		__syncthreads();

		for (int i = threadIdx.x; i < N_OBJS * DIM; i += NT)
		{
			atomicAdd(&obs[i], obs_share[i]);
		}
	}

	__global__ void InitTubeFlow(uint8_t *flag, LBtype *f)
	{
		__shared__ LBtype f_share[NT * NQ];
		__shared__ uint8_t flag_share[NT];
		int x, y, k;
		LBtype u;
		Index_lattice(x, y, k);
		int totalCells = NX * NY;

		flag_share[threadIdx.x] = flag[k];
		for (int i = 0; i < NQ; i++)
		{
			f_share[threadIdx.x + i * NT] = f[k + i * totalCells];
		}

		__syncthreads();

		u = U0 * 1.5 * (1 - 4 * (y - 0.5 * (NY - 1)) * (y - 0.5 * (NY - 1)) / ((NY - 2) * (NY - 2)));
		if (y == 0 || y == NY - 1 || x == 0 || x == NX - 1)
		{
			flag_share[threadIdx.x] = SOLID;
			for (int i = 0; i < NQ; i++)
			{
				f_share[threadIdx.x + i * NT] = 0;
			}
		}
		else
		{
			flag_share[threadIdx.x] = FLUID;
			for (int i = 0; i < NQ; i++)
			{
				f_share[threadIdx.x + i * NT] = w[i] * RHO * (3 * e[i][0] * u + \
					4.5 * e[i][0] * e[i][0] * u * u - 1.5 * u * u);
			}
		}

		__syncthreads();

		flag[k] = flag_share[threadIdx.x];
		for (int i = 0; i < NQ; i++)
		{
			f[k + i * totalCells] = f_share[threadIdx.x + i * NT];
		}
	}

	// __global__ void AddVortex(LBtype *f, int32_t *config)
	// {
	// 	__shared__ LBtype f_share[NT * NQ];
	// 	int x, y, k;
	// 	LBtype u, v, u_vor, v_vor;
	// 	Index_lattice(x, y, k);
	// 	int totalCells = NX * NY;

	// 	for (int i = 0; i < NQ; i++)
	// 	{
	// 		f_share[threadIdx.x + i * NT] = f[k + i * totalCells];
	// 	}

	// 	__syncthreads();

	// 	u = f_share[threadIdx.x + 1 * NT] - f_share[threadIdx.x + 3 * NT] + f_share[threadIdx.x + 5 * NT] - f_share[threadIdx.x + 6 * NT] - f_share[threadIdx.x + 7 * NT] + f_share[threadIdx.x + 8 * NT];
	// 	v = f_share[threadIdx.x + 2 * NT] - f_share[threadIdx.x + 4 * NT] + f_share[threadIdx.x + 5 * NT] + f_share[threadIdx.x + 6 * NT] - f_share[threadIdx.x + 7 * NT] - f_share[threadIdx.x + 8 * NT];

	// 	if type & V_TAYLOR
	// 	{
	// 		u_vor = -2 * PI * U0 * sin(2 * PI * x / NX) * sin(2 * PI * y / NY);
	// 		v_vor = 2 * PI * U0 * cos(2 * PI * x / NX) * cos(2 * PI * y / NY);
	// 	}
	// 	else
	// 	{
	// 		u_vor = 0;
	// 		v_vor = 0;
	// 	}


	// }
}