#include "macros.h"
#include "const.h"

__device__ void Index_lattice(int &x, int &y, int &k) {
	// Only for D2
	x = threadIdx.x + NT * blockIdx.x;
	y = blockIdx.y;
	k = y * NX + x;
}

__device__ void CollisionKernel(LBtype* g, LBtype* m) {
	// Only for D2Q9
	LBtype p, u, v;
    LBtype niu = 1.0 / (0.5 + 3 * VIS);

	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;
	p = (g[0]+g[1]+g[2]+g[3]+g[4]+g[5]+g[6]+g[7]+g[8])/3.0;

	m[0]=   g[0]  +g[1]  +g[2]  +g[3]  +g[4]  +g[5]  +g[6]  +g[7]  +g[8];
	m[1]=-4*g[0]  -g[1]  -g[2]  -g[3]  -g[4]+2*g[5]+2*g[6]+2*g[7]+2*g[8];
	m[2]= 4*g[0]-2*g[1]-2*g[2]-2*g[3]-2*g[4]  +g[5]  +g[6]  +g[7]  +g[8];
	m[3]=          g[1]         -g[3]         +g[5]  -g[6]  -g[7]  +g[8];
	m[4]=       -2*g[1]       +2*g[3]         +g[5]  -g[6]  -g[7]  +g[8];
	m[5]=                 g[2]         -g[4]  +g[5]  +g[6]  -g[7]  -g[8];
	m[6]=              -2*g[2]       +2*g[4]  +g[5]  +g[6]  -g[7]  -g[8];
	m[7]=          g[1]  -g[2]  +g[3]  -g[4];
	m[8]=                                      g[5]  -g[6]  +g[7]  -g[8];	
	
	m[0]=1.00*( 3*p 	            -m[0]);
	m[1]=1.20*(-6*p +3*RHO*(u*u+v*v)-m[1]);
	m[2]=1.20*( 3*p -3*RHO*(u*u+v*v)-m[2]);
	m[3]=1.00*( RHO*u               -m[3]);
	m[4]=1.20*(-RHO*u               -m[4]);
	m[5]=1.00*( RHO*v               -m[5]);
	m[6]=1.20*(-RHO*v               -m[6]);
	m[7]= niu*( RHO*(u*u-v*v)       -m[7]);
	m[8]= niu*( RHO*u*v             -m[8]);

	g[0]=g[0]+(  m[0]  -m[1]  +m[2]                                   )/ 9.0;
	g[1]=g[1]+(4*m[0]  -m[1]-2*m[2]+6*m[3]-6*m[4]              +9*m[7])/36.0;
	g[2]=g[2]+(4*m[0]  -m[1]-2*m[2]              +6*m[5]-6*m[6]-9*m[7])/36.0;
	g[3]=g[3]+(4*m[0]  -m[1]-2*m[2]-6*m[3]+6*m[4]              +9*m[7])/36.0;
	g[4]=g[4]+(4*m[0]  -m[1]-2*m[2]              -6*m[5]+6*m[6]-9*m[7])/36.0;
	g[5]=g[5]+(4*m[0]+2*m[1]  +m[2]+6*m[3]+3*m[4]+6*m[5]+3*m[6]+9*m[8])/36.0;
	g[6]=g[6]+(4*m[0]+2*m[1]  +m[2]-6*m[3]-3*m[4]+6*m[5]+3*m[6]-9*m[8])/36.0;
	g[7]=g[7]+(4*m[0]+2*m[1]  +m[2]-6*m[3]-3*m[4]-6*m[5]-3*m[6]+9*m[8])/36.0;
	g[8]=g[8]+(4*m[0]+2*m[1]  +m[2]+6*m[3]+3*m[4]-6*m[5]-3*m[6]-9*m[8])/36.0;
}

__device__ void ParabolicInlet(LBtype* f, LBtype* f_neb, LBtype y) {
	LBtype p, u, v, yy;
	LBtype feq1, feq5, feq8, feqn1, feqn5, feqn8;

	p=(f_neb[0]+f_neb[1]+f_neb[2]+f_neb[3]+f_neb[4]+f_neb[5]+f_neb[6]+f_neb[7]+f_neb[8])/3.0;
	yy=(y-0.5*(NY-1))/(NY-2.0);
	u=U0*1.5*(1-4*yy*yy);
	v=0.0;

	feq1=(2*p+RHO*(2*u*u+2*u    -v*v)  )/ 6.0;
	feq5=(  p+RHO*(  u*u+3*u*v+u+v*v+v))/12.0;
	feq8=(  p+RHO*(  u*u-3*u*v+u+v*v-v))/12.0;

	u=(f_neb[1]+f_neb[5]+f_neb[8]-f_neb[3]-f_neb[6]-f_neb[7])/RHO;
	v=(f_neb[2]+f_neb[5]+f_neb[6]-f_neb[4]-f_neb[7]-f_neb[8])/RHO;

	feqn1=(2*p+RHO*(2*u*u+2*u    -v*v)  )/ 6.0;
	feqn5=(  p+RHO*(  u*u+3*u*v+u+v*v+v))/12.0;
	feqn8=(  p+RHO*(  u*u-3*u*v+u+v*v-v))/12.0;

	f[1]=f_neb[1]-feqn1+feq1;
	f[5]=f_neb[5]-feqn5+feq5;
	f[8]=f_neb[8]-feqn8+feq8;
}

__device__ void PressureOutlet(LBtype* f, LBtype* f_neb, LBtype y) {
	// Edit to Parabolic Outlet temporarily
	LBtype p, u, v, yy;	
	LBtype feq3, feq6, feq7, feqn3, feqn6, feqn7;

	p=0.0;

	yy=(y-0.5*(NY-1))/(NY-2.0);
	u=U0*1.5*(1-4*yy*yy);
	v=0.0;

	feq3=(2*p-RHO*(-2*u*u+2*u    +v*v)  )/ 6.0;
	feq6=(  p+RHO*(   u*u-3*u*v-u+v*v+v))/12.0;
	feq7=(  p+RHO*(   u*u+3*u*v-u+v*v-v))/12.0;

	u=(f_neb[1]+f_neb[5]+f_neb[8]-f_neb[3]-f_neb[6]-f_neb[7])/RHO;
	v=(f_neb[2]+f_neb[5]+f_neb[6]-f_neb[4]-f_neb[7]-f_neb[8])/RHO;
	// p=(f_neb[0]+f_neb[1]+f_neb[2]+f_neb[3]+f_neb[4]+f_neb[5]+f_neb[6]+f_neb[7]+f_neb[8])/3.0;
	feqn3=(2*p-RHO*(-2*u*u+2*u    +v*v)  )/ 6.0;
	feqn6=(  p+RHO*(   u*u-3*u*v-u+v*v+v))/12.0;
	feqn7=(  p+RHO*(   u*u+3*u*v-u+v*v-v))/12.0;

	f[3]=f_neb[3]-feqn3+feq3;
	f[6]=f_neb[6]-feqn6+feq6;
	f[7]=f_neb[7]-feqn7+feq7;
}