#include <stdio.h>
#include <stdlib.h>
#include "mpi.h"

int main(int argc, char *argv[]) {
        int numtasks, rank, N;
        int result=0;

        //arrays for the two vectors
        int* a;
        int* b;

        double start, end, tot_time;
        start = 0;
        end = 0;

        //initializing MPI
        MPI_Init(&argc, &argv);
        //get size
        MPI_Comm_size(MPI_COMM_WORLD, &numtasks);
        //get rank
        MPI_Comm_rank(MPI_COMM_WORLD, &rank);
        //have to broadcast N still to distribute it
        MPI_Bcast(&N, 1, MPI_INT, 0, MPI_COMM_WORLD);

        if(rank == 0) {
                N = 100; //simulated reading in user input
        }

        if(rank == 0) {
                //pointers to dynamically allocate memory for 2 arrays on each node. done in master thread
                a=(int*)malloc(N*sizeof(int));
                b=(int*)malloc(N*sizeof(int));
                //hardcoding each value of vectors to 2
                for(int i = 0; i < N; i++) {
                        a[i]=b[i]=2;
                }
        }

        //receive buffers
//        int* recA = (int*)malloc(N*sizeof(int));
  //      int* recB = (int*)malloc(N*sizeof(int));

        int new_N;      //new_N: new discrete variable to sum counts
        if(rank < (N%numtasks)) {
                new_N = N/numtasks + 1;
        }
        else {
                new_N = N/numtasks;
        }

        //array to hold count for each rank
        int* sendcount = (int*)malloc(N*sizeof(int));
        //array that holds displacement from start of sendbuf to take the data
        int* displs = (int*)malloc(N*sizeof(int));

        //fill sendcount array
        for(int i = 0; i < new_N; i++) {
                sendcount[i] = new_N;
                displs[i] = i*sendcount[i];
        }
        //takes care of remainder
        for(int i = N%numtasks; i < numtasks; i++) {
                sendcount[i] = N/numtasks;
                displs[i] = i*sendcount[i];
        }

        start = MPI_Wtime();
        MPI_Scatterv(a, sendcount, displs, MPI_INT, a, new_N, MPI_INT, 0, MPI_COMM_WORLD);
        MPI_Scatterv(b, sendcount, displs, MPI_INT, a, new_N, MPI_INT, 0, MPI_COMM_WORLD);
        end = MPI_Wtime();
        tot_time = end - start;
        printf("Scatterv total time elapsed on %d: %lf", rank, tot_time);

        //actually calculating dot product
        for(int i = 0; i < N; i++){
                result+=(a[i] * b[i]);
        }
        //collect individual sums from processors for final result
        MPI_Allreduce(MPI_IN_PLACE, &result, 1, MPI_INT, MPI_SUM, MPI_COMM_WORLD);
        printf("Rank: %d: Result: %d\n", rank, result);
        //finished with MPI
        MPI_Finalize();
        return 0;
}