#include "mpi_test_incl.h"

class Master : public MasterBase
{
public:
    Master(int slaves, int iters, FILE* fpLog) : MasterBase(slaves, iters, fpLog) {}
    // void SendReceiveAsyncP2P();
    // void SendReceiveAsyncP2MP();
    // void SendReceiveAsyncP2MP_2();
    void SendReceiveAsyncP2MP_C();

    virtual void SendEndOfCommunication(MPI::Intracomm* comm = NULL);
protected:
    // int Waitany(int count, MPI::Request* array_of_requests, int* index,  MPI::Status* status);
    int Waitany_C(int count, MPI_Request* array_of_requests, int* outIndex, MPI_Status* status);
};

/*
void Master::SendReceiveAsyncP2MP_2()
{
    MPI_Request send_request, recv_request;
    int         rcvInt, sendInt = 0;
    int         retErr  = MPI_SUCCESS;
    int         rcvIter = 0;
    int         slaveRank;

    int         nodeFail = 0;
    int         opCompleted = 1;
    vector<MPI::Request> rcvRequests(mSlaves);
    MPI::Status status;

    // Init MPI_Irecv requests.
    for(int slaveIdx=0; slaveIdx<mSlaves; ++slaveIdx)
    {
        slaveRank = slaveIdx + 1;
        rcvRequests[slaveIdx] = MPI::COMM_WORLD.Irecv(&(mRcvNumsBuf[slaveIdx]), 1, MPI::INT, slaveRank, MPI::ANY_TAG);
    }

    int reqIdx = -1;
    while(mSlavesFinished < mSlaves)
    {
        retErr = Waitany(rcvRequests.size(), &rcvRequests.front(), &reqIdx, &status);
        slaveRank = reqIdx + 1;

        if (retErr != MPI_SUCCESS)
        {
            handleMPIerror(mFpLog, "Master::SendReceiveAsyncP2MP_2 ", retErr, (MPI_Status*) status);
            mIsSlaveLives[reqIdx] = false;
            ++mSlavesFinished;
            continue;
        }

        // Got something from slave on this iteration
        fprintf(mFpLog, "From rank %d, Got number = %d ", slaveRank, mRcvNumsBuf[reqIdx]);
        fflush(mFpLog);
        ++mSlavesRcvIters[reqIdx];
        if(mSlavesRcvIters[reqIdx] == nIters) // slave sent last iteration
        {
            ++mSlavesFinished;
            fprintf(mFpLog, "\n\nSlave %d finished\n\n", slaveRank);
            fflush(mFpLog);
        }
        else
        {
            mSndNumsBuf[reqIdx] = mRcvNumsBuf[reqIdx] - Scale * slaveRank;
            fprintf(mFpLog, "sent number = %d \n", mSndNumsBuf[reqIdx]);
            fflush(mFpLog);
            MPI_Isend(&(mSndNumsBuf[reqIdx]), 1, MPI_INT, slaveRank, 0, MPI_COMM_WORLD, &send_request); // no need waiting here
            rcvRequests[reqIdx] = MPI::COMM_WORLD.Irecv(&(mRcvNumsBuf[reqIdx]), 1, MPI::INT, slaveRank, MPI::ANY_TAG);
        }
    } // end of while(mSlavesFinished < mSlaves)

    MasterBase::SendEndOfCommunication();

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

int Master::Waitany(int count, MPI::Request* array_of_requests, int* index,  MPI::Status* status)
{
    int funcRetErr = MPI_SUCCESS;

    fprintf(mFpLog, "GES_FTMPIService::Waitany - start    ");
    try
    {
        int opCompleted = 0;

        // Patch to avoid search from same index, use value of index as previous index
        if((*index < 0) || (*index >= count)) // if *index is out of range init it with last index in the vector.
        {
            *index = count-1;
        }

        MPI::Request* req  = array_of_requests + *index;
        MPI_Status*   stat = (MPI_Status*)status;

        int startIndex = *index;
        int isCancelled = 0;
        while( (funcRetErr == MPI_SUCCESS) && (!opCompleted) ) // && (!isCancelled) )
        {
            ++req;
            *index += 1;
            if(*index == count)
            {
                *index  = 0;
                req     = array_of_requests;
                usleep(100);
            }

            if( *req != MPI::Request(MPI_REQUEST_NULL) )
            {
                funcRetErr = MPI_Test((MPI_Request*)*req, &opCompleted, stat);
                // funcRetErr = MPI_Request_get_status(*(MPI_Request*)*req, &opCompleted, stat);
            }
        }

        if (funcRetErr != MPI_SUCCESS)
        {
            fprintf(mFpLog, "GES_FTMPIService::Waitany - Fail of index %d\n", *index);
            *req = MPI::Request(MPI_REQUEST_NULL);
        }
        else // (funcRetErr == MPI_SUCCESS)
        {
            MPI_Test_cancelled((MPI_Status*)*status, &isCancelled);
            if (isCancelled)
            {
                funcRetErr = MPI_ERR_OTHER;
                fprintf(mFpLog, "GES_FTMPIService::Waitany - request %d isCancelled", *index);
            }
            // MPI_Wait((MPI_Request*)*req, stat);
            // MPI_Request_free((MPI_Request*)*req);
        }
    }
    catch(...)
    {
        std::cout << "Master::Waitany Unhandled exception"<<endl  ;
    }

    fprintf(mFpLog, "GES_FTMPIService::Waitany - end\n\n");
    fflush(mFpLog);
    return funcRetErr;
}
*/
void Master::SendReceiveAsyncP2MP_C()
{
    MPI_Request send_request, recv_request;
    int         rcvInt, sendInt = 0;
    int         retErr  = MPI_SUCCESS;
    int         rcvIter = 0;
    int         slaveRank;

    int         nodeFail = 0;
    int         opCompleted = 1;
    vector<MPI_Request> rcvRequests(mSlaves);
    MPI_Status status;

    // Init MPI_Irecv requests.
    for(int slaveIdx=0; slaveIdx<mSlaves; ++slaveIdx)
    {
        slaveRank = slaveIdx + 1;
        MPI_Irecv(&(mRcvNumsBuf[slaveIdx]), 1, MPI::INT, slaveRank, MPI::ANY_TAG, MPI_COMM_WORLD, &(rcvRequests[slaveIdx]));
    }

    int reqIdx = -1;
    while(mSlavesFinished < mSlaves)
    {
        retErr = Waitany_C(rcvRequests.size(), &rcvRequests.front(), &reqIdx, &status);
        slaveRank = reqIdx + 1;

        if ( (retErr != MPI_SUCCESS) && mIsSlaveLives[reqIdx] )
        {
            handleMPIerror(mFpLog, "Master::SendReceiveAsyncP2MP_C ", retErr, &status);
            mIsSlaveLives[reqIdx] = false;
            ++mSlavesFinished;
            continue;
        }

        // Got something from slave on this iteration
        // fprintf(mFpLog, "From rank %d, Got number = %d ", slaveRank, mRcvNumsBuf[reqIdx]);
        // fflush(mFpLog);
        ++mSlavesRcvIters[reqIdx];
        if(mSlavesRcvIters[reqIdx] == nIters) // slave sent last iteration
        {
            ++mSlavesFinished;
            fprintf(mFpLog, "\n\nSlave %d finished\n\n", slaveRank);
            fflush(mFpLog);
        }
        else
        {
            mSndNumsBuf[reqIdx] = mRcvNumsBuf[reqIdx] - Scale * slaveRank;
            // fprintf(mFpLog, "sent number = %d \n", mSndNumsBuf[reqIdx]);
            // fflush(mFpLog);
            MPI_Isend(&(mSndNumsBuf[reqIdx]), 1, MPI_INT, slaveRank, 0, MPI_COMM_WORLD, &send_request); // no need waiting here
            MPI_Irecv(&(mRcvNumsBuf[reqIdx]), 1, MPI::INT, slaveRank, MPI::ANY_TAG, MPI_COMM_WORLD, &(rcvRequests[reqIdx]));
        }
    } // end of while(mSlavesFinished < mSlaves)

    MasterBase::SendEndOfCommunication();

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


int Master::Waitany_C(int count, MPI_Request* array_of_requests, int* outIndex, MPI_Status* status)
{
    int funcRetErr = MPI_SUCCESS;

//    fprintf(mFpLog, "GES_FTMPIService::Waitany_C - start    ");

    int opCompleted = 0;

    // Patch to avoid search from same index, use value of index as previous index
    int           index    = count-1;
    MPI_Request*  request  = array_of_requests + index;


    while( (funcRetErr == MPI_SUCCESS) && (!opCompleted) )
    {
        ++request;
        ++index;
        if(index == count)
        {
            index   = 0;
            request = array_of_requests;
            usleep(100);
        }

        if( *request != MPI_REQUEST_NULL )
        {
            funcRetErr = MPI_Test(request, &opCompleted, status);
        }
    }

    if (funcRetErr != MPI_SUCCESS)
    {
        fprintf(mFpLog, "GES_FTMPIService::WaitanyC - Fail of index %d\n", index);
        // MPI_Request_free(request);
        *request = MPI_REQUEST_NULL;
    }

 //   fprintf(mFpLog, "GES_FTMPIService::Waitany_C - end\n\n");
    fflush(mFpLog);

    *outIndex = index;

    return funcRetErr;
}



/*
void Master::SendReceiveAsyncP2MP()
{
    MPI_Request send_request, recv_request;
    MPI_Status  status;
    int         flag;

    int         rcvInt, sendInt = 0;
    int         retErr  = MPI_SUCCESS;
    int         rcvIter = 0;
    int         slaveRank;

    int         nodeFail = 0;
    int         opCompleted = 1;
    vector<MPI_Request> rcvRequests(mSlaves);
    vector<MPI_Status>  statuses(mSlaves);


    // Init MPI_Irecv requests.
    for(int slaveIdx=0; slaveIdx<mSlaves; ++slaveIdx)
    {
        slaveRank = slaveIdx + 1;
        MPI_Irecv(&(mRcvNumsBuf[slaveIdx]), 1, MPI::INT, slaveRank, MPI::ANY_TAG, MPI_COMM_WORLD, &(rcvRequests[slaveIdx]));
    }

    while(mSlavesFinished < mSlaves)
    {
        int reqIdx = -1;
        retErr = MPI_Testany(rcvRequests.size(), &rcvRequests.front(), &reqIdx, &flag, &status);
        slaveRank = reqIdx + 1;

        if( (flag == true) && (reqIdx == MPI_UNDEFINED) )
        {
            fprintf(mFpLog, "No active requests!!!\n");
            fflush(mFpLog);
            usleep(100000);
            break;
        }

        if( (flag == false) && (reqIdx == MPI_UNDEFINED) )
        {   // no completed requests.
            fprintf(mFpLog, "No completed requests  ");
            fprintf(mFpLog, "Status_ERR=%d ", status.MPI_ERROR);
            fprintf(mFpLog, "Status_Src %d \n", status.MPI_SOURCE);
            fflush(mFpLog);
            usleep(100000);
            continue;
        }
        if (retErr == MPI_ERR_IN_STATUS)
        {
            fprintf(mFpLog, "MPI_ERR_IN_STATUS\n");
            fflush(mFpLog);
            if (status.MPI_ERROR != MPI_SUCCESS)
            {
                mIsSlaveLives[reqIdx] = false;
                ++mSlavesFinished;
                fprintf(mFpLog, "Error: Node %d failed on recieve\n", slaveRank);
                handleMPIerror(mFpLog, "Node failed on recieve", retErr, &status);
                fflush(mFpLog);
                MPI_Request_free(&(rcvRequests[reqIdx]));
            }
            continue;
        }

//        fprintf(mFpLog, "retErr = %d, reqIdx = %d, flag = %d\n", retErr, reqIdx, flag);
//        fflush(mFpLog);

        // Got something from slave on this iteration
        fprintf(mFpLog, "From rank %d, Got number = %d ", slaveRank, mRcvNumsBuf[reqIdx]);
        fflush(mFpLog);
        ++mSlavesRcvIters[reqIdx];
        if(mSlavesRcvIters[reqIdx] == nIters) // slave sent last iteration
        {
            ++mSlavesFinished;
            fprintf(mFpLog, "\n\nSlave %d finished\n\n", slaveRank);
            fflush(mFpLog);
        }
        else
        {
            mSndNumsBuf[reqIdx] = mRcvNumsBuf[reqIdx] - Scale * slaveRank;
            fprintf(mFpLog, "sent number = %d \n", mSndNumsBuf[reqIdx]);
            fflush(mFpLog);
            MPI_Isend(&(mSndNumsBuf[reqIdx]), 1, MPI_INT, slaveRank, 0, MPI_COMM_WORLD, &send_request); // no need waiting here
            MPI_Irecv(&(mRcvNumsBuf[reqIdx]), 1, MPI::INT, slaveRank, MPI::ANY_TAG, MPI_COMM_WORLD, &(rcvRequests[reqIdx]));
        }
    } // end of while(mSlavesFinished < mSlaves)

    MasterBase::SendEndOfCommunication();

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

/*
void Master::SendReceiveAsyncP2P()
{
    MPI_Request send_request, recv_request;
    MPI_Status  status;
    int         flag;

    int         rcvInt, sendInt = 0;
    int         retErr  = MPI_SUCCESS;
    int         rcvIter = 0;
    int         slaveRank = 1;

    int         nodeFail = 0;
    int         opCompleted = 1;

    while( (rcvIter < nIters) && (!nodeFail))
    {
        if(opCompleted)
            MPI_Irecv(&rcvInt, 1, MPI::INT, slaveRank, MPI::ANY_TAG, MPI_COMM_WORLD, &recv_request);

        fprintf(mFpLog, "MPI_Test loop started ...  ");
        fflush(mFpLog);
        do
        {
            nodeFail = (MPI_Test(&recv_request, &opCompleted, &status) != MPI_SUCCESS);
        } while( (!opCompleted) && (!nodeFail) );
        fprintf(mFpLog, "finished \n");
        fflush(mFpLog);


        if(nodeFail)
        {
            fprintf(mFpLog, "Node failed on recieve\n");
            fflush(mFpLog);
            MPI_Request_free(&recv_request);
            break;
        }

        // receive completed OK
        fprintf(mFpLog, "From rank %d, Got number = %d  ", slaveRank, rcvInt);
        ++rcvIter;

        // Send reply
        sendInt = rcvInt - slaveRank * Scale;
        fprintf(mFpLog, "To slave rank = %d  send %d ", slaveRank, sendInt);
        MPI_Isend(&sendInt, 1, MPI_INT, slaveRank, 0, MPI_COMM_WORLD, &send_request);
        retErr = MPI_Wait(&send_request, &status);
        handleMPIerror(mFpLog, "Send error ", retErr);
        if(retErr != MPI_SUCCESS)
           break;

        fprintf(mFpLog, "\n");
        fflush(mFpLog);
    }

    if(!nodeFail)
        SendEndOfCommunication();

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

void SndRcvSlave(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; testRcvNum != 0; ++i)
    {

        testSndNum = rank * Scale + i;
        MPI_Isend(&testSndNum, 1, MPI_INT, masterRank, 0, MPI_COMM_WORLD, &request);
        retErr = MPI_Wait(&request, &status);
        handleMPIerror(fpLog, "W_MpiSendSync error ", retErr);
        // fprintf(fpLog, "Sent number = %d  ", testSndNum);

        FailFunction(rank, i, testSndNum, fpLog);

        MPI_Irecv(&testRcvNum, 1, MPI_INT, masterRank, MPI_ANY_TAG, MPI_COMM_WORLD, &request);
        retErr = MPI_Wait(&request, &status);


        // 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, "Last Sent number = %d\n", testSndNum);

    // Wait for final message from master


    fprintf(fpLog, "\n\n\nEnd of slave subroutine\n");
}

void Master::SendEndOfCommunication(MPI::Intracomm* comm)
{
    int retErr;
    for(int i=0; i<mSlaves; ++i)
    {
        int slaveRank = i + 1;

        if(comm == NULL)
            retErr = MPI_Send((int*)&EndOfSession, 1, MPI_INT, slaveRank, 0, MPI_COMM_WORLD);
        else
            retErr = MPI_Send((int*)&EndOfSession, 1, MPI_INT, slaveRank, 0, *comm);
        fprintf(mFpLog, "  - Sent exit message to slave rank %d\n", slaveRank);
        handleMPIerror(mFpLog, "Exit message error ", retErr);
    }
    fflush(mFpLog);
}




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.SendReceiveAsyncP2MP_C();
            // master.SendReceiveAsyncP2MP_2();
            // master.SendReceiveAsyncP2MP();
            isAllSlavesLive = master.IsAllSlavesLive();
        }
        else
        {
            SndRcvSlave(rank, fpLog);
        }
    }

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

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

    return 0;
}