#include "XRD3D.cuh"
//该kernel函数根据fijk的公式构建其傅里叶变换结果
//其已经考虑近程作用衰减删除，和四阶B样条线插值干扰消除
//
__global__ void Build_fft_fijk_List(const int NzNy_Nx21_numbers, const INT_VECTOR grid_dimension,
	const INT_VECTOR grid_dimension_upper_half, const INT_VECTOR grid_dimension_lower_half_plus_one,
	const VECTOR box_length_inverse, const float s,const float k0,const float scalor,
	const cufftComplex *fft_B_spline, float *fft_fijk)
{
	int gpu_grid_layer_numbers = gridDim.x*blockDim.x;
	int layer_numbers = grid_dimension.int_y*grid_dimension_lower_half_plus_one.int_x;
	float s_inverse = 1. / s;
	float k0_square = k0*k0;

	INT_VECTOR grid_pos;
	INT_VECTOR grid_fft_serial;
	float Pi2ijk;//2.*Pi*|{i/a,j/b,k/c}|
	float Pi2ijk_square;//2.*Pi*|{i/a,j/b,k/c}|

	float rk02;
	float fijk_s;//s
	int temp;
	for (int grid_i = blockIdx.x*blockDim.x + threadIdx.x; grid_i < NzNy_Nx21_numbers; grid_i = grid_i + gpu_grid_layer_numbers)
	{
		grid_pos.int_x = grid_i%grid_dimension_lower_half_plus_one.int_x;
		grid_pos.int_y = (grid_i % layer_numbers) / grid_dimension_lower_half_plus_one.int_x;
		grid_pos.int_z = grid_i / layer_numbers;

		temp = grid_pos.int_x - grid_dimension_upper_half.int_x;
		temp = temp >> 31;
		grid_fft_serial.int_x = (temp& grid_pos.int_x) + ((~temp)&(grid_dimension.int_x - grid_pos.int_x));

		temp = grid_pos.int_y - grid_dimension_upper_half.int_y;
		temp = temp >> 31;
		grid_fft_serial.int_y = (temp& grid_pos.int_y) + ((~temp)&(grid_dimension.int_y - grid_pos.int_y));

		temp = grid_pos.int_z - grid_dimension_upper_half.int_z;
		temp = temp >> 31;
		grid_fft_serial.int_z = (temp& grid_pos.int_z) + ((~temp)&(grid_dimension.int_z - grid_pos.int_z));

		Pi2ijk = (2.*CONSTANT_Pi)*norm3df((float)grid_fft_serial.int_x*box_length_inverse.x,
			(float)grid_fft_serial.int_y*box_length_inverse.y,
			(float)grid_fft_serial.int_z*box_length_inverse.z);
		Pi2ijk_square = Pi2ijk*Pi2ijk;

		Pi2ijk_square = Pi2ijk_square + k0_square;
		rk02 = 2.*Pi2ijk*k0;
		//下面算法会溢出，因此要展开
		//fijk_s = 2. / Pi2ijk*expf(-0.25*s_inverse*(Pi2ijk_square + k0_square))*sinhf(0.5*s_inverse*Pi2ijk*k0);
		fijk_s = 1. / Pi2ijk*(expf(-0.25*s_inverse*(Pi2ijk_square - rk02)) - expf(-0.25*s_inverse*(Pi2ijk_square + rk02)));

		fft_fijk[grid_i] = scalor*fijk_s / (fft_B_spline[grid_i].x*fft_B_spline[grid_i].x);
	}
	fft_fijk[0] = 0.;//无脑将000置为0
}
__global__ void Array_float_Multiply_cufftComplex(const int array_numbers, const float *f, cufftComplex *c)
{
	int gpu_grid_layer_numbers = gridDim.x*blockDim.x;
	for (int i = blockIdx.x*blockDim.x + threadIdx.x; i < array_numbers; i = i + gpu_grid_layer_numbers)
	{
		c[i].x = f[i] * c[i].x;
		c[i].y = f[i] * c[i].y;
	}
}

