#include <mpi.h>

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

#undef min

/* Within the given communicator, assuming that rank 0 is the root of the tree and, at each level, there are `nrBranches` children (assuming there are enough nodes),
 * this function computes, for the current node:
 *   - who is the parent (-1 if we are the root)
 *   - how many nodes the current process has in its sub-tree (including itself)
 * */
void parentAndTotalSubordinates(MPI_Comm comm, int nrBranches, int& parent, int& nrProcs) {
    int me;
    int nrTotalProcs;
    MPI_Comm_rank(comm, &me);
    MPI_Comm_size(comm, &nrTotalProcs);
    if(me == 0) {
        parent = -1;
        nrProcs = nrTotalProcs;
        return;
    }
    parent = 0;
    while(true) {
        int nrChildren = std::min(nrBranches, nrTotalProcs);
        for(int i=0 ; i<nrChildren ; ++i) {
            int begin = i*nrTotalProcs/nrChildren;
            int end = (i+1)*nrTotalProcs/nrChildren;
            if(me >= parent+begin && me<parent+end) {
                if(me == parent + begin) {
                    nrProcs = end-begin;
                    return;
                }
                parent = parent + begin;
                nrTotalProcs = end - begin;
                break; // break from inner for loop
            }
        }
    }
}


enum MpiTags {
    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 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) {
    int parent;
    int nrProcs;
    parentAndTotalSubordinates(comm, nrBranches, parent, nrProcs);
    MPI_Status status;
    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 3

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();
}