/*
comm.c (MPI)

MPI routines
*/

#include "oth.h"
#include "oth_mpi.h"
#include "complex.h"

#ifdef _MPI
#include "mpi.h"
#endif

// initialize
void mpi_init(int argc, char **argv)
{
#ifdef _MPI
	MPI_Init(&argc, &argv);
	MPI_Comm_size(MPI_COMM_WORLD, &commSize);
	MPI_Comm_rank(MPI_COMM_WORLD, &commRank);
#else
	commSize = 1;
	commRank = 0;
	argc = argc;	// dummy
	argv = argv;	// dummy
#endif
}

// close
void mpi_close(void)
{
#ifdef _MPI
	MPI_Finalize();
#endif
}

// check error code
// mode = 0/1 : Bcast/Allreduce
void comm_check(int ierr, int mode, int prompt)
{
#ifdef _MPI
	if (commSize > 1) {
		if (mode == 0) {
			MPI_Bcast(&ierr, 1, MPI_INT, 0, MPI_COMM_WORLD);
		}
		else {
			int g_ierr;
			MPI_Allreduce(&ierr, &g_ierr, 1, MPI_INT, MPI_LOR, MPI_COMM_WORLD);
			ierr = g_ierr;
		}
	}
	if (ierr) {
		MPI_Finalize();
	}
#endif
	mode = mode;  // dummy
	if (ierr) {
		if (prompt && (commRank == 0)) {
			fflush(stdout);
			getchar();
		}
		exit(0);
	}
}

// gather string
void comm_string(const char *str, char *lstr)
{
#ifdef _MPI
	char buff[BUFSIZ];
	if (commRank == 0) {
		MPI_Status status;
		strcpy(lstr, str);
		for (int i = 1; i < commSize; i++) {
			MPI_Recv(buff, BUFSIZ, MPI_CHAR, i, 0, MPI_COMM_WORLD, &status);
			strcat(lstr, "\n");
			strcat(lstr, buff);
		}
	}
	else {
		strcpy(buff, str);
		MPI_Send(buff, BUFSIZ, MPI_CHAR, 0, 0, MPI_COMM_WORLD);
	}
#else
	strcpy(lstr, str);
#endif
}

// get cpu time [sec]
double comm_cputime(void)
{
#ifdef _MPI
	MPI_Barrier(MPI_COMM_WORLD);
	return MPI_Wtime();
#else
#ifdef _WIN32
	return (double)clock() / CLOCKS_PER_SEC;
#else
	struct timespec ts;
	clock_gettime(CLOCK_REALTIME, &ts);
	return (ts.tv_sec + (ts.tv_nsec * 1e-9));
#endif  // _WIN32
#endif  // MPI
}

