#include <mpi.h>
#include <stdio.h>
#include <iostream>
#include <cmath>

using namespace std;

const double eps = 0.001;

int main(int argc, char** argv) {
    srand(time(0));
    int size;
    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);
    if(world_rank == 0) {
        cout<< "Size: ";
        cin >> size;
    }
    MPI_Bcast(&size, 1, MPI_INT, 0, MPI_COMM_WORLD);
    double *A = new double[size*size];
    double *F = new double[size];
    double *X = new double[size]{};
    if(world_rank == 0) {
        for (int i = 0; i < size; i++) {
            for (int j = 0; j < size; j++) {
                A[i*size + j] = rand()%5 + 1;
            }
            for (int j = 1; j < size; j++) {
                A[i*size + i] += A[i*size + j];
            }
            for (int j = 0; j < size; j++) {
                cout << A[i*size + j] << " ";
            }
            cout << endl;
            F[i] = rand()%100 + 1;
        }
        cout << endl;
        cout << "F:";
        for(int i = 0; i < size; i++) {
            cout << F[i] << " ";
        }
        cout << endl;
    }
    int localSize = (size * size)/world_size;
    int localRows = localSize/size;
    double *localF = new double[localRows];
    double *localA = new double[localSize];
    double *localX = new double[localRows];
    MPI_Scatter(A, localSize, MPI_DOUBLE, localA, localSize, MPI_DOUBLE, 0, MPI_COMM_WORLD);
    MPI_Scatter(F, localRows, MPI_DOUBLE, localF, localRows, MPI_DOUBLE, 0, MPI_COMM_WORLD);
    MPI_Scatter(X, localRows, MPI_DOUBLE, localX, localRows, MPI_DOUBLE, 0, MPI_COMM_WORLD);
    MPI_Bcast(X, size, MPI_DOUBLE, 0, MPI_COMM_WORLD);
    double *TempX = new double[localRows];
    double norm;
    double *allNorms = new double[world_size];
    int count = 0;
    do {
        count++;
        for (int i = 0; i < localRows; i++) {
            TempX[i] = localF[i];
            for (int j = 0; j < size; j++) {
                if (world_rank * localRows + i != j) {
                    TempX[i] -= localA[i * size + j] * X[j];
                }
            }
            TempX[i] /= localA[i * size + world_rank * localRows + i];
        }
        norm = fabs(localX[0] - TempX[0]);
        for (int i = 0; i < localRows; i++) {
            if (fabs(localX[i] - TempX[i]) > norm) {
                norm = fabs(localX[i] - TempX[i]);
            }
            localX[i] = TempX[i];
        }
        MPI_Gather(localX, localRows, MPI_DOUBLE, X, localRows, MPI_DOUBLE, 0, MPI_COMM_WORLD);
        MPI_Bcast(X, size, MPI_DOUBLE, 0, MPI_COMM_WORLD);
        MPI_Gather(&norm, 1, MPI_DOUBLE, allNorms, 1, MPI_DOUBLE, 0, MPI_COMM_WORLD);
        if(world_rank == 0) {
            for (int i = 0; i < world_size; i++) {
                if(allNorms[i] > norm) {
                    norm = allNorms[i];
                }
            }
        }
        MPI_Bcast(&norm, 1, MPI_DOUBLE, 0, MPI_COMM_WORLD);
    } while (norm > eps);
    if(world_rank == 0) {
        cout << "iter: " << count << endl;
        for (int i = 0; i < size; i++) {
            cout << "X[" << i << "] = " << X[i] << endl;
        }
    }
    MPI_Finalize();
    return 0;
}