#include <stdlib.h>
#include <unistd.h>
#include <stdio.h>
#include <string.h>
#include <sys/time.h>
#include <math.h>
#include <stdbool.h>
#include <mpi.h>
#define WRITETOGNUPOT //вывод в файл
double dx = 0.1;
double dt = 0.01;
double a = 0.1;
#define LEFT 1
#define RIGHT 2
double dx2;
double a2;
struct timeval tv1,tv2,dtv;
struct timezone tz;
int lengthX;

void time_start() {
	gettimeofday(&tv1, &tz);
}

int time_stop() {
	gettimeofday(&tv2, &tz);
	dtv.tv_sec = tv2.tv_sec - tv1.tv_sec;
	dtv.tv_usec = tv2.tv_usec - tv1.tv_usec;
	if (dtv.tv_usec < 0) {
		dtv.tv_sec--;
		dtv.tv_usec += 1000000;
	}
	return dtv.tv_sec * 1000 + dtv.tv_usec / 1000;
}

FILE *out,*fp;
void writeIntoFile(double* Z, int lengthX, double CurrentTime) { //вывод в файл
	int i;
	for (i = 0; i < lengthX; i++) {
		fprintf(out, "%d\t%lf\n", i, Z[i]);
	}
	fprintf(out, "\n\n");
}

void openFiles(int lengthX, int timeInterval) { //открытие файлов
	out = fopen("out.txt", "w");
	fp = fopen("vgraph0.dat", "w");
	fprintf(fp, "set xrange[0:%d]\nset yrange[-100:100]\n", lengthX - 1);
	fprintf(fp, "do for [i=0:%d]{\n", timeInterval - 1);
	fprintf(fp, "plot 'out.txt' index i using 1:2 smooth bezier\npause 0.1}\npause -1\n");
}

void calculateFisrtTime(double* Z0, double* Z1, int lengthX) { //инициализация массивов
	int i;
	int prev = 0;
	memset(Z0, lengthX * sizeof(double), 0);
	memset(Z1, lengthX * sizeof(double), 0);
	Z0[lengthX - 1] = 50;
	Z1[lengthX - 1] = 50;
}

