/*
matrixEy.cu
u = Av (Ey)
*/

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

__host__ __device__ __forceinline__
static void matrixEy_(
	int i, int j, int k,
	int pre, int ifreq,
	cuDoubleComplex *a, const cuDoubleComplex *v, cuDoubleComplex *u, cuDoubleComplex *b,
	const double *kdzn, const double *kdxn,
	const double *kdyc,
	const double *krzn, const double *krxn,
	const double *krzp, const double *krxp,
	const double *krzm, const double *krxm,
	const double *kryc,
	const material_t *material, const int *me, const int *mh,
	const double *fxpml, const double *fypml, const double *fzpml,
	param_t *q)
{
	int nx = q->Nx;
	int ny = q->Ny;
	int nz = q->Nz;
	int zpbc = q->Zpbc;
	int xpbc = q->Xpbc;
	int lpml = q->Lpml;
	int m0 = ifreq * q->NMaterial;
	int64_t n0 = nEY(q, i, j, k);

	int64_t p[NNZ];
	p[ 0] = nEY(q, i,     j,     k    );
	p[ 1] = nEY(q, i,     j,     k + 1);
	p[ 2] = nEY(q, i,     j,     k - 1);
	p[ 3] = nEY(q, i + 1, j,     k    );
	p[ 4] = nEY(q, i - 1, j,     k    );
	p[ 5] = nEZ(q, i,     j + 1, k    );
	p[ 6] = nEZ(q, i,     j,     k    );
	p[ 7] = nEZ(q, i,     j + 1, k - 1);
	p[ 8] = nEZ(q, i,     j,     k - 1);
	p[ 9] = nEX(q, i,     j + 1, k    );
	p[10] = nEX(q, i,     j,     k    );
	p[11] = nEX(q, i - 1, j + 1, k    );
	p[12] = nEX(q, i - 1, j,     k    );

	cuDoubleComplex c[NNZ];
	memset(c, 0, NNZ * sizeof(cuDoubleComplex));

	if (mEY(q, i, j, k) == PEC) {
		c[0] = make_cuDoubleComplex(1, 0);
	}
	else {
		if ((q->Abc == MUR) && (
			(!zpbc && ((k == 0) || (k == nz))) ||
			(!xpbc && ((i == 0) || (i == nx))))) {
			// Mur
			cuDoubleComplex cv = cuCcast(material[mEY(q, i, j, k) + m0].cv);
			cuDoubleComplex cel = make_cuDoubleComplex(0, 0);
			int mur = 0;
			if      (!zpbc && (k ==  0)) {
				cel = cuCrmul(kdzn[ 0], cv);
				mur = 1;
			}
			else if (!zpbc && (k == nz)) {
				cel = cuCrmul(kdzn[nz], cv);
				mur = 2;
			}
			else if (!xpbc && (i ==  0)) {
				cel = cuCrmul(kdxn[ 0], cv);
				mur = 3;
			}
			else if (!xpbc && (i == nx)) {
				cel = cuCrmul(kdxn[nx], cv);
				mur = 4;
			}
			c[0] = make_cuDoubleComplex(1, 0);
			c[mur] = cuCrmul(-exp(cel.y), cuCexp(-cel.x));
		}
		else if ((q->Abc == PML) &&
			((j < 0) || (j >= ny) || (k <= 0) || (k >= nz) || (i <= 0) || (i >= nx))) {
			// PML
			int ib = MINMAX(0, nx,     i);
			int jb = MINMAX(0, ny - 1, j);
			int kb = MINMAX(0, nz,     k);
			if (mEY(q, ib, jb, kb) != PEC) {
				d_complex_t cv = material[mEY(q, ib, jb, kb) + m0].cv;
				cuDoubleComplex vc = cuCinv(make_cuDoubleComplex(cv.r, cv.i));

				cuDoubleComplex zc = cuCdiv(vc, make_cuDoubleComplex(fzpml[(2 * k + 0) + (2 * lpml)], kdzn[kb]));
				cuDoubleComplex zp = cuCdiv(vc, make_cuDoubleComplex(fzpml[(2 * k + 1) + (2 * lpml)], kdzn[kb]));
				cuDoubleComplex zm = cuCdiv(vc, make_cuDoubleComplex(fzpml[(2 * k - 1) + (2 * lpml)], kdzn[kb]));
				cuDoubleComplex xc = cuCdiv(vc, make_cuDoubleComplex(fxpml[(2 * i + 0) + (2 * lpml)], kdxn[ib]));
				cuDoubleComplex xp = cuCdiv(vc, make_cuDoubleComplex(fxpml[(2 * i + 1) + (2 * lpml)], kdxn[ib]));
				cuDoubleComplex xm = cuCdiv(vc, make_cuDoubleComplex(fxpml[(2 * i - 1) + (2 * lpml)], kdxn[ib]));
				cuDoubleComplex yp = cuCdiv(vc, make_cuDoubleComplex(fypml[(2 * j + 1) + (2 * lpml)], kdyc[jb]));

				c[ 0] = make_cuDoubleComplex(1, 0);
				c[ 1] = cuCneg(cuCmul(zc, zp));
				c[ 2] = cuCneg(cuCmul(zc, zm));
				c[ 3] = cuCneg(cuCmul(xc, xp));
				c[ 4] = cuCneg(cuCmul(xc, xm));
				c[ 5] =        cuCmul(zc, yp);
				c[ 6] = cuCneg(c[ 5]);
				c[ 7] =        c[ 6];
				c[ 8] =        c[ 5];
				c[ 9] =        cuCmul(xc, yp);
				c[10] = cuCneg(c[ 9]);
				c[11] =        c[10];
				c[12] =        c[ 9];
				c[ 0] = cuCsub(c[ 0], cuCadd4(c[ 1], c[ 2], c[ 3], c[ 4]));
			}
		}
		else {
			cuDoubleComplex c0 = cuCcast(material[mEY(q, i,     j,     k    ) + m0].cepsr);
			cuDoubleComplex c1 = cuCcast(material[mHX(q, i,     j,     k    ) + m0].cnur);
			cuDoubleComplex c2 = cuCcast(material[mHX(q, i,     j,     k - 1) + m0].cnur);
			cuDoubleComplex c3 = cuCcast(material[mHZ(q, i,     j,     k    ) + m0].cnur);
			cuDoubleComplex c4 = cuCcast(material[mHZ(q, i - 1, j,     k    ) + m0].cnur);

			c[ 0] = cuCneg(c0);
			c[ 1] = cuCrmul(- krzp[k],           c1);
			c[ 2] = cuCrmul(- krzm[k],           c2);
			c[ 3] = cuCrmul(- krxp[i],           c3);
			c[ 4] = cuCrmul(- krxm[i],           c4);
			c[ 5] = cuCrmul(+ kryc[j] * krzn[k], c1);
			c[ 6] = cuCneg(c[ 5]);
			c[ 7] = cuCrmul(- kryc[j] * krzn[k], c2);
			c[ 8] = cuCneg(c[ 7]);
			c[ 9] = cuCrmul(+ kryc[j] * krxn[i], c3);
			c[10] = cuCneg(c[ 9]);
			c[11] = cuCrmul(- kryc[j] * krxn[i], c4);
			c[12] = cuCneg(c[11]);
			c[ 0] = cuCsub(c[ 0], cuCadd4(c[ 1], c[ 2], c[ 3], c[ 4]));
		}
	}

	cuDoubleComplex sf = q->SCALING ? cuCinv(c[0]) : make_cuDoubleComplex(1, 0);
	if (pre) {
		// setup matrix
		if (q->MATRIX) {
			for (int m = 0; m < NNZ; m++) {
				a[(n0 * NNZ) + m] = cuCmul(sf, c[m]);
			}
		}
		// scaling RHS
		b[n0] = cuCmul(sf, b[n0]);
	}
	else {
		// matrix * vector
		cuDoubleComplex sum = make_cuDoubleComplex(0, 0);
		for (int m = 0; m < NNZ; m++) {
			sum = cuCadd(sum, cuCmul(c[m], v[p[m]]));
		}
		u[n0] = cuCmul(sf, sum);
	}
}

