#include <stdio.h>
#include <math.h>
#include <stdlib.h>
#include <math.h>
#include "mpi.h"

#define N 10000 //Длина стороны сетки
#define NUM_EPOCHS 300 //Количество шагов
#define SIGMA 0.36
#define K_X 1.228
#define K_Y 1.1488

int main(int argc, char **argv)
{
    int   myid, numprocs;
    double startwtime = 0.0, endwtime;

    MPI_Init(&argc, &argv);
    MPI_Comm_size(MPI_COMM_WORLD,&numprocs);
    MPI_Comm_rank(MPI_COMM_WORLD,&myid);
    MPI_Status status;
    MPI_Request request;

    //Каждому потоку по куску сетки
    int myWidth = myid <= N % numprocs ? N / numprocs + 1 : N / numprocs;//ширина куска сетки
    double **myField = (double **) malloc(myWidth * sizeof(double*));
    double **myOldField = (double **) malloc(myWidth * sizeof(double*));
    int i, j;
    for(i = 0; i < myWidth; ++i) {
        myField[i] = (double *) malloc(N * sizeof(double));
        myOldField[i] = (double *) malloc(N * sizeof(double));
        for(j = 0; j < N; ++j) {
            myField[i][j] = 0;
            myOldField[i][j] = sin(2 * M_PI * (K_X * j + K_Y * i));
        }
    }
    double highLine[N], lowLine[N];
    int curEpoch;

    startwtime = MPI_Wtime();

    for(curEpoch = 0; curEpoch < NUM_EPOCHS; ++curEpoch) {
        if(myid < numprocs-1) {
            MPI_Isend(myOldField[myWidth-1], N, MPI_DOUBLE, myid+1, 0, MPI_COMM_WORLD, &request);
            MPI_Barrier(MPI_COMM_WORLD);
            MPI_Recv(lowLine, N, MPI_DOUBLE, myid+1, 0, MPI_COMM_WORLD, &status);
        }
        if(myid > 0) {
            MPI_Isend(myOldField[0], N, MPI_DOUBLE, myid-1, 0, MPI_COMM_WORLD, &request);
	    MPI_Barrier(MPI_COMM_WORLD);
            MPI_Recv(highLine, N, MPI_DOUBLE, myid-1, 0, MPI_COMM_WORLD, &status);
            //printf("%f\n", highLine[0]);
        }

        for(i = 0; i < myWidth; ++i) {
            for(j = 0; j < N; ++j) {
                myField[i][j] = SIGMA * (
                    (i < myWidth - 1) ? myOldField[i+1][j] : lowLine[j] +
                    (i > 0) ? myOldField[i-1][j] : highLine[j] +
                    (j > 0) ? myOldField[i][j-1] : 0 +
                    (j < N - 1) ? myOldField[i][j+1] : 0
                ) + myOldField[i][j];

            }
        }

        for(i = 0; i < myWidth; ++i)
            for(j = 0; j < N; ++j)
                myOldField[i][j] = myField[i][j];

        //MPI_Barrier(MPI_COMM_WORLD);
    }
    if(myid==0)
    {
        endwtime = MPI_Wtime();
        printf("wall clock time =%f\n", endwtime-startwtime);
        fflush(stdout);
    }
    int k;
    FILE* res = fopen("res.txt", "w");
    if (res == NULL){
        perror("File res is not open");
        return 1;
    }
    for(k = 0; k < numprocs; ++k) {
        if(myid == k) {
            fprintf(res, "I am number %d, my chunk:\n", myid);
            for(i = 0; i < myWidth; ++i){
                for(j = 0; j < N; ++j)
                    fprintf(res,"%0.4f ", myOldField[i][j]);
                fprintf(res,"\n");
            }

        }
        MPI_Barrier(MPI_COMM_WORLD);
    }
    fclose(res);
    MPI_Finalize();
    return 0;
}