#include <stdio.h>
#include <mpi.h>
#include <cmath>

double g(const double x)
{
	return  x * 0.9;
}

double uol(const double x)
{
	return sin(x);
}

const double f0 = 0.0;
const double f1 = 0.0;

int main(int arc, char** argv)
{
	int size, rank;
	const auto msgtag = 12;

	auto x = 0.0, time = 0.0;
	const auto tau = 0.2;
	const auto tmax = 1.0;
	const auto L = 10.0;
	auto l = 0;

	const auto h = 0.1;

	const auto r = tau * tau / (h * h);

	const auto N = static_cast<int>(L / h) + 1;

	//шапочка

	if (MPI_Init(&arc, &argv) != MPI_SUCCESS) return 1;

	if (MPI_Comm_size(MPI_COMM_WORLD, &size) != MPI_SUCCESS)
	{
		MPI_Finalize();
		return 2;
	}

	if (MPI_Comm_rank(MPI_COMM_WORLD, &rank) != MPI_SUCCESS)
	{
		MPI_Finalize();
		return 3;
	}

	MPI_Status status;

	//Распределение отрезков по процессам
	auto l1 = N / size;
	auto l2 = N % size;
	const auto kol = new int[size];

	for (auto i = 0; i < size; i++) kol[i] = l1;

	if (l2)
	{
		if (rank == size - 1) l1++;
		l2--;
		kol[size - 1]++;
	}

	if (l2)
	{
		if (rank < l2) l1++;
		for (auto i = 0; i < l2; i++) kol[i]++;
	}

	auto n = l1;

	if (!rank || rank == size - 1) n++; else n += 2;

	auto u = new double[n];
	auto uold = new double[n];

	auto sum = 0;


	if (rank)
	{
		for (auto i = 0; i < rank; i++) sum += kol[i];
		l = sum;
		sum--;
		x = h * sum; //Вычисляем смещение для каждого промежутка
	}

	//Присваиваем начальные значения

	for (auto i = 0; i < n; i++)
	{
		u[i] = g(x);
		uold[i] = u[i] - tau * uol(x);
		x += h;
	}

	if (!rank) uold[0] = u[0] = f0;
	if (rank == size - 1) uold[n - 1] = u[n - 1] = f1;

	const auto tn = MPI_Wtime(); //Начальное время
	do
	{
		const auto xn = h * sum;

		for (auto i = 1; i < n - 1; i++)
		{
			const auto uu = u[i];
			u[i] = r * ((xn + i * h + h) * (xn + i * h + h) * (u[i + 2] - u[i + 1]) - (xn + i * h) * (xn + i * h) * (u[i + 1] - u[i])) + 2.0 * u[i] - uold[i];
			uold[i] = uu;
		}

		//Начальные значения в краевых точках
		if (!rank) u[0] = f0;
		if (rank == size - 1) u[n - 1] = f1;

		//Обмен данными между процессами
		if (rank & 1)
		{
			MPI_Ssend(&u[1], 1, MPI_DOUBLE, rank - 1, msgtag, MPI_COMM_WORLD);
			MPI_Recv(&u[0], 1, MPI_DOUBLE, rank - 1, msgtag, MPI_COMM_WORLD, &status);

			if (rank != size - 1)
			{
				MPI_Ssend(&u[n - 2], 1, MPI_DOUBLE, rank + 1, msgtag, MPI_COMM_WORLD);
				MPI_Recv(&u[n - 1], 1, MPI_DOUBLE, rank + 1, msgtag, MPI_COMM_WORLD, &status);
			}
		}

		else
		{
			if (rank != size - 1) {
				MPI_Recv(&u[n - 1], 1, MPI_DOUBLE, rank + 1, msgtag, MPI_COMM_WORLD, &status);
				MPI_Ssend(&u[n - 2], 1, MPI_DOUBLE, rank + 1, msgtag, MPI_COMM_WORLD);
			}

			if (rank) {
				MPI_Recv(&u[0], 1, MPI_DOUBLE, rank - 1, msgtag, MPI_COMM_WORLD, &status);
				MPI_Ssend(&u[1], 1, MPI_DOUBLE, rank - 1, msgtag, MPI_COMM_WORLD);
			}
		}

		time += tau; //Шаг по времени

	} while (time <= tmax);

	const auto tk = MPI_Wtime(); //Конечное время

	//вывод
	int in, ik;
	if (!rank) in = 0; else in = 1;
	if (rank == size - 1) ik = n; else ik = n - 1;

	for (auto i = 0; i < size; i++)
	{
		if (i == rank)
		{
			const auto f = fopen("Temperature.txt", "a");  // NOLINT(android-cloexec-fopen)
			for (auto j = in; j < ik; j++, l++)
			{
				fprintf(f, "u[%d] = %f\n", l, u[j]);
			}
			fclose(f);

		}
		MPI_Barrier(MPI_COMM_WORLD);
	}

	if (!rank)
	{
		const auto f = fopen("Temperature.txt", "a");  // NOLINT(android-cloexec-fopen)
		fprintf(f, "\nВремя вычислений: %f", tk - tn);
		fclose(f);
	}

	delete[]u;
	delete[]kol;
	MPI_Finalize();
	return 0;
}