void calculate(double* Z0, double* Z1, int lengthX, double CurrentTime, int myrank, int total, double* zleft, double* zright) { //функция расчетов новых значений
	int i = 0;
	if (myrank != 0) {
		Z1[i] = dt * a * ((*zleft - 2 * Z0[i] + Z0[i + 1]) / dx2) + Z0[i];
		*zleft = Z1[i];
	}
	else
		Z1[0] = Z0[1] + 20;


		// printf("%f -2 %f %f  %f\n", Z0[0], Z0[1], Z0[2], Z0[1]);
	for (i = 1; i < lengthX - 1; i++) {
		Z1[i] = dt * a * (Z0[i - 1] - 2 * Z0[i] + Z0[i + 1]) / dx2 + Z0[i];
	}

	if (myrank != total - 1) {
		Z1[i] = dt * a * ((Z0[i - 1] - 2 * Z0[i] + *zright) / dx2) + Z0[i];
		*zright = Z1[i];
	}
	else
		Z1[lengthX - 1] = 50;
	// отправка и получение соседних значений
	if (myrank != total - 1) {
		MPI_Send((void*)(zright), 1, MPI_DOUBLE, myrank + 1, LEFT, MPI_COMM_WORLD);
		MPI_Recv((void*)(zright), 1, MPI_DOUBLE, myrank + 1, RIGHT, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
	}
	if (myrank != 0) {
		MPI_Send((void*)(zleft), 1, MPI_DOUBLE, myrank - 1, RIGHT, MPI_COMM_WORLD);
		MPI_Recv((void*)(zleft), 1, MPI_DOUBLE, myrank - 1, LEFT, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
	}
}

int main(int argc, char **argv) {
	int myrank, total;
	if (argc < 3) {
		printf("./a.out #lengthX #timeInterval\n");
		exit(0);
	}
	lengthX = atoi(argv[1]);
	double timeInterval = atof(argv[2]);
	dx2 = dx * dx;
	double *Z0, *Z1;
	double CurrentTime = 0;
	int znumber = 1;

	MPI_Init(&argc, &argv);
	MPI_Comm_size(MPI_COMM_WORLD, &total);
	MPI_Comm_rank(MPI_COMM_WORLD, &myrank);
	if (lengthX % total != 0) {
		printf("Incorrect!\n");
		exit(0);
	}
	int newlengthX = lengthX / total;
	double* z0;
	double* z1;
	double zleft;
	double zright;
	if (myrank == 0) { //инициализация массивов в основном процессе
		Z0 = (double*)calloc(lengthX, sizeof(double));
		Z1 = (double*)calloc(lengthX, sizeof(double));
		calculateFisrtTime(Z0, Z1, lengthX);
		#ifdef WRITETOGNUPOT
			openFiles(lengthX, timeInterval);
			writeIntoFile(Z0, lengthX, CurrentTime);
		#endif
	}
	CurrentTime += 1;
	z0 = (double*)calloc(newlengthX, sizeof(double));
	z1 = (double*)calloc(newlengthX, sizeof(double));
	int* displsSend;
	displsSend = (int *)calloc(total, sizeof(int));
	int numberOfProc;
	for (numberOfProc = 0; numberOfProc < total; numberOfProc++)
		displsSend[numberOfProc] = numberOfProc * newlengthX;
	int* displsSendCount;
	displsSendCount = (int *)calloc(total, sizeof(int));
	for (numberOfProc = 0; numberOfProc < total; numberOfProc++)
		displsSendCount[numberOfProc] = newlengthX;
	int i;
	for (i = 0; i < total; i++) { //рассылка процессам данных
		MPI_Scatterv((void *)(Z0), displsSendCount, displsSend, MPI_DOUBLE, (void *)(z0), newlengthX, MPI_DOUBLE, 0, MPI_COMM_WORLD);
		MPI_Scatterv((void *)(Z1), displsSendCount, displsSend, MPI_DOUBLE, (void *)(z1), newlengthX, MPI_DOUBLE, 0, MPI_COMM_WORLD);
	}
	zleft = z0[0];
	zright = z0[newlengthX - 1];
	//рассылка соседних значений
	if (myrank != total - 1) {
		MPI_Send((void*)(&zright), 1, MPI_DOUBLE, myrank + 1, LEFT, MPI_COMM_WORLD);
		MPI_Recv((void*)(&zright), 1, MPI_DOUBLE, myrank + 1, RIGHT, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
	}
	if (myrank != 0) {
		MPI_Send((void*)(&zleft), 1, MPI_DOUBLE, myrank - 1, RIGHT, MPI_COMM_WORLD);
		MPI_Recv((void*)(&zleft), 1, MPI_DOUBLE, myrank - 1, LEFT, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
	}
	if (myrank == 0)
		time_start();
	bool flag = 1;
	while (CurrentTime < timeInterval) {
		if (flag)
			calculate(z0, z1, newlengthX, CurrentTime, myrank, total, &zleft, &zright);
		else
			calculate(z1, z0, newlengthX, CurrentTime, myrank, total, &zleft, &zright);
		#ifdef WRITETOGNUPOT
			MPI_Barrier(MPI_COMM_WORLD);
			if (flag)
				MPI_Gatherv((void*)(z1), newlengthX, MPI_DOUBLE, (void*)(Z1), displsSendCount, displsSend, MPI_DOUBLE, 0, MPI_COMM_WORLD);
			else
				MPI_Gatherv((void*)(z0), newlengthX, MPI_DOUBLE, (void*)(Z0), displsSendCount, displsSend, MPI_DOUBLE, 0, MPI_COMM_WORLD);
			MPI_Barrier(MPI_COMM_WORLD);
			if (myrank == 0)
				if (flag)
					writeIntoFile(Z1, lengthX, CurrentTime);
				else
					writeIntoFile(Z0, lengthX, CurrentTime);
		#endif
		flag = !flag;
		CurrentTime += 1;
	}
	if (myrank == 0) {
		int ms = time_stop();
		printf("Time: %d milliseconds\n", ms);
	}
	MPI_Finalize();
	exit(0);
}