#include "mpi.h"
#include <iostream>
#include <ctime>
using namespace std;

#define MAX_VALUE 10
#define VECTOR_SIZE 10
const int MAIN_PROCESS = 0;

class MPI_Process
{
private:
    int rank;
    int processCount;

public:
    MPI_Process(int argc, char *argv[])
    {
        MPI_Init(&argc, &argv);
        MPI_Comm_rank(MPI_COMM_WORLD, &rank);
        MPI_Comm_size(MPI_COMM_WORLD, &processCount);
    }

    ~MPI_Process()
    {
        MPI_Finalize();
    }

    int getRank()
    {
        return rank;
    }

    int getProcessCount()
    {
        return processCount;
    }

    bool isMain()
    {
        return (rank == MAIN_PROCESS);
    }
};

void createVector(double *, int, int);
void printVector(double *, int);

int main(int argc, char *argv[])
{
    srand((unsigned)time(0));

    MPI_Process process(argc, argv);

    int     sendSize = VECTOR_SIZE / process.getProcessCount();
    int     leftSize = VECTOR_SIZE % process.getProcessCount();
    double  *vector1 = nullptr;
    double  *vector2 = nullptr;
    int     recvSize = process.getRank() >= leftSize ? sendSize : sendSize + 1;
    double  *buffer1 = new double[recvSize];
    double  *buffer2 = new double[recvSize];
    int     *sendCounts = nullptr;
    int     *displs = nullptr;

    if (process.isMain())
    {
        vector1 = new double[VECTOR_SIZE];
        vector2 = new double[VECTOR_SIZE];

        createVector(vector1, VECTOR_SIZE, MAX_VALUE);
        createVector(vector2, VECTOR_SIZE, MAX_VALUE);

        cout << "Vector 1 = ";
        printVector(vector1, VECTOR_SIZE);
        cout << "Vector 2 = ";
        printVector(vector2, VECTOR_SIZE);

        sendCounts = new int[process.getProcessCount()];
        displs = new int[process.getProcessCount()];

        for (int i = 0; i < process.getProcessCount(); i++)
        {
            sendCounts[i] = i >= leftSize ? sendSize : sendSize + 1;
            if (i == 0) {
                displs[i] = 0;
            }
            else {
                displs[i] = displs[i - 1] + sendCounts[i - 1];
            }
        }
    }

    MPI_Scatterv(vector1, sendCounts, displs, MPI_DOUBLE, buffer1, recvSize, MPI_DOUBLE, MAIN_PROCESS, MPI_COMM_WORLD);
    MPI_Scatterv(vector2, sendCounts, displs, MPI_DOUBLE, buffer2, recvSize, MPI_DOUBLE, MAIN_PROCESS, MPI_COMM_WORLD);
    for (int i = 0; i < recvSize; i++) {
        buffer1[i] -= buffer2[i];
    }
    delete[] buffer2;

    MPI_Gatherv(buffer1, recvSize, MPI_DOUBLE, vector1, sendCounts, displs, MPI_DOUBLE, MAIN_PROCESS, MPI_COMM_WORLD);
    delete[] buffer1;
    delete[] sendCounts;
    delete[] displs;
    if (process.isMain())
    {
        cout << "Vector1 - Vector2 = ";
        printVector(vector1, VECTOR_SIZE);
        delete[] vector2;
        delete[] vector1;
    }
    return 0;
}

void createVector(double *out, int size, int maxValue)
{
    for (int i = 0; i < size; i++) {
        out[i] = (double)(rand() % (200 * maxValue + 1) - maxValue * 100) / 100;
    }
}

void printVector(double *vec, int size)
{
    cout << "(";
    for (int i = 0; i < size - 1; i++) {
        cout << vec[i] << ", ";
    }
    cout << vec[size - 1] << ")\n";
}