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

int main(int argc, char** argv) {
    MPI_Init(&argc, &argv);
    int world_size;
    MPI_Comm_size(MPI_COMM_WORLD, &world_size);
    int world_rank;
    MPI_Comm_rank(MPI_COMM_WORLD, &world_rank);

    int n;
    int *a, *b;

    if (world_rank == 0) {
        bool f = false;
        while (!f) {
            //printf("n = ");
            scanf("%i", &n);
            //printf("\n");
            if(n % world_size == 0) {
                f = true;
            }
        }

        srand(time(NULL));

        a = new int[n*n];
        b = new int[n];

        printf("a =\n");

        // generating matrix a

        bool converge = false;
        while (!converge) {
            for (int i = 0; i < n; i++) {
                for (int j = 0; j < n; j++) {
                    if (i == j) {
                        a[i*n + j] = 3*n + rand() % 10;
                    }
                    else {
                        a[i*n + j] = rand() % 4;
                    }
                }
            }

            converge = true;
            for (int i = 0; i < n; i++) {
                int sum = 0;
                for (int j = 0; j < n; j++) {
                    if (i != j) {
                        sum += abs(a[i*n + j]);
                    }
                }
                if (abs(a[i*n + i]) < sum) {
                    converge = false;
                    break;
                }
            }
        }

        for (int i = 0; i < n; i++) {
            for (int j = 0; j < n; j++) {
                printf("%2i ", a[i*n + j]);
            }
            printf("\n");
        }
        printf("\n");

        // vector b

        printf("b =\n");
        for (int i = 0; i < n; i++) {
            b[i] = rand() % 10;
            printf("%i ", b[i]);
        }
        printf("\n");
    }

    MPI_Bcast(&n, 1, MPI_INT, 0, MPI_COMM_WORLD);

    int *ra = new int[n*n / world_size];
    int *rb = new int[n / world_size];

    MPI_Scatter(a, n * n / world_size, MPI_INT, ra, n * n / world_size, MPI_INT, 0, MPI_COMM_WORLD);
    MPI_Scatter(b, n / world_size, MPI_INT, rb, n / world_size, MPI_INT, 0, MPI_COMM_WORLD);

    double *x0 = new double[n];
    double *x1 = new double[n];
    double *rx1 = new double[n / world_size];

    for (int i = 0; i < n; i++) {
        x0[i] = 0;
    }

    int end = 1;
    int endend = 1;
    double eps = 0.001;

    do {
        // calculating next iteration
        for (int i = 0; i < n / world_size; i++) {

            int ind = world_rank * n / world_size + i;

            double sum = 0;
            for (int j = 0; j < n; j++) {
                if (ind != j) {
                    sum += ra[i*n + j] * x0[j];
                }
            }

            rx1[i] = 1.0 / ra[i*n + ind] * (rb[i] - sum);
        }

        // checking end condition
        end = 1;
        int tmp = world_rank * n / world_size;
        double maxdiff = fabs(rx1[0] - x0[tmp]);

        x0[tmp] = rx1[0];

         for (int i = 1; i < n / world_size; i++) {
            if (fabs(rx1[i] - x0[tmp + i]) > maxdiff) {
                maxdiff = fabs(rx1[i] - x0[tmp + i]);
            }
            x0[tmp + i] = rx1[i];
        }
        if (maxdiff > eps) {
            end = 0;
        }

        int *res = new int[world_size];

        MPI_Gather(&end, 1, MPI_INT, res, 1, MPI_INT, 0, MPI_COMM_WORLD);
        MPI_Gather(rx1, n / world_size, MPI_DOUBLE, x0, n / world_size, MPI_DOUBLE, 0, MPI_COMM_WORLD);

        if (world_rank == 0) {
            endend = 1;
            for (int i = 0; i < world_size; i++) {
                if (res[i] == 0) {
                    endend = 0;
                    break;
                }
            }

            /*
            end = 1;
            double maxdiff = fabs(x1[0] - x0[0]);
            for (int i = 1; i < n; i++) {
                if (fabs(x1[i] - x0[i]) > maxdiff) {
                    maxdiff = fabs(x1[i] - x0[i]);
                }
            }
            if (maxdiff > eps) {
                end = 0;
            }
            */
            //printf("x1[0] = %lf\n", x1[0]);
            //printf("maxdiff = %lf\n", maxdiff);
            //printf("\nx =\n");
        }

        MPI_Bcast(&endend, 1, MPI_INT, 0, MPI_COMM_WORLD);
        MPI_Bcast(x0, n, MPI_DOUBLE, 0, MPI_COMM_WORLD);
    }
    while (!endend);

    if (world_rank == 0) {
        printf("\nx =\n");
        for (int i = 0; i < n; i++) {
            printf("%lf ", x0[i]);
        }
        printf("\n");
    }

    MPI_Finalize();
}