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

#define a 1e5
#define EPSILON 1e-6
#define SQUARE_SIDE_LENGTH 2.0
#define X_NULL -1
#define Y_NULL -1
#define Z_NULL -1

#define X_SIZE_NODES 384
#define Y_SIZE_NODES 384
#define Z_SIZE_NODES 384

#define UP Y_SIZE_NODES * Z_SIZE_NODES
#define DOWN 0

#define arr(a,b,c) (a) * Y_SIZE_NODES * Z_SIZE_NODES + (b) * Z_SIZE_NODES + (c)


int wsize, wrank, xThrNodesSize;
double xStride, yStride, zStride;

double realPhi(double x, double y, double z){
    return x*x + y*y + z*z;
}

double realRo(double x, double y, double z){
    return 6 - a * realPhi(x, y, z);
}

void initBorders(double* nodePhi, int xStart){
    for(int i = 0, localStart = xStart; i < xThrNodesSize; ++i, ++localStart){
        for(int j = 0; j < Y_SIZE_NODES ; ++j){
            for(int k = 0; k < Z_SIZE_NODES; ++k){
                if(localStart*j*k == 0 || localStart == (X_SIZE_NODES - 1) || j == (Y_SIZE_NODES - 1) || k == (Z_SIZE_NODES - 1))
                    nodePhi[arr(i,j,k)] = realPhi(X_NULL + localStart*xStride, Y_NULL + j*yStride, Z_NULL + k*zStride);
            }
        }
    }
}

double getConstant(){
    return 1.0 / ((2.0 / (xStride * xStride)) + (2.0 / (yStride * yStride)) + (2.0 / (zStride * zStride)) + a);
}

void setLocalRo (double* nodeRo, int xstart){
    for(int i = 0; i < xThrNodesSize; ++i){
        for(int j = 0; j < Y_SIZE_NODES; ++j){
            for(int k = 0; k < Z_SIZE_NODES; ++k){
                nodeRo[arr(i,j,k)] = realRo(X_NULL + (xstart + i) * xStride, Y_NULL + j*yStride, Z_NULL + k*zStride);
            }
        }
    }
}

void sendRecvData(double* localNodePhi , MPI_Request* request, double* borders){
    if(wrank){
        MPI_Isend(&localNodePhi[0], Y_SIZE_NODES * Z_SIZE_NODES,MPI_DOUBLE, wrank - 1,UP, MPI_COMM_WORLD, &request[0]);
        MPI_Irecv(&borders[UP], Y_SIZE_NODES * Z_SIZE_NODES, MPI_DOUBLE, wrank - 1, DOWN,MPI_COMM_WORLD, &request[1]);
    }

    if(wrank != wsize - 1){
        MPI_Isend(&localNodePhi[xThrNodesSize - 1], Y_SIZE_NODES * Z_SIZE_NODES,MPI_DOUBLE, wrank + 1, DOWN, MPI_COMM_WORLD, &request[2]);
        MPI_Irecv(&borders[DOWN], Y_SIZE_NODES * Z_SIZE_NODES, MPI_DOUBLE, wrank + 1, UP, MPI_COMM_WORLD, &request[3]);
    }
}

void waitDataSendRecv(MPI_Request* request){
    if(wrank){
        MPI_Wait(&request[0],MPI_STATUS_IGNORE);
        MPI_Wait(&request[1], MPI_STATUS_IGNORE);
    }

    if(wrank != wsize - 1){
        MPI_Wait(&request[2],MPI_STATUS_IGNORE);
        MPI_Wait(&request[3], MPI_STATUS_IGNORE);
    }

}

void refreshBorders (double* borders, double* localNodePhi, double* nodeRo, double constantBeforePhi){
    int phiX, phiY, phiZ, i;

    for(int j = 1; j < Y_SIZE_NODES - 1; ++j){
        for(int k = 1; k < Z_SIZE_NODES - 1; ++k){
            if(wrank){
                phiX = (localNodePhi[arr(1,j,k)] + (&borders[UP])[j * Z_SIZE_NODES + k])/(xStride * xStride);
                phiY = (localNodePhi[arr(0,j+1,k)] + localNodePhi[arr(0 ,j- 1,k)])/(yStride * yStride);
                phiZ = (localNodePhi[arr(0,j,k+1)] + localNodePhi[arr(0 ,j,k - 1)])/(zStride * zStride);
                localNodePhi[arr(0,j,k)] = constantBeforePhi * (phiX + phiY + phiZ - nodeRo[arr(0,j,k)]);
            }
            if(wrank != wsize - 1){
                i = xThrNodesSize - 1;
                phiX = (localNodePhi[arr(i - 1,j,k )] + (&borders[DOWN])[j * Z_SIZE_NODES + k])/(xStride * xStride);
                phiY = (localNodePhi[arr(i,j+1,k )] + localNodePhi[arr(i ,j- 1,k)])/(yStride * yStride);
                phiZ = (localNodePhi[arr(i,j,k+1 )] + localNodePhi[arr(i ,j,k - 1)])/(zStride * zStride);
                localNodePhi[arr(i,j,k)] = constantBeforePhi * (phiX + phiY + phiZ - nodeRo[arr(i,j,k)]);
            }
        }
    }

}

