#include "mpi_test_incl.h"

using namespace std;


class Master : public MasterBase
{
public:
    Master(int slaves, int iters, FILE* fpLog) : MasterBase(slaves, iters, fpLog)
    {
        mRcvRequests.resize(mSlaves);
    }
    void Rcv_WaitAny();
protected:
    vector<MPI_Request> mRcvRequests;
    virtual void updateOnFail();
};

void Master::updateOnFail()
{
    vector<int>::iterator it = mIsSlaveLives.begin();
    for(int slaveIdx=0; it!= mIsSlaveLives.end(); ++it, ++slaveIdx)
    {
        if(!*it)
            mRcvRequests[slaveIdx] = MPI_REQUEST_NULL;
    }
}

void Master::Rcv_WaitAny()
{
    MPI_Request send_request, recv_request;
    MPI_Status  status;

    int  retErr     = MPI_SUCCESS;
    int  testSndNum;
    int  slaveRank;
    int  slaveIdx;
    int  flag;

    for(slaveIdx=0; slaveIdx<mSlaves; ++slaveIdx)
    {
        slaveRank = slaveIdx + 1;
        MPI_Irecv(&(mRcvNumsBuf[slaveIdx]), 1, MPI::INT, slaveRank, MPI::ANY_TAG, MPI_COMM_WORLD, &(mRcvRequests[slaveIdx]));
    }

    while(mSlavesFinished < mSlaves)
    {
        retErr = MPI_Waitany(mRcvRequests.size(), &*mRcvRequests.begin(), &slaveIdx, &status);
        slaveRank = slaveIdx + 1;

        /* if all requests have been completed, we should have exited the loop already */
        assert(slaveIdx != MPI_UNDEFINED);
        /* if the slave is dead, we should not be able to receive a message from it */
        assert(mIsSlaveLives[slaveIdx]);

        if  (retErr != MPI_SUCCESS)
        {
            // markFailedRanks("MPI_Waitany");
            char Msg[256];
            sprintf(Msg, "From rank %d, fail - request deallocated", slaveRank);
            handleMPIerror(mFpLog, Msg, retErr, &status);
            mRcvRequests[slaveIdx] = MPI_REQUEST_NULL;
            mIsSlaveLives[slaveIdx] = 0;
            ++mSlavesFinished;
            continue;
        }

        // fprintf(mFpLog, "From rank %d, Got number = %d\n ", slaveRank, mRcvNumsBuf[slaveIdx]);
        ++mSlavesRcvIters[slaveIdx];
        if(mSlavesRcvIters[slaveIdx] == nIters)
        {
            ++mSlavesFinished;
            fprintf(mFpLog, "\n\nFrom rank %d, Got number = %d\n ", slaveRank, mRcvNumsBuf[slaveIdx]);
            fprintf(mFpLog, "Slave %d finished\n\n", slaveIdx+1);
        }
        else
        {
            MPI_Irecv(&(mRcvNumsBuf[slaveIdx]), 1, MPI::INT, slaveRank, MPI::ANY_TAG, MPI_COMM_WORLD, &(mRcvRequests[slaveIdx]));
        }
        fflush(mFpLog);
    }

    // SendEndOfCommunication();
    // RecvEndOfCommunication();
    /*
    retErr = MPI_Abort(MPI_COMM_WORLD, 0);
    retErr = MPI_Finalize();
     */

    fprintf(mFpLog, "\n\n\nEnd of master subroutine\n");
}


void SndSyncSlave(int rank, FILE* fpLog)
{
    MPI_Status  status;
    MPI_Request request;

    int masterRank  = 0;
    int testSndNum = -1, testRcvNum = -1;
    int retErr;
    int passIter = 0;

    for(int i=1; i<nIters+1; ++i)
    {
        testSndNum = rank * Scale + i;
        // retErr = W_MpiSendSync(&testSndNum, 1, MPI_INT, masterRank, 0, MPI_COMM_WORLD);
        retErr = MPI_Send(&testSndNum, 1, MPI_INT, masterRank, 0, MPI_COMM_WORLD);
        handleMPIerror(fpLog, "W_MpiSendSync error ", retErr);
        // fprintf(fpLog, "Sent number = %d  \n", testSndNum);

        FailFunction(rank, i, testSndNum, fpLog);
    }

    fprintf(fpLog, "Last Sent number = %d\n", testSndNum);

    // Wait for final message from master
    // retErr = MPI_Recv(&testRcvNum, 1, MPI_INT, masterRank, MPI_ANY_TAG, MPI_COMM_WORLD, &status);
    // retErr = MPI_Send(&testRcvNum, 1, MPI_INT, masterRank, 0, MPI_COMM_WORLD);
    // handleMPIerror(fpLog, "MpiRecv error ", retErr, &status);
    fprintf(fpLog, "Received number = %d\n", testRcvNum);
    fprintf(fpLog, "\n\n\nEnd of slave subroutine\n");
    fflush(fpLog);
    // retErr = MPI_Finalize();
}