//以下三个kernel函数实现xrd原子的四阶B样条线插值与反插值
//来源于PME模块修改
__constant__ float XRD3D_Ma[4] = { 1.0 / 6.0, -0.5, 0.5, -1.0 / 6.0 };
__constant__ float XRD3D_Mb[4] = { 0, 0.5, -1, 0.5 };
__constant__ float XRD3D_Mc[4] = { 0, 0.5, 0, -0.5 };
__constant__ float XRD3D_Md[4] = { 0, 1.0 / 6.0, 4.0 / 6.0, 1.0 / 6.0 };
__constant__ float XRD3D_dMa[4] = { 0.5, -1.5, 1.5, -0.5 };
__constant__ float XRD3D_dMb[4] = { 0, 1, -2, 1 };
__constant__ float XRD3D_dMc[4] = { 0, 0.5, 0, -0.5 };
__global__ void XRD3D_Atom_Near(const int atom_numbers, const UNSIGNED_INT_VECTOR *uint_crd,
	const INT_VECTOR grid_dimension,  const int layer_numbers,
	const VECTOR uint_crd_to_grid_serial_factor,
	const UNSIGNED_INT_VECTOR *XRD3D_kxyz, int **XRD3D_atom_near, UNSIGNED_INT_VECTOR *XRD3D_uxyz, VECTOR *XRD3D_frxyz)
{
	int atom = blockDim.x * blockIdx.x + threadIdx.x;
	if (atom < atom_numbers)
	{
		UNSIGNED_INT_VECTOR *temp_uxyz = &XRD3D_uxyz[atom];
		int k, tempux, tempuy, tempuz;
		float tempf;
		tempf = (float)uint_crd[atom].uint_x * uint_crd_to_grid_serial_factor.x;
		tempux = (int)tempf;
		XRD3D_frxyz[atom].x = tempf - tempux;

		tempf = (float)uint_crd[atom].uint_y * uint_crd_to_grid_serial_factor.y;
		tempuy = (int)tempf;
		XRD3D_frxyz[atom].y = tempf - tempuy;

		tempf = (float)uint_crd[atom].uint_z * uint_crd_to_grid_serial_factor.z;
		tempuz = (int)tempf;
		XRD3D_frxyz[atom].z = tempf - tempuz;

		if (tempux != (*temp_uxyz).uint_x || tempuy != (*temp_uxyz).uint_y || tempuz != (*temp_uxyz).uint_z)
		{
			(*temp_uxyz).uint_x = tempux;
			(*temp_uxyz).uint_y = tempuy;
			(*temp_uxyz).uint_z = tempuz;
			int *temp_near = XRD3D_atom_near[atom];
			int kx, ky, kz;
			for (k = 0; k < 64; k++)
			{
				UNSIGNED_INT_VECTOR temp_kxyz = XRD3D_kxyz[k];
				kx = tempux - temp_kxyz.uint_x;
				if (kx < 0)
					kx += grid_dimension.int_x;
				ky = tempuy - temp_kxyz.uint_y;
				if (ky < 0)
					ky += grid_dimension.int_y;
				kz = tempuz - temp_kxyz.uint_z;
				if (kz < 0)
					kz += grid_dimension.int_z;
				temp_near[k] = kz * layer_numbers + ky * grid_dimension.int_x + kx;
			}
		}
	}
}
__global__ void XRD3D_Q_Spread
(const int atom_numbers, int **XRD3D_atom_near, const float *charge,
const UNSIGNED_INT_VECTOR *XRD3D_kxyz,const VECTOR *XRD3D_frxyz,
float *XRD3D_Q )
{
	int atom = blockDim.x * blockIdx.x + threadIdx.x;

	if (atom < atom_numbers)
	{
		int k;
		float tempf, tempQ, tempf2;
		int *temp_near = XRD3D_atom_near[atom];
		VECTOR temp_frxyz = XRD3D_frxyz[atom];
		float tempcharge = charge[atom];

		UNSIGNED_INT_VECTOR temp_kxyz;
		unsigned int kx;

		for (k = threadIdx.y; k < 64; k = k + blockDim.y)
		{
			temp_kxyz = XRD3D_kxyz[k];
			kx = temp_kxyz.uint_x;
			tempf = (temp_frxyz.x);
			tempf2 = tempf * tempf;
			tempf = XRD3D_Ma[kx] * tempf * tempf2 + XRD3D_Mb[kx] * tempf2 + XRD3D_Mc[kx] * tempf + XRD3D_Md[kx];

			tempQ = tempcharge * tempf;

			kx = temp_kxyz.uint_y;
			tempf = (temp_frxyz.y);
			tempf2 = tempf * tempf;
			tempf = XRD3D_Ma[kx] * tempf * tempf2 + XRD3D_Mb[kx] * tempf2 + XRD3D_Mc[kx] * tempf + XRD3D_Md[kx];

			tempQ = tempQ * tempf;

			kx = temp_kxyz.uint_z;
			tempf = (temp_frxyz.z);
			tempf2 = tempf * tempf;
			tempf = XRD3D_Ma[kx] * tempf * tempf2 + XRD3D_Mb[kx] * tempf2 + XRD3D_Mc[kx] * tempf + XRD3D_Md[kx];
			tempQ = tempQ * tempf;

			atomicAdd(&XRD3D_Q[temp_near[k]], tempQ);
		}
	}
}
__global__ void XRD3D_Force(const int atom_numbers, int **XRD3D_atom_near, const float *charge,
	const VECTOR *XRD3D_frxyz, const UNSIGNED_INT_VECTOR *XRD3D_kxyz, const VECTOR crd_to_grid_serial_factor,
	const float *XRD3D_Q, VECTOR *force)
{
	int atom = blockDim.x * blockIdx.x + threadIdx.x;
	if (atom < atom_numbers)
	{

		int k, kx;
		float tempdQx, tempdQy, tempdQz, tempdx, tempdy, tempdz, tempx, tempy, tempz, tempdQf;
		float tempf, tempf2;
		float temp_charge = charge[atom];
		int *temp_near = XRD3D_atom_near[atom];
		UNSIGNED_INT_VECTOR temp_kxyz;
		VECTOR temp_frxyz = XRD3D_frxyz[atom];
		
		//在copy frc to grad 时会清空frc，因此这里不再重复此操作，且担心多线程冲突
		//force[atom].x = 0., force[atom].y = 0., force[atom].z = 0.;
		VECTOR temp_frc = {0.,0.,0.};
		for (k = threadIdx.y; k < 64; k = k + blockDim.y)
		{
			temp_kxyz = XRD3D_kxyz[k];
			tempdQf = -XRD3D_Q[temp_near[k]];

			kx = temp_kxyz.uint_x;
			tempf = (temp_frxyz.x);
			tempf2 = tempf * tempf;
			tempx = XRD3D_Ma[kx] * tempf * tempf2 + XRD3D_Mb[kx] * tempf2 + XRD3D_Mc[kx] * tempf + XRD3D_Md[kx];
			tempdx = XRD3D_dMa[kx] * tempf2 + XRD3D_dMb[kx] * tempf + XRD3D_dMc[kx];

			kx = temp_kxyz.uint_y;
			tempf = (temp_frxyz.y);
			tempf2 = tempf * tempf;
			tempy = XRD3D_Ma[kx] * tempf * tempf2 + XRD3D_Mb[kx] * tempf2 + XRD3D_Mc[kx] * tempf + XRD3D_Md[kx];
			tempdy = XRD3D_dMa[kx] * tempf2 + XRD3D_dMb[kx] * tempf + XRD3D_dMc[kx];

			kx = temp_kxyz.uint_z;
			tempf = (temp_frxyz.z);
			tempf2 = tempf * tempf;
			tempz = XRD3D_Ma[kx] * tempf * tempf2 + XRD3D_Mb[kx] * tempf2 + XRD3D_Mc[kx] * tempf + XRD3D_Md[kx];
			tempdz = XRD3D_dMa[kx] * tempf2 + XRD3D_dMb[kx] * tempf + XRD3D_dMc[kx];


			tempdQx = tempdx * tempy * tempz;
			tempdQy = tempdy * tempx * tempz;
			tempdQz = tempdz * tempx * tempy;

			temp_frc.x = temp_frc.x + tempdQf * tempdQx;
			temp_frc.y = temp_frc.y + tempdQf * tempdQy;
			temp_frc.z = temp_frc.z + tempdQf * tempdQz;

		}
		atomicAdd(&force[atom].x, temp_frc.x*temp_charge*crd_to_grid_serial_factor.x);
		atomicAdd(&force[atom].y, temp_frc.y*temp_charge*crd_to_grid_serial_factor.y);
		atomicAdd(&force[atom].z, temp_frc.z*temp_charge*crd_to_grid_serial_factor.z);
	}
}

