#include <stdlib.h>
#include <unistd.h>
#include <stdio.h>
#include <string.h>
#include <sys/time.h>
#include <math.h>
#include "/usr/include/mpich/mpi.h"
#define ENABLE_GNUPLOT
double dx;
double dt;
#define A 0.5
#define LEFT_NODE 1
#define RIGHT_NODE 2
#define LENGTH_X 20
#define LENGTH_T 100

double dx2;
double dt2;
double a2 ;
FILE *out,*script;
struct timeval tv1,tv2,dtv;
struct timezone tz;

void writeInFile(double* Z,int lengthX, double currentTime) {
    int i;
    for(i = 0; i < lengthX; i++)
        fprintf(out,"%d\t%lf\n",i * (int)dx, Z[i]);
    fprintf(out,"\n");
    fprintf(out, "\n\n");
}


void generateGnuplotScript(int timeInterval) {
    out = fopen("out.txt","w");
    script = fopen("script.dat","w");
    fprintf(script, "set cbrange [0.9:1]\n");
    fprintf(script, "set xrange [0:%d]\n", (int)LENGTH_X - (int)dx);
    fprintf(script, "set yrange [-10:10]\n");
    fprintf(script, "set palette defined (1 '#ce4c7d')\n");
    fprintf(script, "set style line 1 lc rgb '#b90046' lt 1 lw 0.5\n");
    fprintf(script, "do for [i=0:%d]{\n", (int) LENGTH_T - 1);
    fprintf(script, "plot 'out.txt' index i using 1:2 smooth bezier title 'powered by Artur'\npause 0.1}\npause -1\n");
}



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


int timerStop() {
    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;
}


void setInitialValues(double* Z0, double* Z1, int lengthX) {
    int i;  // номер столбца
    for(i = 0; i < lengthX; i++) {
        Z0[i] = Z1[i] = 0; //sin(i*dx*M_PI / (LENGTH_X-1));
    }
    writeInFile(Z0, lengthX, 0);
}


// Внешняя сила
double F(int x, double curtime) {
    return (curtime <= 5 && (int)x == 3)
}


