/*
pbcX.cu

copy vector (X-PBC)
*/

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

// copy
__host__ __device__ __forceinline__
static void copyXpbcEx_(cuDoubleComplex *v, int j, int k, param_t *q)
{
	v[nEX(q,       - 1, j, k)] = v[nEX(q, q->Nx - 1, j, k)];
	v[nEX(q, q->Nx + 0, j, k)] = v[nEX(q,         0, j, k)];
}
__host__ __device__ __forceinline__
static void copyXpbcEy_(cuDoubleComplex *v, int j, int k, param_t *q)
{
	v[nEY(q,       - 1, j, k)] = v[nEY(q, q->Nx - 1, j, k)];
	v[nEY(q, q->Nx + 1, j, k)] = v[nEY(q,         1, j, k)];
}
__host__ __device__ __forceinline__
static void copyXpbcEz_(cuDoubleComplex *v, int j, int k, param_t *q)
{
	v[nEZ(q,       - 1, j, k)] = v[nEZ(q, q->Nx - 1, j, k)];
	v[nEZ(q, q->Nx + 1, j, k)] = v[nEZ(q,         1, j, k)];
}

// factor
__host__ __device__ __forceinline__
static void factorXpbcEx_(cuDoubleComplex *v, int i0, int i1, int j, int k, param_t *q)
{
	int64_t n;
	int i;
	const cuDoubleComplex fp = q->pbc_fx;
	const cuDoubleComplex fm = cuCinv(fp);
	i =       - 1;
	n = nEX(q, i, j, k);
	if ((i >= i0 - 1) && (i <= i1 + 1)) {  // MPI
		v[n] = cuCmul(v[n], fm);
	}
	i = q->Nx + 0;
	n = nEX(q, i, j, k);
	if ((i >= i0 - 1) && (i <= i1 + 1)) {  // MPI
		v[n] = cuCmul(v[n], fp);
	}
}
__host__ __device__ __forceinline__
static void factorXpbcEy_(cuDoubleComplex *v, int i0, int i1, int j, int k, param_t *q)
{
	int64_t n;
	int i;
	const cuDoubleComplex fp = q->pbc_fx;
	const cuDoubleComplex fm = cuCinv(fp);
	i =       - 1;
	n = nEY(q, i, j, k);
	if ((i >= i0 - 1) && (i <= i1 + 1)) {  // MPI
		v[n] = cuCmul(v[n], fm);
	}
	i = q->Nx + 1;
	n = nEY(q, i, j, k);
	if ((i >= i0 - 1) && (i <= i1 + 1)) {  // MPI
		v[n] = cuCmul(v[n], fp);
	}
}
__host__ __device__ __forceinline__
static void factorXpbcEz_(cuDoubleComplex *v, int i0, int i1, int j, int k, param_t *q)
{
	int64_t n;
	int i;
	const cuDoubleComplex fp = q->pbc_fx;
	const cuDoubleComplex fm = cuCinv(fp);
	i =       - 1;
	n = nEZ(q, i, j, k);
	if ((i >= i0 - 1) && (i <= i1 + 1)) {  // MPI
		v[n] = cuCmul(v[n], fm);
	}
	i = q->Nx + 1;
	n = nEZ(q, i, j, k);
	if ((i >= i0 - 1) && (i <= i1 + 1)) {  // MPI
		v[n] = cuCmul(v[n], fp);
	}
}

// (GPU)
__global__
static void copyXpbc_gpu(char c, cuDoubleComplex *v, int j0, int j1, int k0, int k1)
{
	int j = (j0 - 1) + threadIdx.x + (blockIdx.x * blockDim.x);
	int k = (k0 - 1) + threadIdx.y + (blockIdx.y * blockDim.y);
	if ((j <= j1 + 1) &&
	    (k <= k1 + 1)) {
		if (c == 'X') copyXpbcEx_(v, j, k, &d_Param);
		if (c == 'Y') copyXpbcEy_(v, j, k, &d_Param);
		if (c == 'Z') copyXpbcEz_(v, j, k, &d_Param);
	}
}
__global__
static void factorXpbc_gpu(char c, cuDoubleComplex *v, int i0, int i1, int j0, int j1, int k0, int k1)
{
	int j = (j0 - 1) + threadIdx.x + (blockIdx.x * blockDim.x);
	int k = (k0 - 1) + threadIdx.y + (blockIdx.y * blockDim.y);
	if ((j <= j1 + 1) &&
	    (k <= k1 + 1)) {
		if (c == 'X') factorXpbcEx_(v, i0, i1, j, k, &d_Param);
		if (c == 'Y') factorXpbcEy_(v, i0, i1, j, k, &d_Param);
		if (c == 'Z') factorXpbcEz_(v, i0, i1, j, k, &d_Param);
	}
}

// (CPU)
static void copyXpbc_cpu(char c, cuDoubleComplex *v, int j0, int j1, int k0, int k1)
{
	for (int j = j0 - 1; j <= j1 + 1; j++) {
	for (int k = k0 - 1; k <= k1 + 1; k++) {
		if (c == 'X') copyXpbcEx_(v, j, k, &h_Param);
		if (c == 'Y') copyXpbcEy_(v, j, k, &h_Param);
		if (c == 'Z') copyXpbcEz_(v, j, k, &h_Param);
	}
	}
}
static void factorXpbc_cpu(char c, cuDoubleComplex *v, int i0, int i1, int j0, int j1, int k0, int k1)
{
	for (int j = j0 - 1; j <= j1 + 1; j++) {
	for (int k = k0 - 1; k <= k1 + 1; k++) {
		if (c == 'X') factorXpbcEx_(v, i0, i1, j, k, &h_Param);
		if (c == 'Y') factorXpbcEy_(v, i0, i1, j, k, &h_Param);
		if (c == 'Z') factorXpbcEz_(v, i0, i1, j, k, &h_Param);
	}
	}
}

void copyXpbc(cuDoubleComplex *v)
{
	assert(Pbc.xpbc);

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

void factorXpbc(cuDoubleComplex *v)
{
	assert(Pbc.xpbc);

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