#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) {
        std::cout<< "Size: ";
        std::cin >> size;
    }
    MPI_Bcast(&size, 1, MPI_INT, 0, MPI_COMM_WORLD);
    double *A = new double[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()%10;
            }
            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()%10;
            X[i] = F[i];
        }
        cout << endl;
        for(int i = 0; i < size; i++) {
            cout << F[i] << " ";
        }
        cout << endl;
    }
    int localSize = (size * size)/world_size;
    double *localA = new double[localSize];
    MPI_Scatter(A, localSize, MPI_DOUBLE, localA, localSize, MPI_DOUBLE, 0, MPI_COMM_WORLD);
    MPI_Bcast(F, size, MPI_DOUBLE, 0, MPI_COMM_WORLD);
    double *TempX = new double[size];
    double norm;
    double *allNorms = new double[world_size];
    double *allX = new double[size];
    int count = 0;
    do {
        count++;
        for (int i = 0; i < localSize/size; i++) {
            TempX[i] = F[i];
            for (int j = 0; j < size; j++) {
                if (i != j) {
                    TempX[i] -= localA[i*size+j] * X[j];
                }
            }
            TempX[i] /= localA[i*size+i];
        }
        MPI_Gather(TempX, localSize/size, MPI_DOUBLE, TempX, localSize/size, MPI_DOUBLE, 0, MPI_COMM_WORLD);
        MPI_Bcast(TempX, size, MPI_DOUBLE, 0, MPI_COMM_WORLD);
        norm = fabs(X[0] - TempX[0]);
        for (int i = 0; i < localSize/size; i++) {
            if (fabs(X[i] - TempX[i]) > norm) {
                norm = fabs(X[i] - TempX[i]);
            }
            X[i] = TempX[i];
        }
        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);
    MPI_Gather(X, localSize/size, MPI_DOUBLE, allX, localSize/size, MPI_DOUBLE, 0, MPI_COMM_WORLD);
    if(world_rank == 0) {
        cout << "iter: " << count << endl;
        for (int i = 0; i < size; i++) {
            cout << "X[" << i << "] = " << allX[i] << endl;
        }
    }
    MPI_Finalize();
}