__global__ void XRD3D_Energy_Product(const int element_number, const float* list1, const float* list2, float *sum)
{
	if (threadIdx.x == 0 && blockIdx.x==0)
	{
		sum[0] = 0.;
	}
	__syncthreads();
	int gpu_grid_layer_numbers = gridDim.x*blockDim.x;
	float temp = 0.;
	for (int i = blockIdx.x*blockDim.x + threadIdx.x; i < element_number; i = i + gpu_grid_layer_numbers)
	{
		temp = temp + list1[i] * list2[i];
	}
	atomicAdd(sum, temp);
}

__global__ void Copy_frc_To_System_Grad(const int atom_numbers, const int *atom_serial, VECTOR *system_grad, VECTOR *frc)
{
	int gpu_grid_layer_numbers = gridDim.x*blockDim.x;
	for (int i = blockIdx.x*blockDim.x + threadIdx.x; i < atom_numbers; i = i + gpu_grid_layer_numbers)
	{
		system_grad[atom_serial[i]] = -2.*frc[i];//能量梯度是能量导致力的负数，同时，由于求和对i>j和i<j均进行一次，因此梯度实际是力的两倍
		//printf("%d %f %f %f\n", atom_serial[i], system_grad[atom_serial[i]].x, system_grad[atom_serial[i]].y, system_grad[atom_serial[i]].z);
		frc[i].x = 0., frc[i].y = 0., frc[i].z = 0.;
	}
}
void XRD3D::FFT::FFT_One_Time()
{
	XRD3D_Atom_Near << <ceilf((float)xrd3d->xatom_numbers / 32.), 32 >> >(xrd3d->xatom_numbers, xrd3d->atom_uint_crd,
		xrd3d->grid_dimension, xrd3d->layer_numbers,
		uint_crd_to_grid_serial_factor,
		XRD3D_kxyz, XRD3D_atom_near, XRD3D_uxyz, XRD3D_frxyz);

	Reset_List << < ceilf((float)xrd3d->grid_numbers / 1024 ), 1024 >> >(xrd3d->grid_numbers, rho, 0);

	XRD3D_Q_Spread << < ceilf((float)xrd3d->xatom_numbers / thread_XRD3D.x), thread_XRD3D >> >
		(xrd3d->xatom_numbers, XRD3D_atom_near, xrd3d->d_fi,
		XRD3D_kxyz, XRD3D_frxyz,
		rho);

	cufftExecR2C(plan_3d_r2c, rho, fft_rho);

	Array_float_Multiply_cufftComplex << <20, 32 >> >
		(NzNy_Nx21_numbers, fft_fijk, fft_rho);

	cufftExecC2R(plan_3d_c2r, fft_rho, phi);

	XRD3D_Energy_Product << <20, 32 >> >
		(xrd3d->grid_numbers, rho, phi, xrd3d->d_sum_of_reciprocal_energy);

	XRD3D_Force << < ceilf((float)xrd3d->xatom_numbers / thread_XRD3D.x), thread_XRD3D >> >
		(xrd3d->xatom_numbers, XRD3D_atom_near, xrd3d->d_fi,
		XRD3D_frxyz, XRD3D_kxyz, crd_to_grid_serial_factor,
		phi, xrd3d->atom_frc);

	Copy_frc_To_System_Grad << <20, 32 >> >
		(xrd3d->xatom_numbers, xrd3d->d_atom_serial, xrd3d->system_atom_crd_grad, xrd3d->atom_frc);
}
void XRD3D::FFT::Refresh_FFT_Volume()
{
	box_length_inverse.x = 1. / xrd3d->box_length.x;
	box_length_inverse.y = 1. / xrd3d->box_length.y;
	box_length_inverse.z = 1. / xrd3d->box_length.z;

	crd_to_grid_serial_factor.x = 1. / xrd3d->box_length.x*xrd3d->grid_dimension.int_x;
	crd_to_grid_serial_factor.y = 1. / xrd3d->box_length.y*xrd3d->grid_dimension.int_y;
	crd_to_grid_serial_factor.z = 1. / xrd3d->box_length.z*xrd3d->grid_dimension.int_z;

	scalor = CONSTANT_Pi*sqrtf(CONSTANT_Pi) / xrd3d->volume / xrd3d->k0 / sqrtf(xrd3d->s) / xrd3d->xatom_numbers;
	Build_fft_fijk_List << <20, 32 >> >
		(NzNy_Nx21_numbers, xrd3d->grid_dimension,
		grid_dimension_upper_half, grid_dimension_lower_half_plus_one,
		box_length_inverse, xrd3d->s, xrd3d->k0, scalor,
		fft_B_spline, fft_fijk);
}
void XRD3D::FFT::Initial(CONTROLLER *controller, XRD3D *xrd3d)
{
	this->xrd3d = xrd3d;

	thread_XRD3D.x = 8;
	thread_XRD3D.y = 8;

	cuda_result=cufftPlan3d(&plan_3d_r2c, xrd3d->grid_dimension.int_z, xrd3d->grid_dimension.int_y, xrd3d->grid_dimension.int_x, CUFFT_R2C);
	cuda_result=cufftPlan3d(&plan_3d_c2r, xrd3d->grid_dimension.int_z, xrd3d->grid_dimension.int_y, xrd3d->grid_dimension.int_x, CUFFT_C2R);
	if (cuda_result != CUFFT_SUCCESS)
	{
		controller[0].printf("    FFT initial failed, cuda_result %d\n", cuda_result);
		getchar();
	}

	grid_dimension_lower_half_plus_one.int_x = (xrd3d->grid_dimension.int_x >> 1) + 1;
	grid_dimension_lower_half_plus_one.int_y = (xrd3d->grid_dimension.int_y >> 1) + 1;
	grid_dimension_lower_half_plus_one.int_z = (xrd3d->grid_dimension.int_z >> 1) + 1;

	grid_dimension_upper_half.int_x = (xrd3d->grid_dimension.int_x + 1) >> 1;
	grid_dimension_upper_half.int_y = (xrd3d->grid_dimension.int_y + 1) >> 1;
	grid_dimension_upper_half.int_z = (xrd3d->grid_dimension.int_z + 1) >> 1;

	box_length_inverse.x = 1. / xrd3d->box_length.x;
	box_length_inverse.y = 1. / xrd3d->box_length.y;
	box_length_inverse.z = 1. / xrd3d->box_length.z;

	crd_to_grid_serial_factor.x = 1. / xrd3d->box_length.x*xrd3d->grid_dimension.int_x;
	crd_to_grid_serial_factor.y = 1. / xrd3d->box_length.y*xrd3d->grid_dimension.int_y;
	crd_to_grid_serial_factor.z = 1. / xrd3d->box_length.z*xrd3d->grid_dimension.int_z;

	uint_crd_to_grid_serial_factor.x = 1. / CONSTANT_UINT_MAX_FLOAT*xrd3d->grid_dimension.int_x;
	uint_crd_to_grid_serial_factor.y = 1. / CONSTANT_UINT_MAX_FLOAT*xrd3d->grid_dimension.int_y;
	uint_crd_to_grid_serial_factor.z = 1. / CONSTANT_UINT_MAX_FLOAT*xrd3d->grid_dimension.int_z;

	NzNy_Nx21_numbers = xrd3d->grid_dimension.int_z*xrd3d->grid_dimension.int_y*grid_dimension_lower_half_plus_one.int_x;
	Cuda_Malloc_Safely((void**)&fft_B_spline, sizeof(cufftComplex)*NzNy_Nx21_numbers);
	Cuda_Malloc_Safely((void**)&fft_fijk, sizeof(float)*NzNy_Nx21_numbers);
	Cuda_Malloc_Safely((void**)&d_B_spline_grid, sizeof(float)*xrd3d->grid_numbers);
	Malloc_Safely((void**)&h_B_spline_grid, sizeof(float)*xrd3d->grid_numbers);

	//构建B样条线插值的卷积格点
	for (int i = 0; i < xrd3d->grid_numbers; i = i + 1)
	{
		h_B_spline_grid[i] = 0.;
	}
	float temp_b_spline[3] = { 1. / 6., 2. / 3., 1. / 6. };
	for (int k = -1; k <= 1; k = k + 1)
	{
		for (int j = -1; j <= 1; j = j + 1)
		{
			for (int i = -1; i <= 1; i = i + 1)
			{
				float weight = temp_b_spline[k + 1] * temp_b_spline[j + 1] * temp_b_spline[i + 1];
				int kk, jj, ii;
				if (k < 0)
				{
					kk = k + xrd3d->grid_dimension.int_z;
				}
				else
				{
					kk = k;
				}
				if (j < 0)
				{
					jj = j + xrd3d->grid_dimension.int_y;
				}
				else
				{
					jj = j;
				}
				if (i < 0)
				{
					ii = i + xrd3d->grid_dimension.int_x;
				}
				else
				{
					ii = i;
				}
				int grid_serial = kk*xrd3d->layer_numbers + jj*xrd3d->grid_dimension.int_x + ii;
				h_B_spline_grid[grid_serial] = weight;
			}
		}
	}
	cudaMemcpy(d_B_spline_grid, h_B_spline_grid, sizeof(float)*xrd3d->grid_numbers, cudaMemcpyHostToDevice);

	//进行fft获得fft B spline
	cufftExecR2C(plan_3d_r2c, d_B_spline_grid, fft_B_spline);
	scalor = CONSTANT_Pi*sqrtf(CONSTANT_Pi) / xrd3d->volume / xrd3d->k0 / sqrtf(xrd3d->s)/xrd3d->xatom_numbers;
	Build_fft_fijk_List << <20, 32 >> >
		(NzNy_Nx21_numbers, xrd3d->grid_dimension,
		grid_dimension_upper_half, grid_dimension_lower_half_plus_one,
		box_length_inverse, xrd3d->s, xrd3d->k0, scalor,
		fft_B_spline, fft_fijk);

	Cuda_Malloc_Safely((void**)&fft_rho, sizeof(cufftComplex)*NzNy_Nx21_numbers);
	Cuda_Malloc_Safely((void**)&rho, sizeof(float)*xrd3d->grid_numbers);
	Cuda_Malloc_Safely((void**)&phi, sizeof(float)*xrd3d->grid_numbers);

	//初始化样条线插值变量
	Cuda_Malloc_Safely((void**)&XRD3D_uxyz, sizeof(UNSIGNED_INT_VECTOR)* xrd3d->xatom_numbers);
	Cuda_Malloc_Safely((void**)&XRD3D_frxyz, sizeof(VECTOR)* xrd3d->xatom_numbers);
	Reset_List << <3. * xrd3d->xatom_numbers / 32 + 1, 32 >> >(3 * xrd3d->xatom_numbers, (int*)XRD3D_uxyz, 1 << 30);

	int **atom_near_cpu;
	Malloc_Safely((void**)&atom_near_cpu, sizeof(int*)* xrd3d->xatom_numbers);
	Cuda_Malloc_Safely((void**)&XRD3D_atom_near, sizeof(int*)* xrd3d->xatom_numbers);
	for (int i = 0; i < xrd3d->xatom_numbers; i++)
	{
		Cuda_Malloc_Safely((void**)&atom_near_cpu[i], sizeof(int)* 64);
	}
	cudaMemcpy(XRD3D_atom_near, atom_near_cpu, sizeof(int*)* xrd3d->xatom_numbers, cudaMemcpyHostToDevice);
	free(atom_near_cpu);

	UNSIGNED_INT_VECTOR *XRD3D_kxyz_cpu;
	Cuda_Malloc_Safely((void**)&XRD3D_kxyz, sizeof(UNSIGNED_INT_VECTOR)* 64);
	Malloc_Safely((void**)&XRD3D_kxyz_cpu, sizeof(UNSIGNED_INT_VECTOR)* 64);
	for (int kz = 0; kz < 4; kz++)
	{

		for (int ky = 0; ky < 4; ky++)
		{

			for (int kx = 0; kx < 4; kx++)
			{
				int index = kz * 16 + ky * 4 + kx;
				XRD3D_kxyz_cpu[index].uint_x = kx;
				XRD3D_kxyz_cpu[index].uint_y = ky;
				XRD3D_kxyz_cpu[index].uint_z = kz;
			}
		}
	}
	cudaMemcpy(XRD3D_kxyz, XRD3D_kxyz_cpu, sizeof(UNSIGNED_INT_VECTOR)* 64, cudaMemcpyHostToDevice);
	free(XRD3D_kxyz_cpu);
}