/*
pbcZ.cu

copy vector (Z-PBC)
*/

#include "oth.h"
#include "oth_cuda.h"
#include "cuda_complex.h"

// copy
__host__ __device__ __forceinline__
static void copyZpbcEz_(cuDoubleComplex *v, int i, int j, param_t *q)
{
	v[nEZ(q, i, j,       - 1)] = v[nEZ(q, i, j, q->Nz - 1)];
	v[nEZ(q, i, j, q->Nz + 0)] = v[nEZ(q, i, j,         0)];
}
__host__ __device__ __forceinline__
static void copyZpbcEx_(cuDoubleComplex *v, int i, int j, param_t *q)
{
	v[nEX(q, i, j,       - 1)] = v[nEX(q, i, j, q->Nz - 1)];
	v[nEX(q, i, j, q->Nz + 1)] = v[nEX(q, i, j,         1)];
}
__host__ __device__ __forceinline__
static void copyZpbcEy_(cuDoubleComplex *v, int i, int j, param_t *q)
{
	v[nEY(q, i, j,       - 1)] = v[nEY(q, i, j, q->Nz - 1)];
	v[nEY(q, i, j, q->Nz + 1)] = v[nEY(q, i, j,         1)];
}

// factor
__host__ __device__ __forceinline__
static void factorZpbcEz_(cuDoubleComplex *v, int k0, int k1, int i, int j, param_t *q)
{
	int64_t n;
	int k;
	const cuDoubleComplex fp = q->pbc_fz;
	const cuDoubleComplex fm = cuCinv(fp);
	k =       - 1;
	n = nEZ(q, i, j, k);
	if ((k >= k0 - 1) && (k <= k1 + 1)) {  // MPI
		v[n] = cuCmul(v[n], fm);
	}
	k = q->Nz + 0;
	n = nEZ(q, i, j, k);
	if ((k >= k0 - 1) && (k <= k1 + 1)) {  // MPI
		v[n] = cuCmul(v[n], fp);
	}
}
__host__ __device__ __forceinline__
static void factorZpbcEx_(cuDoubleComplex *v, int k0, int k1, int i, int j, param_t *q)
{
	int64_t n;
	int k;
	const cuDoubleComplex fp = q->pbc_fz;
	const cuDoubleComplex fm = cuCinv(fp);
	k =       - 1;
	n = nEX(q, i, j, k);
	if ((k >= k0 - 1) && (k <= k1 + 1)) {  // MPI
		v[n] = cuCmul(v[n], fm);
	}
	k = q->Nz + 1;
	n = nEX(q, i, j, k);
	if ((k >= k0 - 1) && (k <= k1 + 1)) {  // MPI
		v[n] = cuCmul(v[n], fp);
	}
}
__host__ __device__ __forceinline__
static void factorZpbcEy_(cuDoubleComplex *v, int k0, int k1, int i, int j, param_t *q)
{
	int64_t n;
	int k;
	const cuDoubleComplex fp = q->pbc_fz;
	const cuDoubleComplex fm = cuCinv(fp);
	k =       - 1;
	n = nEY(q, i, j, k);
	if ((k >= k0 - 1) && (k <= k1 + 1)) {  // MPI
		v[n] = cuCmul(v[n], fm);
	}
	k = q->Nz + 1;
	n = nEY(q, i, j, k);
	if ((k >= k0 - 1) && (k <= k1 + 1)) {  // MPI
		v[n] = cuCmul(v[n], fp);
	}
}

// (GPU)
__global__
static void copyZpbc_gpu(char c, cuDoubleComplex *v, int i0, int i1, int j0, int j1)
{
	int i = (i0 - 1) + threadIdx.x + (blockIdx.x * blockDim.x);
	int j = (j0 - 1) + threadIdx.y + (blockIdx.y * blockDim.y);
	if ((i <= i1 + 1) &&
	    (j <= j1 + 1)) {
		if (c == 'X') copyZpbcEx_(v, i, j, &d_Param);
		if (c == 'Y') copyZpbcEy_(v, i, j, &d_Param);
		if (c == 'Z') copyZpbcEz_(v, i, j, &d_Param);
	}
}
__global__
static void factorZpbc_gpu(char c, cuDoubleComplex *v, int i0, int i1, int j0, int j1, int k0, int k1)
{
	int i = (i0 - 1) + threadIdx.x + (blockIdx.x * blockDim.x);
	int j = (j0 - 1) + threadIdx.y + (blockIdx.y * blockDim.y);
	if ((i <= i1 + 1) &&
	    (j <= j1 + 1)) {
		if (c == 'X') factorZpbcEx_(v, k0, k1, i, j, &d_Param);
		if (c == 'Y') factorZpbcEy_(v, k0, k1, i, j, &d_Param);
		if (c == 'Z') factorZpbcEz_(v, k0, k1, i, j, &d_Param);
	}
}

