#include <mpi.h>

#include <algorithm>
#include <iostream>
#include <vector>

#undef min

enum MpiTags {
    mpiTagNrProcs = 1,
    mpiTagSizeToProcess = 2,
    mpiTagInputData = 3,
    mpiTagResult = 4
};

/** Processes as the root of the recursive decomposition tree.
 *
 * Parameters:
 *   - comm - communicator used
 *   - nrProcs - number of processes used for processing. The current process shall use processes with ranks from the current one to current rank + nrProcs -1.
 *   - nrBranches - number of brances at each level of the tree
 *   - pInput - pointer to the input data
 *   - inputSize - the size of the input data
 * Returns: the result (the sum of values in the input array)
 * */
int sumVectorN(MPI_Comm comm, int nrProcs, int nrBranches, int const* pInput, int inputSize) {
    if(nrProcs == 1) {
        int s = 0;
        for(int i=0 ; i<inputSize ; ++i) {
            s += pInput[i];
        }
        return s;
    }

    int me;
    MPI_Comm_rank(comm, &me);

    int nrChildren = std::min(nrBranches, nrProcs);
    for(int i=1 ; i<nrChildren ; ++i) {
        int childOffset = i*nrProcs/nrChildren;
        int child = me + childOffset;
        int nrDescendentsOfChild = (i+1)*nrProcs/nrChildren - childOffset;
        MPI_Ssend(&nrDescendentsOfChild, 1, MPI_INT, child, mpiTagNrProcs, comm);
        int childStartIndex = i*inputSize/nrChildren;
        int sizeForChild = (i+1)*inputSize/nrChildren - childStartIndex;
        MPI_Ssend(&sizeForChild, 1, MPI_INT, child, mpiTagSizeToProcess, comm);
        MPI_Ssend(pInput + childStartIndex, sizeForChild, MPI_INT, child, mpiTagInputData, comm);
    }

    int childrenNextLevel = nrProcs / nrChildren;
    int sizeLocal = inputSize / nrChildren;
    int sum = sumVectorN(comm, childrenNextLevel, nrBranches, pInput, sizeLocal);
    
    for(int i=1 ; i<nrChildren ; ++i) {
        int child = me + i*nrProcs/nrChildren;
        int childSum;
        MPI_Status status;
        MPI_Recv(&childSum, 1, MPI_INT, child, mpiTagResult, comm, &status);
        sum += childSum;
    }
    return sum;
}

void worker(MPI_Comm comm, int nrBranches) {
    MPI_Status status;
    int nrProcs;
    MPI_Recv(&nrProcs, 1, MPI_INT, MPI_ANY_SOURCE, mpiTagNrProcs, comm, &status);
    int parent = status.MPI_SOURCE;
    int sizeToRecv;
    MPI_Recv(&sizeToRecv, 1, MPI_INT, parent, mpiTagSizeToProcess, comm, &status);
    std::vector<int> data(sizeToRecv);
    MPI_Recv(data.data(), sizeToRecv, MPI_INT, parent, mpiTagInputData, comm, &status);

    int sum = sumVectorN(comm, nrProcs, nrBranches, data.data(), sizeToRecv);

    MPI_Ssend(&sum, 1, MPI_INT, parent, mpiTagResult, comm);
}

#define NR_BRANCHES 5

int main() {
    MPI_Init(0, 0);
    int me;
    MPI_Comm_rank(MPI_COMM_WORLD, &me);
    if(me == 0) {
        int nrElements = 1000;
        std::vector<int> input(nrElements);
        for(int i=0 ; i<nrElements ; ++i) {
            input[i] = i;
        }
        int nrProcs;
        MPI_Comm_size(MPI_COMM_WORLD, &nrProcs);
        int v = sumVectorN(MPI_COMM_WORLD, nrProcs, NR_BRANCHES, input.data(), nrElements);
        bool ok = (v == nrElements*(nrElements-1)/2);
        std::cout << "Result=" << v << (ok ? " Ok\n" : " ERROR!!!\n");
    } else {
        worker(MPI_COMM_WORLD, NR_BRANCHES);
    }
    MPI_Finalize();
}