double calculate(double* nodesLocal, double* oldNodePhi, double* nodeRo, double constant ){
    double localMax = 0, phiX, phiY, phiZ, diff;
    for(int i = 1; i < xThrNodesSize - 1; ++i){
        for(int j = 0; j < Y_SIZE_NODES - 1; ++j){
            for(int k = 0; k < Z_SIZE_NODES - 1; ++k){
                oldNodePhi[arr(i,j,k)] = nodesLocal[arr(i,j,k)];
                phiX = (nodesLocal[arr(i+1,j,k)] + nodesLocal[arr(i - 1,j,k)])/(xStride * xStride);
                phiY = (nodesLocal[arr(i,j+1,k)] + nodesLocal[arr(i ,j- 1,k)])/(yStride * yStride);
                phiZ = (nodesLocal[arr(i,j,k+1)] + nodesLocal[arr(i ,j,k - 1)])/(zStride * zStride);
                nodesLocal[arr(i,j,k)] = constant * (phiX + phiY + phiZ - nodeRo[arr(i,j,k)]);
                diff = fabs(nodesLocal[arr(i,j,k)] - oldNodePhi[arr(i,j,k)]);
                if(diff > localMax){
                    localMax = diff;
                }
            }
        }
    }
    return localMax;
}

double getRealDiff(double* nodePhi,int xThrNodesSize,int xstart){
    double localMax = 0, globalMax, diff, realphi;

    for(int i = 1; i < xThrNodesSize - 1; ++i){
        for(int j = 1; j < Y_SIZE_NODES - 1; ++j){
            for(int k = 1; k < Z_SIZE_NODES - 1; ++k){
                realphi = realPhi(xstart + i*xStride, Y_NULL + j*yStride, Z_NULL + k*zStride);
                diff = nodePhi[arr(i,j,k)] - realphi;
                if(diff > localMax){
                    localMax = diff;
                }
            }
        }
    }
    MPI_Allreduce(&localMax,&globalMax,1,MPI_DOUBLE, MPI_MAX,MPI_COMM_WORLD);
    return globalMax;
}

void Jakobi(){
    MPI_Comm_size(MPI_COMM_WORLD, &wsize);
    MPI_Comm_rank(MPI_COMM_WORLD, &wrank);

    // SHARED
    xStride = SQUARE_SIDE_LENGTH / (X_SIZE_NODES - 1);
    yStride = SQUARE_SIDE_LENGTH / (Y_SIZE_NODES - 1);
    zStride = SQUARE_SIDE_LENGTH / (Z_SIZE_NODES - 1);

    // line count per thr
    xThrNodesSize = X_SIZE_NODES / wsize;
    // start line for thr
    int xStart = wrank * xThrNodesSize;

    int thrNodesSize = xThrNodesSize * Y_SIZE_NODES * Z_SIZE_NODES ;
    double* nodesLocal = calloc(thrNodesSize ,sizeof(double));
    double* nodeRo = calloc(thrNodesSize, sizeof(double));
    double* oldNodePhi = calloc(thrNodesSize, sizeof(double));

    initBorders(nodesLocal, xStart);

    double constant = getConstant();
    double* borders = malloc(sizeof(double) * 2 * Y_SIZE_NODES * Z_SIZE_NODES);

    double start = MPI_Wtime();
    setLocalRo(nodeRo,xStart);

    int cycleFlag = 0;

    //double phiX, phiY, phiZ;
    MPI_Request request[4];
    double localMax = 0, globalMax;

    do{
        sendRecvData(nodesLocal,request,borders);

        localMax = calculate(nodesLocal,oldNodePhi, nodeRo,constant);

        waitDataSendRecv(request);

        refreshBorders(borders,nodesLocal, nodeRo, constant);

        MPI_Allreduce(&localMax, &globalMax, 1, MPI_DOUBLE, MPI_MAX, MPI_COMM_WORLD);
        cycleFlag = globalMax > EPSILON;

        localMax = 0;

    } while (cycleFlag);

    double end = MPI_Wtime();

    double realDiff = getRealDiff(nodesLocal, xThrNodesSize, (X_NULL + xStart*xStride));
    if(!wrank){
        printf("realDiff: %lf\n", realDiff);
        printf("Time taken: %lf\n", end - start);
    }

    free(nodesLocal);
    free(borders);
    free(oldNodePhi);
    free(nodeRo);
}

int main(int argc, char* argv[]){
    MPI_Init(&argc, &argv);
    Jakobi();
    MPI_Finalize();
    return 0;
}