#include <iostream>
#include "mpi.h"
#include <random>
#include <cassert>

using namespace std;

#define VECTOR_SIZE 1000

inline int generateRandomInRange(int a, int b) {
    random_device rd;
    uniform_int_distribution gen(a, b);
    return gen(rd);
}

void fill_vector_values(int *vec){
    for(size_t i = 0; i < VECTOR_SIZE; i++){
        vec[i] = generateRandomInRange(1, 100);
    }
}


void print_vec(int *vec, int n){
    for(size_t i = 0; i < n; i++){
        std::cout << vec[i] << endl;
    }
    std::cout << "---------\n";
}


int main(int argc, char **argv){
    int rank, size;
    int VEC_SIZE = VECTOR_SIZE;
    auto *vec1 = new int[VECTOR_SIZE];
    auto *vec2 = new int[VECTOR_SIZE];
    fill_vector_values(vec1);
    fill_vector_values(vec2);

    MPI_Init(&argc, &argv);
    MPI_Comm_size(MPI_COMM_WORLD, &size);
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    size_t sub_size = VEC_SIZE / size;
    size_t rest = VEC_SIZE % size;
    int result = 0;

    if(rank == 0){
        for(size_t i = 0; i < size; i++){
            //std::cout << sub_size*i + 1 << " " << sub_size*i + sub_size  << endl;
            MPI_Send(&vec1[sub_size*i + 1 ], static_cast<int>(sub_size), MPI_INT, static_cast<int>(i), (int)i, MPI_COMM_WORLD);
            MPI_Send(&vec2[sub_size*i + 1 ], static_cast<int>(sub_size), MPI_INT, static_cast<int>(i), (int)i, MPI_COMM_WORLD);
        }


        for(size_t i = 0; i < size; i++){
            int local_res;
            MPI_Recv(&local_res, 1, MPI_INT, MPI_ANY_SOURCE, MPI_ANY_TAG, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
            result += local_res;
        }

        std::cout << result;
    }

    if(rank != 0){
        MPI_Recv(&vec1[0], (int)sub_size, MPI_INT, 0, rank, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
        MPI_Recv(&vec2[0], (int)sub_size, MPI_INT, 0, rank, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
        int curr_res = 0;
        for(size_t j = 0; j < sub_size; j++){
            curr_res += vec1[j]*vec2[j];
        }
        MPI_Send(&curr_res, 1, MPI_INT, 0, rank, MPI_COMM_WORLD);
    }


    MPI_Finalize();
    delete []vec1;
    delete []vec2;
    return 0;
}