__global__
static void matrixEy_gpu(
	int i0, int i1, int j0, int j1, int k0, int k1,
	int pre, int ifreq,
	cuDoubleComplex *a, const cuDoubleComplex *v, cuDoubleComplex *u, cuDoubleComplex *b,
	const double *kdzn, const double *kdxn,
	const double *kdyc,
	const double *krzn, const double *krxn,
	const double *krzp, const double *krxp,
	const double *krzm, const double *krxm,
	const double *kryc,
	const material_t *material, const int *me, const int *mh,
	const double *fxpml, const double *fypml, const double *fzpml)
{
	int i = i0 + threadIdx.z + (blockIdx.z * blockDim.z);
	int j = j0 + threadIdx.y + (blockIdx.y * blockDim.y);
	int k = k0 + threadIdx.x + (blockIdx.x * blockDim.x);
	if ((i <= i1) &&
	    (j <= j1) &&
	    (k <= k1)) {
		matrixEy_(
			i, j, k,
			pre, ifreq,
			a, v, u, b,
			kdzn, kdxn,
			kdyc,
			krzn, krxn,
			krzp, krxp,
			krzm, krxm,
			kryc,
			material, me, mh,
			fxpml, fypml, fzpml,
			&d_Param);
	}
}

static void matrixEy_cpu(
	int i0, int i1, int j0, int j1, int k0, int k1,
	int pre, int ifreq,
	cuDoubleComplex *a, const cuDoubleComplex *v, cuDoubleComplex *u, cuDoubleComplex *b,
	const double *kdzn, const double *kdxn,
	const double *kdyc,
	const double *krzn, const double *krxn,
	const double *krzp, const double *krxp,
	const double *krzm, const double *krxm,
	const double *kryc,
	const material_t *material, const int *me, const int *mh,
	const double *fxpml, const double *fypml, const double *fzpml)
{
	for (int i = i0; i <= i1; i++) {
	for (int j = j0; j <= j1; j++) {
	for (int k = k0; k <= k1; k++) {
		matrixEy_(
			i, j, k,
			pre, ifreq,
			a, v, u, b,
			kdzn, kdxn,
			kdyc,
			krzn, krxn,
			krzp, krxp,
			krzm, krxm,
			kryc,
			material, me, mh,
			fxpml, fypml, fzpml,
			&h_Param);
	}
	}
	}
}

