#include <mpi.h>

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

#undef min

/** Processes the current node, as the member of the given communicator.
 * If the rank is 0, this process is the leader (the root of the sub-tree): in that case, the 'input' argument contains the data to be processed
 * and the function returns the result (the sum of elements in the input vector, for this problem).
 * Otherwise, this function acts as a worker
 *
 * Implementation:
 *   - if the rank is 0, this function splits the input, sends it to children, processes its own part (by doing a recursive call),
 *     receives the part results from children, combines them and, finally, returns the result.
 *   - if the rank in 1, this acts as a worker, receives the data from the parent, processes its own part (by a recursive call), and sends back the result
 *   - otherwise, this does nothing at this level, but still calls itself recursively for the next level
 *   - in all cases, this splits the communicator
 * */
int sumVector2(MPI_Comm comm, int const* pInput, int inputSize) {
    int me;
    int nrProcs;
    MPI_Comm_size(comm, &nrProcs);
    MPI_Comm_rank(comm, &me);
    if(nrProcs == 1) {
        int s = 0;
        for(int i=0 ; i<inputSize ; ++i) {
            s += pInput[i];
        }
        return s;
    }

    MPI_Comm subComm;
    MPI_Comm_split(comm, (me%2) + 1, me, &subComm);

    if(me == 0) {
        int sizeLocal = inputSize / 2;
        int sizeToSend = inputSize - sizeLocal;
        MPI_Ssend(&sizeToSend, 1, MPI_INT, 1, 1, comm);
        MPI_Ssend(pInput+sizeLocal, sizeToSend, MPI_INT, 1, 2, comm);
        int localSum = sumVector2(subComm, pInput, sizeLocal);
        MPI_Status status;
        int childSum;
        MPI_Recv(&childSum, 1, MPI_INT, 1, 3, comm, &status);
        MPI_Comm_free(&subComm);
        return localSum + childSum;
    }
    if(me == 1) {
        MPI_Status status;
        int sizeToRecv;
        MPI_Recv(&sizeToRecv, 1, MPI_INT, 0, 1, comm, &status);
        std::vector<int> data(sizeToRecv);
        MPI_Recv(data.data(), sizeToRecv, MPI_INT, 0, 2, comm, &status);
        int localSum = sumVector2(subComm, data.data(), sizeToRecv);
        MPI_Ssend(&localSum, 1, MPI_INT, 0, 3, comm);
    } else {
        sumVector2(subComm, nullptr, 0);
    }
    MPI_Comm_free(&subComm);
    return 0;
}

/** Processes the current node, as the member of the given communicator.
 * If the rank is 0, this process is the leader (the root of the sub-tree): in that case, the 'input' argument contains the data to be processed
 * and the function returns the result (the sum of elements in the input vector, for this problem).
 * Otherwise, this function acts as a worker
 *
 * Implementation:
 *   - if the rank is 0, this function splits the input, sends it to children, processes its own part (by doing a recursive call),
 *     receives the part results from children, combines them and, finally, returns the result.
 *   - if the rank in 1 through `nrBranches`-1, this acts as a worker, receives the data from the parent, processes its own part (by a recursive call), and sends back the result
 *   - otherwise, this does nothing at this level, but still calls itself recursively for the next level
 *   - in all cases, this splits the communicator
 * */
int sumVectorN(MPI_Comm comm, int nrBranches, int const* pInput, int inputSize) {
    int me;
    int nrProcs;
    MPI_Comm_size(comm, &nrProcs);
    MPI_Comm_rank(comm, &me);
    if(nrProcs == 1) {
        int s = 0;
        for(int i=0 ; i<inputSize ; ++i) {
            s += pInput[i];
        }
        return s;
    }

    MPI_Comm subComm;
    MPI_Comm_split(comm, (me % nrBranches) + 1, me, &subComm);

    if(me == 0) {
        int nrChildren = std::min(nrBranches, nrProcs);
        for(int i=1 ; i<nrChildren ; ++i) {
            int childStartIndex = i*inputSize/nrChildren;
            int sizeForChild = (i+1)*inputSize/nrChildren - childStartIndex;
            MPI_Ssend(&sizeForChild, 1, MPI_INT, i, 1, comm);
            MPI_Ssend(pInput + childStartIndex, sizeForChild, MPI_INT, i, 2, comm);
        }
        int sum = sumVectorN(subComm, nrBranches, pInput, inputSize/nrChildren);
        for(int i=1 ; i<nrChildren ; ++i) {
            MPI_Status status;
            int childSum;
            MPI_Recv(&childSum, 1, MPI_INT, i, 3, comm, &status);
            sum += childSum;
        }
        MPI_Comm_free(&subComm);
        return sum;
    }
    if(me < nrBranches) {
        MPI_Status status;
        int sizeToRecv;
        MPI_Recv(&sizeToRecv, 1, MPI_INT, 0, 1, comm, &status);
        std::vector<int> data(sizeToRecv);
        MPI_Recv(data.data(), sizeToRecv, MPI_INT, 0, 2, comm, &status);
        int localSum = sumVectorN(subComm, nrBranches, data.data(), sizeToRecv);
        MPI_Ssend(&localSum, 1, MPI_INT, 0, 3, comm);
    } else {
        sumVectorN(subComm, nrBranches, nullptr, 0);
    }
    MPI_Comm_free(&subComm);
    return 0;
}

#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;
        }
#ifdef NR_BRANCHES
        int v = sumVectorN(MPI_COMM_WORLD, NR_BRANCHES, input.data(), nrElements);
#else
        int v = sumVector2(MPI_COMM_WORLD, input.data(), nrElements);
#endif
        bool ok = (v == nrElements*(nrElements-1)/2);
        std::cout << "Result=" << v << (ok ? " Ok\n" : " ERROR!!!\n");
    } else {
#ifdef NR_BRANCHES
        sumVectorN(MPI_COMM_WORLD, NR_BRANCHES, nullptr, 0);
#else
        sumVector2(MPI_COMM_WORLD, nullptr, 0);
#endif
    }
    MPI_Finalize();
}