/*
setupIndex_mpi.c (MPI)
*/

#include "oth.h"
#include "oth_mpi.h"
#include "oth_prototype.h"

void setupIndex_mpi(int commsize, int commrank)
{
/*
	E : array index
	Ex : i=-L...Nx+L-1, j=-L...Ny+L,   k=-L...Nz+L
	Ey : i=-L...Nx+L,   j=-L...Ny+L-1, k=-L...Nz+L
	Ez : i=-L...Nx+L,   j=-L...Ny+L,   k=-L...Nz+L-1
*/
	const int lx = (Abc == PML) ? Lpml : 1;
	const int i0 =    - lx;
	const int i1 = Nx + lx;

	int *lsum = (int *)calloc(commsize, sizeof(int));
	int lrank = 0;
	for (int i = i0; i <= i1; i++) {
		lsum[lrank]++;
		lrank++;
		if (lrank == commsize) lrank = 0;
	}

	int *imin = (int *)malloc(commsize * sizeof(int));
	imin[0] = i0;
	for (int rank = 1; rank < commsize; rank++) {
		imin[rank] = imin[rank - 1] + lsum[rank];
	}
	free(lsum);

	int *imax = (int *)malloc(commsize * sizeof(int));
	for (int rank = 0; rank < commsize - 1; rank++) {
		imax[rank] = imin[rank + 1];
	}
	imax[commsize - 1] = i1;

	for (int rank = 0; rank < commsize; rank++) {
		if (rank > 0) imin[rank]--;
		if (rank < commsize - 1) imax[rank]++;
	}

	const int i0ex = imin[commrank];
	const int i0ey = imin[commrank];
	const int i0ez = imin[commrank];
	const int i1ex = imax[commrank] - 1;
	const int i1ey = imax[commrank];
	const int i1ez = imax[commrank];
	//printf("%d %d %d %d (%d) %d %d (%d) %d %d (%d)\n", commsize, commrank, i0ex, i1ex, (i1ex - i0ex + 1), i0ey, i1ey, (i1ey - i0ey + 1), i0ez, i1ez, (i1ez - i0ez + 1)); fflush(stdout); exit(0);
	//printf("%3d %3d %3d %3d\n", commsize, commrank, (i1ex - i0ex + 1), (i1ey - i0ey + 1)); fflush(stdout); exit(0);
	free(imin);
	free(imax);

	const int ly = (Abc == PML) ? Lpml : 1;
	const int j0ex =    - ly;
	const int j0ey =    - ly;
	const int j0ez =    - ly;
	const int j1ex = Ny + ly;
	const int j1ey = Ny + ly - 1;
	const int j1ez = Ny + ly;

	const int lz = (Abc == PML) ? Lpml : 1;
	const int k0ex =    - lz;
	const int k0ey =    - lz;
	const int k0ez =    - lz;
	const int k1ex = Nz + lz;
	const int k1ey = Nz + lz;
	const int k1ez = Nz + lz - 1;

	NE = arrayIndex(
		i0ex, i1ex, j0ex, j1ex, k0ex, k1ex,
		i0ey, i1ey, j0ey, j1ey, k0ey, k1ey,
		i0ez, i1ez, j0ez, j1ez, k0ez, k1ez);
	//printf("%d %d %zd\n", commsize, commrank, NE); fflush(stdout);

	assert(NEY(i0ey, j0ey, k0ey)     == 0);
	assert(NEZ(i0ez, j0ez, k0ez)     == NEY(i0ey, j1ey, k1ey) + 1);
	assert(NEX(i0ex, j0ex, k0ex)     == NEZ(i0ez, j1ez, k1ez) + 1);
	assert(NEZ(i1ez, j1ez, k1ez) + 1 == NE);
	assert(NEY(i1ey, j1ey, k1ey) + 1 == NEZ(i1ez, j0ez, k0ez));
	assert(NEX(i1ex, j1ex, k1ex) + 1 == NEY(i1ey, j0ey, k0ey));

	// E : loop index

	i0Ex = i0ex + ((commrank == 0)            && (Abc == PML) ? 0 : 1);
	i0Ey = i0ey + 1;
	i0Ez = i0ez + 1;
	i1Ex = i1ex - ((commrank == commsize - 1) && (Abc == PML) ? 0 : 1);
	i1Ey = i1ey - 1;
	i1Ez = i1ez - 1;
	if (commrank < commsize - 1) {
		i1Ey--;
		i1Ez--;
	}

	j0Ex = j0ex + 1;
	j0Ey = j0ey + ((Abc == PML) ? 0 : 1);
	j0Ez = j0ez + 1;
	j1Ex = j1ex - 1;
	j1Ey = j1ey - ((Abc == PML) ? 0 : 1);
	j1Ez = j1ez - 1;

	k0Ex = k0ex + 1;
	k0Ey = k0ey + 1;
	k0Ez = k0ez + ((Abc == PML) ? 0 : 1);
	k1Ex = k1ex - 1;
	k1Ey = k1ey - 1;
	k1Ez = k1ez - ((Abc == PML) ? 0 : 1);
	//printf("E i %d %d Ex %d - %d Ey %d - %d Ez %d - %d\n", commsize, commrank, i0Ex, i1Ex, i0Ey, i1Ey, i0Ez, i1Ez); fflush(stdout);
	//printf("E j %d %d Ex %d - %d Ey %d - %d Ez %d - %d\n", commsize, commrank, j0Ex, j1Ex, j0Ey, j1Ey, j0Ez, j1Ez); fflush(stdout);
	//printf("E k %d %d Ex %d - %d Ey %d - %d Ez %d - %d\n", commsize, commrank, k0Ex, k1Ex, k0Ey, k1Ey, k0Ez, k1Ez); fflush(stdout);

	// H array index = E array index

	NHXI = NEXI;
	NHXJ = NEXJ;
	NHXK = NEXK;
	NHX0 = NEX0;

	NHYI = NEYI;
	NHYJ = NEYJ;
	NHYK = NEYK;
	NHY0 = NEY0;

	NHZI = NEZI;
	NHZJ = NEZJ;
	NHZK = NEZK;
	NHZ0 = NEZ0;

	NH = NE;

	// H array bound = to be addressed

	i0Hx = (commrank == 0) ?  0 : i0Ex;
	i0Hy = (commrank == 0) ? -1 : i0Ey - 1;
	i0Hz = (commrank == 0) ? -1 : i0Ez - 1;
	i1Hx = (commrank == commsize - 1) ? Nx : i1Ex + 1;
	i1Hy = (commrank == commsize - 1) ? Nx : i1Ey + 1;
	i1Hz = (commrank == commsize - 1) ? Nx : i1Ez + 1;
	//printf("H i %d %d Hx %d - %d Hy %d - %d Hz %d - %d\n", commsize, commrank, i0Hx, i1Hx, i0Hy, i1Hy, i0Hz, i1Hz); fflush(stdout);

	j0Hx = -1;
	j0Hy = 0;
	j0Hz = -1;
	j1Hx = Ny;
	j1Hy = Ny;
	j1Hz = Ny;

	k0Hx = -1;
	k0Hy = -1;
	k0Hz = 0;
	k1Hx = Nz;
	k1Hy = Nz;
	k1Hz = Nz;
/*
	printf("%d %d %zd %zd %zd\n", commsize, commrank, NHX(i0Hx, j0Hx, k0Hx), NHX(i1Hx, j1Hx, k1Hx), NH); fflush(stdout);
	printf("%d %d %zd %zd %zd\n", commsize, commrank, NHY(i0Hy, j0Hy, k0Hy), NHY(i1Hy, j1Hy, k1Hy), NH); fflush(stdout);
	printf("%d %d %zd %zd %zd\n", commsize, commrank, NHZ(i0Hz, j0Hz, k0Hz), NHZ(i1Hz, j1Hz, k1Hz), NH); fflush(stdout);
*/
	assert(NHX(i0Hx, j0Hx, k0Hx) >= 0 && NHX(i1Hx, j1Hx, k1Hx) < NH);
	assert(NHY(i0Hy, j0Hy, k0Hy) >= 0 && NHY(i1Hy, j1Hy, k1Hy) < NH);
	assert(NHZ(i0Hz, j0Hz, k0Hz) >= 0 && NHZ(i1Hz, j1Hz, k1Hz) < NH);

	// MPI, index and alloc

	// share boundary
	if (commsize > 1) {
		if (commrank > 0) {
			// Ey->Ez
			const int64_t send0 = NEY(i0Ey, j0Ey, k0Ey);
			const int64_t send1 = NEZ(i0Ez, j1Ez, k1Ez);
			Send0_m       = send0;
			Sendcount_m   = (int)(send1 - send0 + 1);
			Sendcomplex_m = (d_complex_t *)malloc(Sendcount_m * sizeof(d_complex_t));
			Sendbuf_m     = (double *)malloc(2 * Sendcount_m * sizeof(double));

			// Ey->Ez->Ex
			const int64_t recv0 = NEY(i0Ey - 1, j0Ey, k0Ey);
			const int64_t recv1 = NEX(i0Ex - 1, j1Ex, k1Ex);
			Recv0_m       = recv0;
			Recvcount_m   = (int)(recv1 - recv0 + 1);
			Recvcomplex_m = (d_complex_t *)malloc(Recvcount_m * sizeof(d_complex_t));
			Recvbuf_m     = (double *)malloc(2 * Recvcount_m * sizeof(double));
		}
		if (commrank < commsize - 1) {
			// Ey->Ez->Ex
			const int64_t send0 = NEY(i1Ey, j0Ey, k0Ey);
			const int64_t send1 = NEX(i1Ex, j1Ex, k1Ex);
			Send0_p       = send0;
			Sendcount_p   = (int)(send1 - send0 + 1);
			Sendcomplex_p = (d_complex_t *)malloc(Sendcount_p * sizeof(d_complex_t));
			Sendbuf_p     = (double *)malloc(2 * Sendcount_p * sizeof(double));

			// Ey->Ez
			const int64_t recv0 = NEY(i1Ey + 1, j0Ey, k0Ey);
			const int64_t recv1 = NEZ(i1Ez + 1, j1Ez, k1Ez);
			//const int64_t recv0 = NEY(i1Ey + 1, j0Ey - (Ypbc ? 1 : 0), k0Ey);
			//const int64_t recv1 = NEZ(i1Ez + 1, j1Ez + (Ypbc ? 1 : 0), k1Ez);
			Recv0_p       = recv0;
			Recvcount_p   = (int)(recv1 - recv0 + 1);
			Recvcomplex_p = (d_complex_t *)malloc(Recvcount_p * sizeof(d_complex_t));
			Recvbuf_p     = (double *)malloc(2 * Recvcount_p * sizeof(double));
		}
	}

	// X-PBC
	if (commsize > 1) {
		if (commrank == 0) {
			// Ex->Ey->Ez
			const int64_t send0 = NEX(0, j0Ex, k0Ex);
			const int64_t send1 = NEZ(1, j1Ez, k1Ez);
			Send0_Xpbc_m       = send0;
			Sendcount_Xpbc_m   = (int)(send1 - send0 + 1);
			Sendcomplex_Xpbc_m = (d_complex_t *)malloc(Sendcount_Xpbc_m * sizeof(d_complex_t));
			Sendbuf_Xpbc_m     = (double *)malloc(2 * Sendcount_Xpbc_m * sizeof(double));

			// Ey->Ez->Ex
			const int64_t recv0 = NEY(-1, j0Ey, k0Ey);
			const int64_t recv1 = NEX(-1, j1Ex, k1Ex);
			Recv0_Xpbc_m       = recv0;
			Recvcount_Xpbc_m   = (int)(recv1 - recv0 + 1);
			Recvcomplex_Xpbc_m = (d_complex_t *)malloc(Recvcount_Xpbc_m * sizeof(d_complex_t));
			Recvbuf_Xpbc_m     = (double *)malloc(2 * Recvcount_Xpbc_m * sizeof(double));
		}
		if (commrank == commsize - 1) {
			// Ey->Ez->Ex
			const int64_t send0 = NEY(Nx - 1, j0Ey, k0Ey);
			const int64_t send1 = NEX(Nx - 1, j1Ex, k1Ex);
			Send0_Xpbc_p       = send0;
			Sendcount_Xpbc_p   = (int)(send1 - send0 + 1);
			Sendcomplex_Xpbc_p = (d_complex_t *)malloc(Sendcount_Xpbc_p * sizeof(d_complex_t));
			Sendbuf_Xpbc_p     = (double *)malloc(2 * Sendcount_Xpbc_p * sizeof(double));

			// Ex->Ey->Ez
			const int64_t recv0 = NEX(Nx + 0, j0Ex, k0Ex);
			const int64_t recv1 = NEZ(Nx + 1, j1Ez, k1Ez);
			Recv0_Xpbc_p       = recv0;
			Recvcount_Xpbc_p   = (int)(recv1 - recv0 + 1);
			Recvcomplex_Xpbc_p = (d_complex_t *)malloc(Recvcount_Xpbc_p * sizeof(d_complex_t));
			Recvbuf_Xpbc_p     = (double *)malloc(2 * Recvcount_Xpbc_p * sizeof(double));
		}
	}
}