// (CPU)
static void copyZpbc_cpu(char c, cuDoubleComplex *v, int i0, int i1, int j0, int j1)
{
	for (int i = i0 - 1; i <= i1 + 1; i++) {
	for (int j = j0 - 1; j <= j1 + 1; j++) {
		if (c == 'X') copyZpbcEx_(v, i, j, &h_Param);
		if (c == 'Y') copyZpbcEy_(v, i, j, &h_Param);
		if (c == 'Z') copyZpbcEz_(v, i, j, &h_Param);
	}
	}
}
static void factorZpbc_cpu(char c, cuDoubleComplex *v, int i0, int i1, int j0, int j1, int k0, int k1)
{
	for (int i = i0 - 1; i <= i1 + 1; i++) {
	for (int j = j0 - 1; j <= j1 + 1; j++) {
		if (c == 'X') factorZpbcEx_(v, k0, k1, i, j, &h_Param);
		if (c == 'Y') factorZpbcEy_(v, k0, k1, i, j, &h_Param);
		if (c == 'Z') factorZpbcEz_(v, k0, k1, i, j, &h_Param);
	}
	}
}

void copyZpbc(cuDoubleComplex *v)
{
	assert(Pbc.zpbc);

	if (GPU) {
		dim3 block(PBC_BLOCK_X, PBC_BLOCK_Y);
		dim3 gridx(CEIL(i1Ex - i0Ex + 3, block.x),
		           CEIL(j1Ex - j0Ex + 3, block.y));
		dim3 gridy(CEIL(i1Ey - i0Ey + 3, block.x),
		           CEIL(j1Ey - j0Ey + 3, block.y));
		dim3 gridz(CEIL(i1Ez - i0Ez + 3, block.x),
		           CEIL(j1Ez - j0Ez + 3, block.y));
		cudaMemcpyToSymbol(d_Param, &h_Param, sizeof(param_t));
		copyZpbc_gpu<<<gridx, block>>>('X', v, i0Ex, i1Ex, j0Ex, j1Ex);
		copyZpbc_gpu<<<gridy, block>>>('Y', v, i0Ey, i1Ey, j0Ey, j1Ey);
		copyZpbc_gpu<<<gridz, block>>>('Z', v, i0Ez, i1Ez, j0Ez, j1Ez);
		if (UM) cudaDeviceSynchronize();
	}
	else {
		copyZpbc_cpu('X', v, i0Ex, i1Ex, j0Ex, j1Ex);
		copyZpbc_cpu('Y', v, i0Ey, i1Ey, j0Ey, j1Ey);
		copyZpbc_cpu('Z', v, i0Ez, i1Ez, j0Ez, j1Ez);
	}
}

void factorZpbc(cuDoubleComplex *v)
{
	assert(Pbc.zpbc);

	if (GPU) {
		dim3 block(PBC_BLOCK_X, PBC_BLOCK_Y);
		dim3 gridx(CEIL(i1Ex - i0Ex + 3, block.x),
		           CEIL(j1Ex - j0Ex + 3, block.y));
		dim3 gridy(CEIL(i1Ey - i0Ey + 3, block.x),
		           CEIL(j1Ey - j0Ey + 3, block.y));
		dim3 gridz(CEIL(i1Ez - i0Ez + 3, block.x),
		           CEIL(j1Ez - j0Ez + 3, block.y));
		cudaMemcpyToSymbol(d_Param, &h_Param, sizeof(param_t));
		factorZpbc_gpu<<<gridx, block>>>('X', v, i0Ex, i1Ex, j0Ex, j1Ex, k0Ex, k1Ex);
		factorZpbc_gpu<<<gridy, block>>>('Y', v, i0Ey, i1Ey, j0Ey, j1Ey, k0Ey, k1Ey);
		factorZpbc_gpu<<<gridz, block>>>('Z', v, i0Ez, i1Ez, j0Ez, j1Ez, k0Ez, k1Ez);
		if (UM) cudaDeviceSynchronize();
	}
	else {
		factorZpbc_cpu('X', v, i0Ex, i1Ex, j0Ex, j1Ex, k0Ex, k1Ex);
		factorZpbc_cpu('Y', v, i0Ey, i1Ey, j0Ey, j1Ey, k0Ey, k1Ey);
		factorZpbc_cpu('Z', v, i0Ez, i1Ez, j0Ez, j1Ez, k0Ez, k1Ez);
	}
}
