/* -*- Mode: C++; c-basic-offset:4 ; -*- */
/*
 *
 *  (C) 2003 by Argonne National Laboratory.
 *      See COPYRIGHT in top-level directory.
 */
#include "mpi.h"
#include "mpitestconf.h"
#ifdef HAVE_IOSTREAM
// Not all C++ compilers have iostream instead of iostream.h
#include <iostream>
#ifdef HAVE_NAMESPACE_STD
// Those that do need the std namespace; otherwise, a bare "cout"
// is likely to fail to compile
using namespace std;
#endif
#else
#include <iostream.h>
#endif
#include "mpitestcxx.h"

static char MTEST_Descrip[] = "A simple test of Comm_spawn, followed by intercomm merge";

int main(int argc, char *argv[])
{
    int errs = 0, err;
    int rank, size, rsize, i;
    int np = 2;
    int errcodes[2];
    MPI::Intercomm parentcomm, intercomm;
    MPI::Intracomm intracomm, intracomm2, intracomm3;
    bool isChild = false;
    MPI::Status status;

    MTest_Init();

    parentcomm = MPI::Comm::Get_parent();

    if (parentcomm == MPI::COMM_NULL) {
        /* Create 2 more processes */
        intercomm = MPI::COMM_WORLD.Spawn("./spawnintrax", MPI::ARGV_NULL, np,
                                          MPI::INFO_NULL, 0, errcodes);
    } else
        intercomm = parentcomm;

    /* We now have a valid intercomm */

    rsize = intercomm.Get_remote_size();
    size = intercomm.Get_size();
    rank = intercomm.Get_rank();

    if (parentcomm == MPI::COMM_NULL) {
        /* Master */
        if (rsize != np) {
            errs++;
            cout << "Did not create " << np << " processes (got " << rsize << ")\n";
        }
        if (rank == 0) {
            for (i = 0; i < rsize; i++) {
                intercomm.Send(&i, 1, MPI::INT, i, 0);
            }
        }
    } else {
        /* Child */
        isChild = true;
        if (size != np) {
            errs++;
            cout << "(Child) Did not create " << np << " processes (got " << size << ")\n";
        }
        intercomm.Recv(&i, 1, MPI::INT, 0, 0, status);
        if (i != rank) {
            errs++;
            cout << "Unexpected rank on child " << rank << " (" << i << ")\n";
        }
    }

    /* At this point, try to form the intracommunicator */
    intracomm = intercomm.Merge(isChild);

    /* Check on the intra comm */
    {
        int icsize, icrank, wrank;

        icsize = intracomm.Get_size();
        icrank = intracomm.Get_rank();
        wrank = MPI::COMM_WORLD.Get_rank();

        if (icsize != rsize + size) {
            errs++;
            cout << "Intracomm rank " << icrank << " thinks size is " <<
                icsize << ", not " << rsize + size << "\n";
        }
        /* Make sure that the processes are ordered correctly */
        if (isChild) {
            int psize;
            psize = parentcomm.Get_remote_size();
            if (icrank != psize + wrank) {
                errs++;
                cout << "Intracomm rank " << icrank <<
                    " (from child) should have rank " << psize + wrank << "\n";
            }
        } else {
            if (icrank != wrank) {
                errs++;
                cout << "Intracomm rank " << icrank <<
                    " (from parent) should have rank " << wrank << "\n";
            }
        }
    }

    /* At this point, try to form the intracommunicator, with the other
     * processes first */
    intracomm2 = intercomm.Merge(!isChild);

    /* Check on the intra comm */
    {
        int icsize, icrank, wrank;

        icsize = intracomm2.Get_size();
        icrank = intracomm2.Get_rank();
        wrank = MPI::COMM_WORLD.Get_rank();

        if (icsize != rsize + size) {
            errs++;
            cout << "(2)Intracomm rank " << icrank <<
                " thinks size is " << icsize << ", not " << rsize + size << "\n";
        }
        /* Make sure that the processes are ordered correctly */
        if (isChild) {
            if (icrank != wrank) {
                errs++;
                cout << "(2)Intracomm rank " << icrank <<
                    " (from child) should have rank " << wrank << "\n";
            }
        } else {
            int csize;
            csize = intercomm.Get_remote_size();
            if (icrank != wrank + csize) {
                errs++;
                cout << "(2)Intracomm rank " << icrank <<
                    " (from parent) should have rank " << wrank + csize << "\n";
            }
        }
    }

    /* At this point, try to form the intracommunicator, with an
     * arbitrary choice for the first group of processes */
    intracomm3 = intercomm.Merge(false);
    /* Check on the intra comm */
    {
        int icsize, icrank, wrank;

        icsize = intracomm3.Get_size();
        icrank = intracomm3.Get_rank();
        wrank = MPI::COMM_WORLD.Get_rank();

        if (icsize != rsize + size) {
            errs++;
            cout << "(3)Intracomm rank " << icrank <<
                " thinks size is " << icsize << ", not " << rsize + size << "\n";
        }
        /* Eventually, we should test that the processes are ordered
         * correctly, by groups (must be one of the two cases above) */
    }

    /* Update error count */
    if (isChild) {
        /* Send the errs back to the master process */
        intercomm.Ssend(&errs, 1, MPI::INT, 0, 1);
    } else {
        if (rank == 0) {
            /* We could use intercomm reduce to get the errors from the
             * children, but we'll use a simpler loop to make sure that
             * we get valid data */
            for (i = 0; i < rsize; i++) {
                intercomm.Recv(&err, 1, MPI::INT, i, 1);
                errs += err;
            }
        }
    }

    /* It isn't necessary to free the intracomms, but it should not hurt */
    intracomm.Free();
    intracomm2.Free();
    intracomm3.Free();

    /* It isn't necessary to free the intercomm, but it should not hurt */
    intercomm.Free();

    /* Note that the MTest_Finalize get errs only over COMM_WORLD */
    /* Note also that both the parent and child will generate "No Errors"
     * if both call MTest_Finalize */
    if (parentcomm == MPI::COMM_NULL) {
        MTest_Finalize(errs);
    } else {
        MPI_Finalize();
    }

    return 0;
}