void calculate(double* Z0,double* Z1,int lengthX, double currentTime,int myRank, int totalProcesses, double* zleft, double* zright) {
    // Z0 всегда текущий, а Z1 - предыдущий
    // но Z0 может быть как z0 так и z1;
    int i;
    int index;
    if (myRank != 0) {
        index = 0;
        Z1[index] = dt2 * (a2*(*zleft-2*Z0[index]+Z0[index+1])/dx2 + F(myRank*lengthX + index, currentTime)) + 2*Z0[index]-Z1[index];
        *zleft = Z1[index];
    }
    for(i = 1; i < lengthX - 1; i++) {  // левый и правый столбец тоже для него не хватает данных
        index = i;
        Z1[index] = dt2 * (a2*(Z0[index-1]-2*Z0[index]+Z0[index+1])/dx2 + F(myRank*lengthX + index, currentTime) ) + 2*Z0[index]-Z1[index];
    }
    if (myRank != totalProcesses - 1) {
        index = lengthX-1;
        Z1[index] = dt2 * (a2*(Z0[index-1]-2*Z0[index]+*zright)/dx2 + F(myRank*lengthX + index, currentTime) ) + 2*Z0[index]-Z1[index];
        *zright = Z1[index];
    }
    // высылаем крайние узлы
    if(myRank != totalProcesses - 1) {  // если это не последний процесс, то нужен правый ряд
        // кидаем правую часть следующему процессу
        MPI_Send((void*)zright, 1, MPI_DOUBLE, myRank + 1, LEFT_NODE, MPI_COMM_WORLD);
        // ловим левую часть следующего процесса, для вычисляющего процесса это будет правой
        MPI_Recv((void*)zright, 1, MPI_DOUBLE, myRank + 1, RIGHT_NODE, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
    }
    if(myRank != 0) {  // если это не нулевой процесс, то нужно получить левый ряд
        // кидаем левую часть предыдущему процессу
        MPI_Send((void*)zleft, 1, MPI_DOUBLE, myRank - 1, RIGHT_NODE,MPI_COMM_WORLD);
        // ловим правую часть предыдущего процесса, для вычисляющего процесса это будет левой
        MPI_Recv((void*)zleft, 1, MPI_DOUBLE, myRank-1, LEFT_NODE, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
    }
}


int main(int argc, char **argv) {
    int myRank, totalProcesses;
    if (argc < 3) {
        printf("./a.out lengthX timeInterval\n");
        exit(0);
    }
    int lengthX = atoi(argv[1]);
    double timeInterval = atof(argv[2]); // временной интервал моделирования

    dx = (lengthX > LENGTH_X) ? dx = 1.0 : dx = LENGTH_X / lengthX;

    dt = (timeInterval > LENGTH_T) ? dt = 1.0 : dt = timeInterval / LENGTH_T;

    dx2 = dx * dx;
    dt2 = dt * dt;
    a2 = A * A;

    double* Z0;
    double* Z1;
    double currentTime = dt; // текущий момент времени
    int znumber = 1; // в какой массив записывать

    // Инициализация коммуникационных средств MPI
    MPI_Init (&argc, &argv);
    MPI_Comm_size (MPI_COMM_WORLD, &totalProcesses);
    MPI_Comm_rank (MPI_COMM_WORLD, &myRank);

    int newlengthX = lengthX / totalProcesses; // узлов на один процесс
    int lengthXmod = lengthX % totalProcesses; // остаток
    double* z0;
    double* z1;
    double zleft;
    double zright;
    double* tmp;

    if (myRank == 0) {
        generateGnuplotScript(timeInterval);
        Z0 = (double*) calloc (lengthX, sizeof(double)); // значение узлов в текущий момент времени
        Z1 = (double*) calloc (lengthX, sizeof(double)); // значение узлов в предыдущий момент времени
        setInitialValues(Z0, Z1, lengthX);
    }

    if(myRank == totalProcesses - 1) {
        newlengthX += lengthXmod;
    }

    // для каждого процесса выделяем память
    z0 = (double*) calloc (newlengthX, sizeof(double));
    z1 = (double*) calloc (newlengthX, sizeof(double));

    // каждому процессу передадим необходимую часть
    int* sendArr;
    sendArr = (int *)calloc(totalProcesses,sizeof(int));

    int i;
    for(i = 0; i < totalProcesses; i++)
        sendArr[i] = i * newlengthX;

    int* sendArrCount;
    sendArrCount = (int *)calloc(totalProcesses,sizeof(int));
    for(i = 0; i < totalProcesses; i++)
        sendArrCount[i] = newlengthX;
    sendArrCount[totalProcesses - 1] += lengthXmod;

    MPI_Scatterv((void *)(Z0), sendArrCount, sendArr, MPI_DOUBLE,(void *)(z0), newlengthX, MPI_DOUBLE, 0, MPI_COMM_WORLD); // разбрызгиватель
    MPI_Scatterv((void *)(Z1), sendArrCount, sendArr, MPI_DOUBLE,(void *)(z1), newlengthX, MPI_DOUBLE, 0, MPI_COMM_WORLD);

    zleft = z0[0];
    zright = z0[newlengthX-1];

    if(myRank != totalProcesses - 1) {  // если это не последний процесс, то нужен правый ряд
        // кидаем правую часть следующему процессу
        MPI_Send((void*)&zright, 1, MPI_DOUBLE, myRank + 1, LEFT_NODE, MPI_COMM_WORLD);
        // ловим левую часть следующего процесса, для вычисляющего процесса это будет правой
        MPI_Recv((void*)&zright, 1, MPI_DOUBLE, myRank + 1, RIGHT_NODE, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
    }

    if(myRank!= 0) {  // если это не нулевой процесс, то нужно получить левый ряд
        // кидаем левую часть предыдущему процессу
        MPI_Send((void*)&zleft, 1, MPI_DOUBLE, myRank - 1, RIGHT_NODE, MPI_COMM_WORLD);
        // ловим правую часть предыдущего процесса, для вычисляющего процесса это будет левой
        MPI_Recv((void*)&zleft, 1, MPI_DOUBLE, myRank-1, LEFT_NODE, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
    }

    if (myRank == 0)
        timerStart();
    while(currentTime < timeInterval) {
        if(znumber == 1)
            calculate(z0, z1, newlengthX, currentTime, myRank, totalProcesses, &zleft, &zright);
        else
            calculate(z1, z0, newlengthX, currentTime, myRank, totalProcesses, &zleft, &zright);

        #ifdef ENABLE_GNUPLOT
            if(znumber == 1) {
                MPI_Gatherv((void *)z1, newlengthX, MPI_DOUBLE,(void *)(Z1),sendArrCount, sendArr, MPI_DOUBLE, 0, MPI_COMM_WORLD);
                if (myRank == 0)
                    writeInFile(Z1, lengthX, currentTime); // запись в gnuplot
            }
            else {
                MPI_Gatherv((void *)z0, newlengthX, MPI_DOUBLE,(void *)(Z0),sendArrCount, sendArr, MPI_DOUBLE, 0, MPI_COMM_WORLD); // совок
                if (myRank == 0)
                    writeInFile(Z0, lengthX, currentTime); // запись в gnuplot
            }
            MPI_Barrier(MPI_COMM_WORLD); // ждем запись
        #endif

        if(!myRank) {
            znumber *= -1;
            currentTime += dt;
        }

        MPI_Bcast((void *)&znumber, 1, MPI_INT, 0, MPI_COMM_WORLD);
        // Передаём текущее значение всем процессам на текущем шаге для того, чтобы правильно рассчитывалось внешнее воздействие на каждый узел, которое зависит от координаты и времени
        MPI_Bcast((void *)&currentTime, 1, MPI_DOUBLE, 0, MPI_COMM_WORLD);
    }

    // Вывод времени работы
    if (myRank == 0) {
        int ms = timerStop();
        printf("Time: %d milliseconds\n", ms);
        free(Z0);
        free(Z1);
    }

    free(z0);
    free(z1);
    free(sendArr);
    free(sendArrCount);
    MPI_Finalize();
    exit(0);
}