#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, testErr = 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(true) // endless loop
    {
        retErr = MPI_Waitany(mRcvRequests.size(), &*mRcvRequests.begin(), &slaveIdx, &status);
        slaveRank = slaveIdx + 1;

        if  (retErr != MPI_SUCCESS)
        {
            markFailedRanks("MPI_Waitany");

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

        if(mIsSlaveLives[slaveIdx])
        {
            // 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);
            }
            MPI_Irecv(&(mRcvNumsBuf[slaveIdx]), 1, MPI::INT, slaveRank, MPI::ANY_TAG, MPI_COMM_WORLD, &(mRcvRequests[slaveIdx]));
        }
        fflush(mFpLog);

        if(mSlavesFinished == mSlaves)
            break;
    }

    SendEndOfCommunication();

    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);
    handleMPIerror(fpLog, "MpiRecv error ", retErr, &status);
    fprintf(fpLog, "Received number = %d\n", testRcvNum);
    fprintf(fpLog, "\n\n\nEnd of slave subroutine\n");
    fflush(fpLog);
}


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
        {
            SndSyncSlave(rank, fpLog);
        }
    }

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


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

    return 0;
}