// broadcast input data
void comm_broadcast(void)
{
#ifdef _MPI
	if (commSize < 2) return;

	int    *i_buf = NULL;
	double *d_buf = NULL;
	char   *c_buf = NULL;
	int    i_num = 0;
	int    d_num = 0;
	int    c_num = 0;

	// copy variables to buffers (root only)

	if (commRank == 0) {
		// number of data
		i_num = 15                                + (2 * NGeometry) + (3 * NFeed) +  (3 * IPlanewave);
		d_num =  3 + (8 * NMaterial * NFrequency) + (8 * NGeometry) + (6 * NFeed) + (10 * IPlanewave) + NFrequency + (Nx + 1) + (Ny + 1) + (Nz + 1) + Nx + Ny + Nz;
		c_num =                                                       (1 * NFeed);

		// alloc
		if (i_num > 0) {
			i_buf =    (int *)malloc(i_num * sizeof(int));
		}
		if (d_num > 0) {
			d_buf = (double *)malloc(d_num * sizeof(double));
		}
		if (c_num > 0) {
			c_buf =   (char *)malloc(c_num * sizeof(char));
		}

		int i_id = 0;
		int d_id = 0;
		int c_id = 0;

		i_buf[i_id++] = Nx;
		i_buf[i_id++] = Ny;
		i_buf[i_id++] = Nz;
		i_buf[i_id++] = NMaterial;
		i_buf[i_id++] = NGeometry;
		i_buf[i_id++] = NFeed;
		i_buf[i_id++] = IPlanewave;
		i_buf[i_id++] = Abc;
		i_buf[i_id++] = Lpml;
		i_buf[i_id++] = Pbc.xpbc;
		i_buf[i_id++] = Pbc.ypbc;
		i_buf[i_id++] = Pbc.zpbc;
		i_buf[i_id++] = NFrequency;
		i_buf[i_id++] = Solver.maxiter;
		i_buf[i_id++] = Solver.nout;

		d_buf[d_id++] = Mpml;
		d_buf[d_id++] = R0pml;
		d_buf[d_id++] = Solver.converg;

		for (int i = 0; i <= Nx; i++) {
			d_buf[d_id++] = Xn[i];
		}
		for (int j = 0; j <= Ny; j++) {
			d_buf[d_id++] = Yn[j];
		}
		for (int k = 0; k <= Nz; k++) {
			d_buf[d_id++] = Zn[k];
		}

		for (int i = 0; i < Nx; i++) {
			d_buf[d_id++] = Xc[i];
		}
		for (int j = 0; j < Ny; j++) {
			d_buf[d_id++] = Yc[j];
		}
		for (int k = 0; k < Nz; k++) {
			d_buf[d_id++] = Zc[k];
		}

		for (int n = 0; n < NMaterial * NFrequency; n++) {
			d_buf[d_id++] = Material[n].cepsr.r;
			d_buf[d_id++] = Material[n].cepsr.i;
			d_buf[d_id++] = Material[n].cmur.r;
			d_buf[d_id++] = Material[n].cmur.i;
			d_buf[d_id++] = Material[n].cnur.r;
			d_buf[d_id++] = Material[n].cnur.i;
			d_buf[d_id++] = Material[n].cv.r;
			d_buf[d_id++] = Material[n].cv.i;
		}

		for (int n = 0; n < NGeometry; n++) {
			i_buf[i_id++] = Geometry[n].m;
			i_buf[i_id++] = Geometry[n].shape;
			for (int i = 0; i < 8; i++) {
				d_buf[d_id++] = Geometry[n].g[i];
			}
		}

		for (int n = 0; n < NFeed; n++) {
			c_buf[c_id++] = Feed[n].dir;
			i_buf[i_id++] = Feed[n].i;
			i_buf[i_id++] = Feed[n].j;
			i_buf[i_id++] = Feed[n].k;
			d_buf[d_id++] = Feed[n].volt.r;
			d_buf[d_id++] = Feed[n].volt.i;
			d_buf[d_id++] = Feed[n].z0;
			d_buf[d_id++] = Feed[n].dx;
			d_buf[d_id++] = Feed[n].dy;
			d_buf[d_id++] = Feed[n].dz;
		}

		if (IPlanewave) {
			i_buf[i_id++] = Planewave.pol;
			for (int m = 0; m < 3; m++) {
				d_buf[d_id++] = Planewave.r1[m];
				d_buf[d_id++] = Planewave.t1[m];
				d_buf[d_id++] = Planewave.p1[m];
			}
			i_buf[i_id++] = Planewave.isemi;
			d_buf[d_id++] = Planewave.psemi;
			i_buf[i_id++] = Planewave.ifull;
		}

		for (int n = 0; n < NFrequency; n++) {
			d_buf[d_id++] = Frequency[n];
		}

		// check
		assert(i_id == i_num);
		assert(d_id == d_num);
		assert(c_id == c_num);
	}

	// broadcast (root to non-root)

	MPI_Bcast(&i_num, 1, MPI_INT, 0, MPI_COMM_WORLD);
	MPI_Bcast(&d_num, 1, MPI_INT, 0, MPI_COMM_WORLD);
	MPI_Bcast(&c_num, 1, MPI_INT, 0, MPI_COMM_WORLD);

	// alloc (non-root)
	if (commRank > 0) {
		if (i_num > 0) {
			i_buf =    (int *)malloc(i_num * sizeof(int));
		}
		if (d_num > 0) {
			d_buf = (double *)malloc(d_num * sizeof(double));
		}
		if (c_num > 0) {
			c_buf =   (char *)malloc(c_num * sizeof(char));
		}
	}

	if (i_num > 0) {
		MPI_Bcast(i_buf, i_num, MPI_INT,    0, MPI_COMM_WORLD);
	}
	if (d_num > 0) {
		MPI_Bcast(d_buf, d_num, MPI_DOUBLE, 0, MPI_COMM_WORLD);
	}
	if (c_num > 0) {
		MPI_Bcast(c_buf, c_num, MPI_CHAR,   0, MPI_COMM_WORLD);
	}

	// copy buffers to variables (non-root)

	if (commRank > 0) {
		int i_id = 0;
		int d_id = 0;
		int c_id = 0;

		Nx             = i_buf[i_id++];
		Ny             = i_buf[i_id++];
		Nz             = i_buf[i_id++];
		NMaterial      = i_buf[i_id++];
		NGeometry      = i_buf[i_id++];
		NFeed          = i_buf[i_id++];
		IPlanewave     = i_buf[i_id++];
		Abc            = i_buf[i_id++];
		Lpml           = i_buf[i_id++];
		Pbc.xpbc       = i_buf[i_id++];
		Pbc.ypbc       = i_buf[i_id++];
		Pbc.zpbc       = i_buf[i_id++];
		NFrequency     = i_buf[i_id++];
		Solver.maxiter = i_buf[i_id++];
		Solver.nout    = i_buf[i_id++];

		Mpml           = d_buf[d_id++];
		R0pml          = d_buf[d_id++];
		Solver.converg = d_buf[d_id++];

		Xn = (double *)malloc((Nx + 1) * sizeof(double));
		Yn = (double *)malloc((Ny + 1) * sizeof(double));
		Zn = (double *)malloc((Nz + 1) * sizeof(double));
		for (int i = 0; i <= Nx; i++) {
			Xn[i] = d_buf[d_id++];
		}
		for (int j = 0; j <= Ny; j++) {
			Yn[j] = d_buf[d_id++];
		}
		for (int k = 0; k <= Nz; k++) {
			Zn[k] = d_buf[d_id++];
		}

		Xc = (double *)malloc(Nx * sizeof(double));
		Yc = (double *)malloc(Ny * sizeof(double));
		Zc = (double *)malloc(Nz * sizeof(double));
		for (int i = 0; i < Nx; i++) {
			Xc[i] = d_buf[d_id++];
		}
		for (int j = 0; j < Ny; j++) {
			Yc[j] = d_buf[d_id++];
		}
		for (int k = 0; k < Nz; k++) {
			Zc[k] = d_buf[d_id++];
		}

		Material = (material_t *)malloc(NMaterial * NFrequency * sizeof(material_t));
		for (int n = 0; n < NMaterial * NFrequency; n++) {
			Material[n].cepsr.r = d_buf[d_id++];
			Material[n].cepsr.i = d_buf[d_id++];
			Material[n].cmur.r  = d_buf[d_id++];
			Material[n].cmur.i  = d_buf[d_id++];
			Material[n].cnur.r  = d_buf[d_id++];
			Material[n].cnur.i  = d_buf[d_id++];
			Material[n].cv.r    = d_buf[d_id++];
			Material[n].cv.i    = d_buf[d_id++];
		}

		if (NGeometry > 0) {
			Geometry = (geometry_t *)malloc(NGeometry * sizeof(geometry_t));
			for (int n = 0; n < NGeometry; n++) {
				Geometry[n].m     = i_buf[i_id++];
				Geometry[n].shape = i_buf[i_id++];
				for (int i = 0; i < 8; i++) {
					Geometry[n].g[i] = d_buf[d_id++];
				}
			}
		}

		if (NFeed > 0) {
			Feed = (feed_t *)malloc(NFeed * sizeof(feed_t));
			for (int n = 0; n < NFeed; n++) {
				Feed[n].dir    = c_buf[c_id++];
				Feed[n].i      = i_buf[i_id++];
				Feed[n].j      = i_buf[i_id++];
				Feed[n].k      = i_buf[i_id++];
				Feed[n].volt.r = d_buf[d_id++];
				Feed[n].volt.i = d_buf[d_id++];
				Feed[n].z0     = d_buf[d_id++];
				Feed[n].dx     = d_buf[d_id++];
				Feed[n].dy     = d_buf[d_id++];
				Feed[n].dz     = d_buf[d_id++];
			}
		}

		if (IPlanewave) {
			Planewave.pol = i_buf[i_id++];
			for (int m = 0; m < 3; m++) {
				Planewave.r1[m] = d_buf[d_id++];
				Planewave.t1[m] = d_buf[d_id++];
				Planewave.p1[m] = d_buf[d_id++];
			}
			Planewave.isemi  = i_buf[i_id++];
			Planewave.psemi  = d_buf[d_id++];
			Planewave.ifull  = i_buf[i_id++];
		}

		if (NFrequency > 0) {
			Frequency = (double *)malloc(NFrequency * sizeof(double));
			for (int n = 0; n < NFrequency; n++) {
				Frequency[n] = d_buf[d_id++];
			}
		}

		// check
		assert(i_id == i_num);
		assert(d_id == d_num);
		assert(c_id == c_num);
	}

	// free
	if (i_buf != NULL) {
		free(i_buf);
	}
	if (d_buf != NULL) {
		free(d_buf);
	}
	if (c_buf != NULL) {
		free(c_buf);
	}
/*
	// debug
	//printf("%d %d\n", commSize, commRank);
	//printf("%d %d %d %d %d\n", NMaterial, NGeometry, NFeed, NFrequency, NPoint);
	//printf("%d %d %e\n", Solver.maxiter, Solver.nout, Solver.converg);
	//for (int i = 0; i <= Nx; i++) printf("%d Xn[%d]=%.5f\n", commRank, i, Xn[i] * 1e3);
	//for (int j = 0; j <= Ny; j++) printf("%d Yn[%d]=%.5f\n", commRank, j, Yn[j] * 1e3);
	//for (int k = 0; k <= Nz; k++) printf("%d Zn[%d]=%.5f\n", commRank, k, Zn[k] * 1e3);
	//for (int n = 0; n < NMaterial; n++) printf("%d %d %e %e %e %e\n", commRank, n, Material[n].cepsr.r, Material[n].cepsr.i, Material[n].cmur.r, Material[n].cmur.i);
	//for (int n = 0; n < NGeometry; n++) printf("%d %d %d %e %e %e %e %e %e\n", commRank, n, Geometry[n].m, Geometry[n].g[0], Geometry[n].g[1], Geometry[n].g[2], Geometry[n].g[3], Geometry[n].g[4], Geometry[n].g[5]);
	//for (int n = 0; n < NFeed; n++) printf("%d %d %c %d %d %d %e\n", commRank, n, Feed[n].dir, Feed[n].i, Feed[n].j, Feed[n].k, Feed[n].volt);
	//for (int n = 0; n < NFrequency; n++) printf("%d %d %e\n", commRank, n, Frequency[n]);
	//for (int n = 0; n < NPoint + 2; n++) printf("%d %d %c %d %d %d\n", commRank, n, Point[n].dir, Point[n].i, Point[n].j, Point[n].k);
	fflush(stdout);
*/
#endif
}

