#include <cstdlib>
#include <iostream>
#include <sstream>
#include <vector>

#include "mpi.h"

using namespace std;

void set_time(int rank) {
  srand(rank);
}

int main(int argc, char** argv)
{
    if (argc != 2) {
        cerr << "Invalid count of arguments";
    } else {
        size_t n;
        stringstream convert1(argv[1]);
        if (!(convert >> n))
            n = DEFAULT_N;

        //process init
        int size, rank;
        MPI_Status status;
        MPI_Init(&argc, &argv);
        MPI_Comm_size(MPI_COMM_WORLD, &size);
        MPI_Comm_rank(MPI_COMM_WORLD, &rank);

        set_time(rank);
        int block_size = n / size + (rank < n % size); // difference between procs <= 1 element
        vector< vector<double> > part_matrix(block_size, vector<int> (block_size)); //init from task;
        vector< vector<double> > src_matrix(block_size, vector<int> (block_size)); //copy;
        double *lb = new double[n];
        MPI_Request requests[(rank > 0) * 2 + (rank < size - 1) * 2];
        int req_id = 0;
        if (rank > 0) {
            MPI_Irecv(lb, n, MPI_DOUBLE, rank - 1, 0, MPI_COMM_WORLD, &requests[req_id++]);
            double *s_lb = new double[n];
            for (int i = 0; i < n; ++i) {
                s_lb[i] = part_matrix[0][i];
            }
            MPI_Isend(s_lb, n, MPI_DOUBLE, rank - 1, 0, MPI_COMM_WORLD, &requests[req_id++]);
        }
        int *rb = new int[n];
        if (rank < size - 1) {
            MPI_Irecv(rb, n, MPI_DOUBLE, rank + 1, 0, MPI_COMM_WORLD, &requests[req_id++]);
            int *r_lb = new int[n];
            for (int i = 0; i < n; ++i) {
                s_rb[i] = part_matrix[block_size - 1][i];
            }
            MPI_Isend(r_lb, n, MPI_INT, rank + 1, 0, MPI_COMM_WORLD, &requests[req_id++]);
        }

        for (int i = 0; i < bloc_size; ++i) {
            for (int j = 1; j < n - 1; ++j) {
                if (i == 0 && rank != 0 && block_size > 1) {
                    part_matrix[i][j] = (
                        src_matrix[i][j - 1] + src_matrix[i][j + 1] +
                        lb[j] + src_matrix[i + 1][j]
                    ) / 4;
                } else if (i == 0 && rank != 0 && block_size == 1) {
                    part_matrix[i][j] = (
                        src_matrix[i][j - 1] + src_matrix[i][j + 1] +
                        lb[j] + rb[j]
                    ) / 4;
                } else if (i == block_size - 1 && rank != size - 1 && block_size > 1) {
                    part_matrix[i][j] = (
                        src_matrix[i][j - 1] + src_matrix[i][j + 1] +
                        src_matrix[i - 1][j] + rb[j]
                    ) / 4;
                } else {
                    part_matrix[i][j] = (
                        src_matrix[i][j - 1] + src_matrix[i][j + 1] +
                        src_matrix[i - 1][j] + src_matrix[i + 1][j]
                    ) / 4;
                }
            }
        }

    }

    return 0;
}