void SndSyncSlave2(int rank, int size, FILE* fpLog)
{
    vector<MPI_Request> rcvRequests;
    vector<int>         rcvNums;

    MPI_Status  status;
    MPI_Request request;
    int  completed;

    int masterRank  = 0;
    int testSndNum = -1, testRcvNum = -1;
    int retErr;
    int passIter = 0;

    rcvRequests.resize(size);
    rcvNums.resize(size);

    for(int i=0; i<size; ++i) {
        MPI_Irecv(&(rcvNums[i]), 1, MPI::INT, 0, MPI::ANY_TAG, MPI_COMM_WORLD, &(rcvRequests[i]));
    }

    for(int i=1; i<nIters+1; ++i)
    {
        testSndNum = rank * Scale + i;
        // retErr = W_MpiSendSync(&testSndNum, 1, MPI_INT, masterRank, 0, MPI_COMM_WORLD);
        retErr = MPI_Send(&testSndNum, 1, MPI_INT, masterRank, 0, MPI_COMM_WORLD);
        handleMPIerror(fpLog, "W_MpiSendSync error ", retErr);
        // fprintf(fpLog, "Sent number = %d  \n", testSndNum);

        FailFunction(rank, i, testSndNum, fpLog);
        // fprintf(fpLog, "Got number = %d  \n", testRcvNum);
    }

    fprintf(fpLog, "Last Sent number = %d\n", testSndNum);
    fflush(fpLog);

    int slaveIdx;
    retErr = MPI_Waitany(rcvRequests.size(), &rcvRequests.front(), &slaveIdx, &status);
    if(retErr == MPI_SUCCESS) {
        fprintf(fpLog, "Got number = %d, rank = %d, tag = %d \n", testRcvNum, status.MPI_SOURCE, status.MPI_TAG);
    }
    else {
        fprintf(fpLog, "Error in MPI_Waitany: rank = %d, tag = %d, error = %d\n", status.MPI_SOURCE, status.MPI_TAG, status.MPI_ERROR);
    }

    // Wait for final message from master
    // retErr = MPI_Recv(&testRcvNum, 1, MPI_INT, masterRank, MPI_ANY_TAG, MPI_COMM_WORLD, &status);
    // retErr = MPI_Send(&testRcvNum, 1, MPI_INT, masterRank, 0, MPI_COMM_WORLD);
    // handleMPIerror(fpLog, "MpiRecv error ", retErr, &status);
    // fprintf(fpLog, "Received number = %d\n", testRcvNum);

    fprintf(fpLog, "\n\n\nEnd of slave subroutine\n");
    fflush(fpLog);
    // retErr = MPI_Finalize();
}


int main(int argc, char *argv[])
{
    int  size = 1, rank = 0;
    bool isAllSlavesLive = true;
    FILE* fpLog;

    int retErr = mainInit(argc, argv, &size, &rank, &fpLog);
    if (retErr < 0)
        return 0;

    try
    {
        if(rank == 0)
        {
            Master master(size-1, nIters, fpLog);
            master.Rcv_WaitAny();
            isAllSlavesLive = master.IsAllSlavesLive();
        }
        else
        {
            // SndSyncSlave2(rank, size, fpLog);
            SndSyncSlave(rank, fpLog);
        }
    }

    catch(...)
    {
        fprintf(fpLog, "\n\n\nunexpected C++ exception\n\n");
    }

    bool isAbort = ( (rank == 0) && (!isAllSlavesLive) );
    // MainFinalize(fpLog,  isAbort);
    MainFinalize2(fpLog, isAbort, (rank == 0));

    return 0;
}







/*
            for(int testIdx=0; testIdx<mSlaves; ++testIdx)
            {
                if(!mIsSlaveLives[slaveIdx])
                    continue;

                if ( (testErr = MPI_Test(&(mRcvRequests[testIdx]), &flag, &status) ) != MPI_SUCCESS)
                {
                    char Msg[256];
                    sprintf(Msg, "From rank %d, fail - request deallocated", testIdx+1);
                    handleMPIerror(mFpLog, Msg, testErr, &status);
                    mRcvRequests[testIdx] = MPI_REQUEST_NULL;
                    mIsSlaveLives[testIdx] = 0;
                    ++mSlavesFinished;
                }
            }

            // markFailedRanks("MPI_Waitany");
*/