// sum (double complex)
d_complex_t comm_zsum(d_complex_t l_sum)
{
	d_complex_t g_sum;
	g_sum.r = l_sum.r;
	g_sum.i = l_sum.i;
#ifdef _MPI
	if (commSize > 1) {
		MPI_Allreduce(&l_sum.r, &g_sum.r, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);
		MPI_Allreduce(&l_sum.i, &g_sum.i, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);
	}
#endif
	return g_sum;
}

// sum (double)
double comm_dsum(double l_sum)
{
	double g_sum = l_sum;
#ifdef _MPI
	if (commSize > 1) {
		MPI_Allreduce(&l_sum, &g_sum, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);
	}
#endif
	return g_sum;
}

// sum (int)
int comm_isum(int l_sum)
{
	int g_sum = l_sum;
#ifdef _MPI
	if (commSize > 1) {
		MPI_Allreduce(&l_sum, &g_sum, 1, MPI_INT, MPI_SUM, MPI_COMM_WORLD);
	}
#endif
	return g_sum;
}

// gather E to Ex+Ey+Ez
void comm_gather(const d_complex_t *e, f_complex_t ***ex, f_complex_t ***ey, f_complex_t ***ez)
{
#ifdef _MPI
	if (commSize == 1) return;

	const int tag = 0;
	if (!commRank) {
		// root

		// self copy
		for (int i = i0Ex; i <= i1Ex; i++) {
		for (int j = j0Ex; j <= j1Ex; j++) {
		for (int k = k0Ex; k <= k1Ex; k++) {
			if ((i >= 0) && (i < Nx) && (j >= 0) && (j <= Ny) && (k >= 0) && (k <= Nz)) {
				ex[i][j][k] = f_cast(e[NEX(i, j, k)]);
			}
		}
		}
		}
		for (int i = i0Ey; i <= i1Ey; i++) {
		for (int j = j0Ey; j <= j1Ey; j++) {
		for (int k = k0Ey; k <= k1Ey; k++) {
			if ((i >= 0) && (i <= Nx) && (j >= 0) && (j < Ny) && (k >= 0) && (k <= Nz)) {
				ey[i][j][k] = f_cast(e[NEY(i, j, k)]);
			}
		}
		}
		}
		for (int i = i0Ez; i <= i1Ez; i++) {
		for (int j = j0Ez; j <= j1Ez; j++) {
		for (int k = k0Ez; k <= k1Ez; k++) {
			if ((i >= 0) && (i <= Nx) && (j >= 0) && (j <= Ny) && (k >= 0) && (k < Nz)) {
				ez[i][j][k] = f_cast(e[NEZ(i, j, k)]);
			}
		}
		}
		}

		// recv
		MPI_Status status;
		int nrecv[6];
		for (int rank = 1; rank < commSize; rank++) {
			// recv i-range
			MPI_Recv(nrecv, 6, MPI_INT, rank, 0, MPI_COMM_WORLD, &status);
			//printf("%d %d %d %d %d %d %d\n", rank, nrecv[0], nrecv[1], nrecv[2], nrecv[3], nrecv[4], nrecv[5]); fflush(stdout);

			// recv E
			const int recvcount = 2 * (
				(nrecv[3] - nrecv[0] + 1) * (j1Ex - j0Ex + 1) * (k1Ex - k0Ex + 1) +
				(nrecv[4] - nrecv[1] + 1) * (j1Ey - j0Ey + 1) * (k1Ey - k0Ey + 1) +
				(nrecv[5] - nrecv[2] + 1) * (j1Ez - j0Ez + 1) * (k1Ez - k0Ez + 1));
			float *recvbuf = (float *)malloc(recvcount * sizeof(float));
			//double *recvbuf = (double *)malloc(recvcount * sizeof(double));
			//printf("%d %d\n", rank, recvcount); fflush(stdout);
			MPI_Recv(recvbuf, recvcount, MPI_FLOAT, rank, tag, MPI_COMM_WORLD, &status);
			//MPI_Recv(recvbuf, recvcount, MPI_DOUBLE, rank, tag, MPI_COMM_WORLD, &status);

			int irecv = 0;
			for (int i = nrecv[0]; i <= nrecv[3]; i++) {
			for (int j = j0Ex; j <= j1Ex; j++) {
			for (int k = k0Ex; k <= k1Ex; k++) {
				float re = recvbuf[irecv++];
				float im = recvbuf[irecv++];
				if ((i >= 0) && (i < Nx) && (j >= 0) && (j <= Ny) && (k >= 0) && (k <= Nz)) {
					ex[i][j][k] = f_complex(re, im);
				}
			}
			}
			}
			for (int i = nrecv[1]; i <= nrecv[4]; i++) {
			for (int j = j0Ey; j <= j1Ey; j++) {
			for (int k = k0Ey; k <= k1Ey; k++) {
				float re = recvbuf[irecv++];
				float im = recvbuf[irecv++];
				if ((i >= 0) && (i <= Nx) && (j >= 0) && (j < Ny) && (k >= 0) && (k <= Nz)) {
					ey[i][j][k] = f_complex(re, im);
				}
			}
			}
			}
			for (int i = nrecv[2]; i <= nrecv[5]; i++) {
			for (int j = j0Ez; j <= j1Ez; j++) {
			for (int k = k0Ez; k <= k1Ez; k++) {
				float re = recvbuf[irecv++];
				float im = recvbuf[irecv++];
				if ((i >= 0) && (i <= Nx) && (j >= 0) && (j <= Ny) && (k >= 0) && (k < Nz)) {
					ez[i][j][k] = f_complex(re, im);
				}
			}
			}
			}
			assert(irecv == recvcount);
		}

	}
	else {
		// non-root : send

		// send i-range
		int nsend[6];
		nsend[0] = i0Ex;
		nsend[1] = i0Ey;
		nsend[2] = i0Ez;
		nsend[3] = i1Ex;
		nsend[4] = i1Ey;
		nsend[5] = i1Ez;
		MPI_Send(nsend, 6, MPI_INT, 0, tag, MPI_COMM_WORLD);

		// send E
		const int sendcount = 2 * (
			(i1Ex - i0Ex + 1) * (j1Ex - j0Ex + 1) * (k1Ex - k0Ex + 1) +
			(i1Ey - i0Ey + 1) * (j1Ey - j0Ey + 1) * (k1Ey - k0Ey + 1) +
			(i1Ez - i0Ez + 1) * (j1Ez - j0Ez + 1) * (k1Ez - k0Ez + 1));
		//printf("%d %d\n", commRank, sendcount); fflush(stdout);
		float *sendbuf = (float *)malloc(sendcount * sizeof(float));
		//double *sendbuf = (double *)malloc(sendcount * sizeof(double));
		int isend = 0;
		for (int i = i0Ex; i <= i1Ex; i++) {
		for (int j = j0Ex; j <= j1Ex; j++) {
		for (int k = k0Ex; k <= k1Ex; k++) {
			sendbuf[isend++] = (float)e[NEX(i, j, k)].r;
			sendbuf[isend++] = (float)e[NEX(i, j, k)].i;
		}
		}
		}
		for (int i = i0Ey; i <= i1Ey; i++) {
		for (int j = j0Ey; j <= j1Ey; j++) {
		for (int k = k0Ey; k <= k1Ey; k++) {
			sendbuf[isend++] = (float)e[NEY(i, j, k)].r;
			sendbuf[isend++] = (float)e[NEY(i, j, k)].i;
		}
		}
		}
		for (int i = i0Ez; i <= i1Ez; i++) {
		for (int j = j0Ez; j <= j1Ez; j++) {
		for (int k = k0Ez; k <= k1Ez; k++) {
			sendbuf[isend++] = (float)e[NEZ(i, j, k)].r;
			sendbuf[isend++] = (float)e[NEZ(i, j, k)].i;
		}
		}
		}
		assert(isend == sendcount);
		MPI_Send(sendbuf, sendcount, MPI_FLOAT, 0, tag, MPI_COMM_WORLD);
		//MPI_Send(sendbuf, sendcount, MPI_DOUBLE, 0, tag, MPI_COMM_WORLD);

	}
#endif
}