void matrixEy(int pre, int ifreq, cuDoubleComplex *a, const cuDoubleComplex *v, cuDoubleComplex *u, cuDoubleComplex *b)
{
	if (GPU) {
		cudaMemcpyToSymbol(d_Param, &h_Param, sizeof(param_t));
		dim3 block(BLOCK_X, BLOCK_Y, 1);
		dim3 grid(
			CEIL(k1Ey - k0Ey + 1, block.x),
			CEIL(j1Ey - j0Ey + 1, block.y),
			CEIL(i1Ey - i0Ey + 1, block.z));
		matrixEy_gpu<<<grid, block>>>(
			i0Ey, i1Ey, j0Ey, j1Ey, k0Ey, k1Ey,
			pre, ifreq,
			a, v, u, b,
			d_KDZn, d_KDXn,
			d_KDYc,
			d_KRZn, d_KRXn,
			d_KRZp, d_KRXp,
			d_KRZm, d_KRXm,
			d_KRYc,
			d_Material, d_mE, d_mH,
			d_Fxpml, d_Fypml, d_Fzpml);
		if (UM) cudaDeviceSynchronize();
	}
	else {
		matrixEy_cpu(
			i0Ey, i1Ey, j0Ey, j1Ey, k0Ey, k1Ey,
			pre, ifreq,
			a, v, u, b,
			KDZn, KDXn,
			KDYc,
			KRZn, KRXn,
			KRZp, KRXp,
			KRZm, KRXm,
			KRYc,
			Material, mE, mH,
			Fxpml, Fypml, Fzpml);
	}
}
