#include <iostream>
#include <fstream>
#include <vector>
#include <cmath>
#include <mpi.h>

using namespace std;

void readArraysFromFile(string filename, double* aMatrix, double* fMatrix, int size) {
    ifstream fin(filename);
    int m, n;
    fin >> m >> n;

    m = 0;
    n = 0;
    for (int i = 0; i < size * (size + 1); ++i) {
        if (i == size * (m + 1) + m) {
            fin >> fMatrix[m];
            m++;
        } else {
            fin >> aMatrix[n];
            n++;
        }
    }
}

void readArrayFromFile(string filename, double* xMatrix) {
    ifstream fin(filename);
    int m;
    fin >> m;

    for (int i = 0; i < m; ++i) {
        fin >> xMatrix[i];
    }
}

void createMatrixFile(string filename, int m, int n) {
    ofstream fout;
    fout.open(filename);
    fout << m;
    if (n != 1){
        fout << " " << n;
    }
    fout << endl;

    for (int i = 0; i < m; i++) {
        for (int j = 0; j < n; j++) {
            if (i == j) {
                fout << rand() % 9 + m * 10 << " ";
            } else {
                fout << rand() % 9 + 1 << " ";
            }
        }
        fout << endl;
    }
}

void writeMatrixInFile(string filename, double* matrix, int size) {
    ofstream fout;
    fout.open(filename);
    fout << size << endl;

    for (int i = 0; i < size; i++ ) {
        fout << matrix[i] << endl;
    }
}

double* Jacobi (int N, double* A, double* F, double* X, double eps, int M, int rank, int* receiveCount, int* displs)
{
    auto* TempX = new double[N];
    double norm = 0.0;
    double globalNorm = 0.0;

    do {
        for (int i = 0; i < N; i++) {
            TempX[i] = F[i];
            for (int g = 0; g < M; g++) {
                if (i + rank * N != g) {
                    TempX[i] -= A[i * M + g] * X[g];
                }
            }
            TempX[i] /= A[i * M + i + displs[rank]];
        }

        norm = fabs(X[displs[rank]] - TempX[0]);
        for (int h = 0; h < N; h++) {
            if (fabs(X[h + displs[rank]] - TempX[h]) > norm) {
                norm = fabs(X[h + displs[rank]] - TempX[h]);
            }
            X[h + displs[rank]] = TempX[h];
        }

        for (int i = 0; i < N; i++) {
            //cout << rank << " " << TempX[i] << endl;
        }

        MPI_Allgatherv(TempX, N, MPI_DOUBLE, X, receiveCount, displs, MPI_DOUBLE, MPI_COMM_WORLD);
        if (X[0] > -200 && X[0] < 200 && rank == 0) {
            for (int i = 0; i < M; i++) {
                //cout << rank << " " << X[i] << endl;
            }
            //cout << endl;
        }
        MPI_Reduce(&norm, &globalNorm, 1, MPI_DOUBLE, MPI_MAX, 0, MPI_COMM_WORLD);
        MPI_Bcast(&globalNorm, 1, MPI_DOUBLE, 0, MPI_COMM_WORLD);
        norm = globalNorm;

    } while (norm > eps);
    delete[] TempX;
    return X;
}

int main(int argc, char** argv) {
    int m = atoi(argv[1]);
    string aMatrixFile = argv[2];
    string matrixFile = argv[3];
    string answerFile = argv[4];
    double eps = atof(argv[5]);

    createMatrixFile(aMatrixFile, m, m + 1);
    createMatrixFile(matrixFile, m, 1);

    MPI_Init(&argc, &argv);
    auto* aMatrix = new double[m * m]; // матрица коэфициентов
    auto* fMatrix = new double[m]; // свободные члены
    auto* xMatrix = new double[m]; // начальное приближение

    double time = 0.0; // счетчик времени

    int rank;
    int worldSize;
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &worldSize);

    if (rank == 0) {
        readArraysFromFile(aMatrixFile, aMatrix, fMatrix, m);
        readArrayFromFile(matrixFile, xMatrix);
        time = MPI_Wtime();
    }
    MPI_Bcast(xMatrix, m, MPI_DOUBLE, 0, MPI_COMM_WORLD);

    auto* sendcountsA = new int[worldSize];
    auto* sendcountsF = new int[worldSize];
    auto* displsA = new int[worldSize];
    auto* displsF = new int[worldSize];
    int remainder = m % worldSize;
    int sum = 0;

    for (int i = 0; i < worldSize; i++) {
        sendcountsF[i] = m / worldSize;
        if (remainder > 0) {
            sendcountsF[i]++;
            remainder--;
        }
        sendcountsA[i] = sendcountsF[i] * m;
        displsF[i] = sum;
        displsA[i] = sum * m;
        sum += sendcountsF[i];
    }

    auto* subAMatrix = new double[sendcountsA[rank]]; // подматрица коэфициентов
    auto* subFMatrix = new double[sendcountsF[rank]]; // подматрица свободных членов

    MPI_Scatterv(aMatrix, sendcountsA, displsA, MPI_DOUBLE, subAMatrix, sendcountsA[rank], MPI_DOUBLE, 0, MPI_COMM_WORLD);
    MPI_Scatterv(fMatrix, sendcountsF, displsF, MPI_DOUBLE, subFMatrix, sendcountsF[rank], MPI_DOUBLE, 0, MPI_COMM_WORLD);

    for (int i = 0; i < sendcountsA[rank]; i++) {
        cout << rank << " " << subAMatrix[i] << endl;
    }
    for (int i = 0; i < sendcountsF[rank]; i++) {
        //cout << rank << " " << subFMatrix[i] << endl;
    }

    xMatrix = Jacobi(sendcountsF[rank], subAMatrix, subFMatrix, xMatrix, eps, m, rank, sendcountsF, displsF);

    if (rank == 0) {
        time = MPI_Wtime() - time;
        cout << time << endl;
        writeMatrixInFile(answerFile, xMatrix, m);
    }

    MPI_Finalize();
    return 0;
}