#define _GNU_SOURCE
#include "mpi.h"
#include <unistd.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <math.h>



#define MATRIX_SIZE 10
#define RANDOM_MAX 10
#define RANDOM_MIN 0


float myRand() {
    return (float)(RANDOM_MIN + (rand() % (RANDOM_MAX - RANDOM_MIN + 1)));
}


int main(int argc, char** argv) {

    setbuf(stdout, NULL);
    int err;
    err = MPI_Init(&argc, &argv);

    int rank;
    int size;

    err = MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    err = MPI_Comm_size(MPI_COMM_WORLD, &size);

    //srand(time(NULL));

    ///   Buffer init
    int layerSize;
    MPI_Pack_size(MATRIX_SIZE, MPI_FLOAT, MPI_COMM_WORLD, &layerSize);
    int bufferSize = 10 * (layerSize + MPI_BSEND_OVERHEAD);
    char* buffer = (char*)malloc(bufferSize * sizeof(float));
    MPI_Buffer_attach(buffer, bufferSize);


    int colCnt = (MATRIX_SIZE + 1) / size;
    int rest = (MATRIX_SIZE + 1) % size;
    if(rank < rest) {
        colCnt++;
    }

    ///   Process that takes b
    int procB = (MATRIX_SIZE + 1) % size - 1;
    procB = procB < 0 ? size - 1 : procB;

    float columns[MATRIX_SIZE + 1][colCnt];
//    float columns[colCnt][MATRIX_SIZE];


    ///   Number of columns in each proc
    int* counts = NULL;


    if(!rank) {
        counts = (int*)malloc(sizeof(int) * size);

        for(int i = 0; i < size; ++i) {
            counts[i] = 0;
        }

        ///   Generate random matrix
        float A[MATRIX_SIZE][MATRIX_SIZE + 1];
        //float b[MATRIX_SIZE];
        for(int i = 0; i < MATRIX_SIZE; ++i) {
            for(int j = 0; j < MATRIX_SIZE + 1; ++j) {
                A[i][j] = myRand();
            }
            //b[i] = myRand();
        }

        ///   Print matrix
        for(int i = 0; i < MATRIX_SIZE; ++i) {
            for(int j = 0; j < MATRIX_SIZE; ++j) {
                printf("%.2f ", A[i][j]);
            }
            printf(" |  %.2f\n", A[i][MATRIX_SIZE]);
        }


        ///   Send cols
        for(int i = 0; i < MATRIX_SIZE + 1; ++i) {
            if(i % size) {
                float buffer[MATRIX_SIZE + 1];

                for(int j = 0; j < MATRIX_SIZE; ++j) {
                    buffer[j] = A[j][i];
                }
                buffer[MATRIX_SIZE] = i;

                MPI_Bsend(&(buffer[0]), MATRIX_SIZE + 1, MPI_FLOAT, i % size,
                          0, MPI_COMM_WORLD);
                counts[i % size]++;
            }
            else {
                for(int j = 0; j < MATRIX_SIZE; ++j) {
                    columns[j][i / size] = A[j][i];
                }
                columns[MATRIX_SIZE][i / size] = i;
            }
        }
    }
    else {
        /// Receive matrix
        for(int i = 0; i < colCnt; ++i) {
            float buffer[MATRIX_SIZE + 1];
            MPI_Recv(&(buffer[0]), MATRIX_SIZE + 1, MPI_FLOAT, 0,
                     0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
            int j = buffer[MATRIX_SIZE] / size;

            for(int i = 0; i < MATRIX_SIZE + 1; ++i) {
                columns[i][j] = buffer[i];
            }
        }
    }
    int cnt = 0;
    MPI_Barrier(MPI_COMM_WORLD);



    double startTime = MPI_Wtime();


    ///   Forward pass
    cnt = 0;
    for(int i = 0; i < MATRIX_SIZE - 1; ++i) {
        MPI_Barrier(MPI_COMM_WORLD);
        if(i % size == rank) {
            ///   Max finding
            double max = 0;
            int maxInd = 0;
            for(int j = i; j < MATRIX_SIZE; ++j) {
                if(max < abs(columns[j][cnt])) {
                    max = abs(columns[j][cnt]);
                    maxInd = j;
                }
            }

            int indices[2];
            indices[0] = i;
            indices[1] = maxInd;

            ///   Send swap to other processes
            MPI_Bcast(&(indices[0]), 2, MPI_INT, i % size, MPI_COMM_WORLD);

            ///   Perform swaps
            for(int l = 0; l < colCnt; ++l) {
                float swap = columns[indices[0]][l];
                columns[indices[0]][l] = columns[indices[1]][l];
                columns[indices[1]][l] = swap;
            }


            float coeffs[MATRIX_SIZE - i - 1];
            ///   Subtracting
            for(int j = i + 1; j < MATRIX_SIZE; ++j) {
                coeffs[j - i - 1] = -columns[j][cnt] / columns[i][cnt];
                columns[j][cnt] = 0;
            }


            ///   Send coefficients
            for(int j = 0; j < size; ++j) {
                if(rank != j) {
                    MPI_Bsend(&(coeffs[0]), MATRIX_SIZE - i - 1, MPI_FLOAT, j,
                              0, MPI_COMM_WORLD);
                }
            }
            for(int j = cnt + 1; j < colCnt; ++j) {
                for(int k = i + 1; k < MATRIX_SIZE; ++k) {
                    columns[k][j] += coeffs[k - i - 1] * columns[i][j];
                }
            }

            cnt++;
        }
        else {
            int swaps[2];
            ///   Get swap
            MPI_Bcast(&(swaps[0]), 2, MPI_INT, i % size, MPI_COMM_WORLD);

            ///   Perform swap
            for(int j = 0; j < colCnt; ++j) {
                float tmp = columns[swaps[0]][j];
                columns[swaps[0]][j] = columns[swaps[1]][j];
                columns[swaps[1]][j] = tmp;
            }

            ///   Get coefficients
            float coeffs[MATRIX_SIZE - i - 1];
            MPI_Recv(&(coeffs[0]), MATRIX_SIZE - i - 1, MPI_FLOAT, i % size,
                     0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);


            ///   Perform subtracting
            for(int j = cnt; j < colCnt; ++j) {
                for(int k = i + 1; k < MATRIX_SIZE; ++k) {
                    columns[k][j] += coeffs[k - i - 1] * columns[i][j];
                }
            }
        }
    }

    cnt = colCnt - 1;
    if(rank == MATRIX_SIZE % size) {
        cnt--;
    }

    ///   Backward pass
    for(int i = MATRIX_SIZE - 1; i >= 0; --i) {
        MPI_Barrier(MPI_COMM_WORLD);
        if(i % size == rank && rank != MATRIX_SIZE % size) {
            float coeffs[i + 1];
            ///   Perform subtracting
            coeffs[i] = columns[i][cnt];
            columns[i][cnt] = 1.;
            for(int j = i - 1; j >= 0; --j) {
                coeffs[j] = -columns[j][cnt];
                columns[j][cnt] = 0;
            }

            ///   Send coefficients
            MPI_Bsend(&(coeffs[0]), i + 1, MPI_FLOAT, MATRIX_SIZE % size,
                      0, MPI_COMM_WORLD);

            cnt--;
        }
        else if(rank == MATRIX_SIZE % size) {
            if(i % size != rank && size != 1) {
                ///   Get coefficients
                float coeffs[i + 1];

                MPI_Recv(&(coeffs[0]), i + 1, MPI_FLOAT, i % size,
                         0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);

                columns[i][colCnt - 1] /= coeffs[i];
                for(int j = i - 1; j >= 0; --j) {
                    columns[j][colCnt - 1] += coeffs[j] * columns[i][colCnt - 1];
                }
            }
            else {
                columns[i][colCnt - 1] /= columns[i][cnt];
                for(int j = i - 1; j >= 0; --j) {
                    columns[j][colCnt - 1] -= columns[i][colCnt - 1] * columns[j][cnt];
                    columns[j][cnt] = 0;
                }
                cnt--;
            }
        }
    }

    MPI_Barrier(MPI_COMM_WORLD);
    if(rank == MATRIX_SIZE % size) {
        for(int i = 0; i < MATRIX_SIZE; ++i) {
            printf("x%d: %.6f\n", i + 1, columns[i][colCnt - 1]);
        }
        fflush(stdout);
    }

    free(counts);
    MPI_Barrier(MPI_COMM_WORLD);


    double endTime = MPI_Wtime();
    if(rank == 0) {
        usleep(500*1000);
        printf("Time: %.6f\n", endTime - startTime);
    }


    err = MPI_Finalize();
    return 0;
}