#include "mpi.h"
#include "mpix.h"
#include <complex>
#include <iostream>
#include <sstream>
#include <stdexcept>
#include <vector>
using namespace std;

int 
main( int argc, char* argv[] )
{
    MPI_Init( &argc, &argv );

    // Extract the total number of processes
    int p;
    MPI_Comm_size( MPI_COMM_WORLD, &p );

    // Create custom BG/P 4D torus communicator
    int torus_rank;
    int result;
    MPI_Comm torus_comm;
    MPIX_Cart_comm_create( &torus_comm );
    MPI_Comm_rank( torus_comm, &torus_rank  );

    result = 0;
    MPIX_Get_property(torus_comm, MPIDO_RECT_COMM, &result);
    if (torus_rank == 0) {
      if(result) fprintf(stderr, "Get_property: world is a rect. comm\n");
    }

    // Set up the MatrixCol communicator
    // (Recall that MPIX_Cart_comm_create orders [T,Z,Y,X]) 
    int remainingDimensions[4];
    remainingDimensions[0] = true; // keep T
    remainingDimensions[1] = true; // keep Z
    remainingDimensions[2] = false;
    remainingDimensions[3] = false;
    int r;
    int matrix_col_rank;
    MPI_Comm matrix_col_comm;
    MPI_Cart_sub( torus_comm, remainingDimensions, &matrix_col_comm );
    MPI_Comm_rank( matrix_col_comm, &matrix_col_rank );
    MPI_Comm_size( matrix_col_comm, &r );

    result = 0;
    MPIX_Get_property(matrix_col_comm, MPIDO_RECT_COMM, &result);
    if (matrix_col_rank == 0) {
      if(result) fprintf(stderr, "Get_property: matrix_col_comm is a rect. sub. comm\n");
    }

    // Set up the MatrixRow communicator
    remainingDimensions[0] = false;
    remainingDimensions[1] = false;
    remainingDimensions[2] = true; // keep Y
    remainingDimensions[3] = true; // keep X
    int c;
    int matrix_row_rank;
    MPI_Comm matrix_row_comm;
    MPI_Cart_sub( torus_comm, remainingDimensions, &matrix_row_comm );
    MPI_Comm_rank( matrix_row_comm, &matrix_row_rank );
    MPI_Comm_size( matrix_row_comm, &c );

    MPIX_Get_property(matrix_row_comm, MPIDO_RECT_COMM, &result);
    if (matrix_row_rank == 0) {
      if(result) fprintf(stderr, "Get_property: matrix_row_comm is a rect. sub. comm\n");
    }

    if( torus_rank == 0 )
    {
        ostringstream msg;
        msg << "Using " << r << " x " << c << " grid." << endl;
        cout << msg.str();
    }

    // We will simulate an Allgather of a column panel of size 40,000 x 64
    int m = 40000;
    int b = 64;
    int num_tests = 20;
    vector<double> local_data;
    vector<double> gathered_data;
    double start_time, stop_time, average_time, bandwidth;

/*
#define MPIDO_USE_ABCAST_ALLGATHER                                         26
#define MPIDO_USE_ABINOM_BCAST_ALLGATHER                                   27
#define MPIDO_USE_ALLTOALL_ALLGATHER                                       28
#define MPIDO_USE_ARECT_BCAST_ALLGATHER                                    29
#define MPIDO_USE_ALLREDUCE_ALLGATHER                                      30
#define MPIDO_USE_BCAST_ALLGATHER                                          31
#define MPIDO_USE_BINOM_BCAST_ALLGATHER                                    32
#define MPIDO_USE_RECT_BCAST_ALLGATHER                                     33
#define MPIDO_USE_RECT_DPUT_ALLGATHER                                      34
#define MPIDO_USE_MPICH_ALLGATHER                                          35  
*/
    vector<int> allgather_codes(10);
    for( unsigned i=0; i<allgather_codes.size(); ++i )
        allgather_codes[i] = 26+i;

    vector<int> allgather_props(allgather_codes.size());
    for( unsigned i=0; i<allgather_codes.size(); ++i )
    {
        MPIX_Get_property
        ( torus_comm, allgather_codes[i], &allgather_props[i] );
    }

    if( torus_rank == 0 )
    {
        ostringstream msg;
        msg << "torus_comm properties: " << endl;
        for( unsigned j=0; j<allgather_codes.size(); ++j )
            msg << "  property " << allgather_codes[j] << ": " 
                << allgather_props[j] << endl;
        cout << msg.str() << endl;
    }

    // Create a list of all of the nonzero properties
    vector<int> avail_props(0);
    for( unsigned i=0; i<allgather_codes.size(); ++i )
    {
        if( allgather_props[i] )
            avail_props.push_back(i);
    }

    // Benchmark default Allgather on torus_comm. 
    int local_size = ((m+(p-1))/p*b)/c;
    int gathered_size = local_size*r*c;
    double Mb = 64.*gathered_size/1.e6;
    local_data.resize( local_size );
    gathered_data.resize( gathered_size );
    MPI_Barrier( torus_comm );
    start_time = MPI_Wtime();
    for( unsigned i=0; i<num_tests; ++i )
    {
        MPI_Allgather
        ( &local_data[0], local_size, MPI_DOUBLE,
          &gathered_data[0], local_size, MPI_DOUBLE, torus_comm );
    }
    stop_time = MPI_Wtime();
    average_time = (stop_time-start_time)/num_tests;
    bandwidth = 64.*gathered_size/(1.e6*average_time);
    if( torus_rank == 0 )
    {
        ostringstream msg;
        msg << "world Allgather (default): " 
            << average_time << " seconds, " 
            << Mb << " Megabits, "
            << bandwidth << " Mbps." << endl;
        cout << msg.str() << endl;
    }

    // Force only one property at a time to be nonzero
    // for( unsigned i=0; i<avail_props.size(); ++i )
//     {
//         for( unsigned j=0; j<allgather_codes.size(); ++j )
//         {
//             if( avail_props[i] == j )
//                 MPIX_Set_property( torus_comm, allgather_codes[j], 1 );
//             else
//                 MPIX_Set_property( torus_comm, allgather_codes[j], 0 );
//         }

//         MPI_Barrier( torus_comm );
//         start_time = MPI_Wtime();
//         for( unsigned i=0; i<num_tests; ++i )
//         {
//             MPI_Allgather
//             ( &local_data[0], local_size, MPI_DOUBLE,
//               &gathered_data[0], local_size, MPI_DOUBLE, torus_comm );
//         }
//         stop_time = MPI_Wtime();
//         average_time = (stop_time-start_time)/num_tests;
//         bandwidth = 64.*gathered_size/(1.e6*average_time);
//         if( torus_rank == 0 )
//         {
//             ostringstream msg;
//             msg << "world Allgather (property " 
//                 << allgather_codes[avail_props[i]] << "): "
//                 << average_time << " seconds, "
//                 << Mb << " Megabits, "
//                 << bandwidth << " Mbps." << endl;
//             cout << msg.str() << endl;
//         }
//     }

//     // Force the MPIDO_USE_RECT_BCAST flag to 1
//     for( unsigned j=0; j<allgather_codes.size(); ++j )
//         MPIX_Set_property( torus_comm, allgather_codes[j], 0 );
//     MPIX_Set_property( torus_comm, 34, 1 );
//     MPI_Barrier( torus_comm );
//     start_time = MPI_Wtime();
//     for( unsigned i=0; i<num_tests; ++i )
//     {
//         MPI_Allgather
//         ( &local_data[0], local_size, MPI_DOUBLE,
//           &gathered_data[0], local_size, MPI_DOUBLE, torus_comm );
//     }
//     stop_time = MPI_Wtime();
//     average_time = (stop_time-start_time)/num_tests;
//     bandwidth = 64.*gathered_size/(1.e6*average_time);
//     if( torus_rank == 0 )
//     {
//         ostringstream msg;
//         msg << "torus Allgather (property 34): " 
//             << average_time << " seconds, "
//             << Mb << " Megabits, "
//             << bandwidth << " Mbps." << endl;
//         cout << msg.str() << endl;
//     }

    // Get and print the matrix_col_comm properties
    for( unsigned i=0; i<allgather_codes.size(); ++i )
    {
        MPIX_Get_property
        ( matrix_col_comm, allgather_codes[i], &allgather_props[i] );
    }

    if( torus_rank == 0 )
    {
        ostringstream msg;
        msg << "matrix_col_comm properties: " << endl;
        for( unsigned j=0; j<allgather_codes.size(); ++j )
            msg << "  property " << allgather_codes[j] << ": " 
                << allgather_props[j] << endl;
        cout << msg.str() << endl;
    }

    // Create a list of all of the nonzero properties
    avail_props.resize(0);
    for( unsigned i=0; i<allgather_codes.size(); ++i ) {
        if( allgather_props[i] )
            avail_props.push_back(i);
    }

    // Benchmark default Allgather on matrix_col_comm. [VR,* ] <- [MR,* ]
    local_size = (m+(p-1))/p*b;
    gathered_size = local_size*r;
    Mb = 64.*gathered_size/1.e6;
    local_data.resize( local_size );
    gathered_data.resize( gathered_size );
    MPI_Barrier( matrix_col_comm );
    start_time = MPI_Wtime();
    for( unsigned i=0; i<num_tests; ++i )
    {
        MPI_Allgather
        ( &local_data[0], local_size, MPI_DOUBLE,
          &gathered_data[0], local_size, MPI_DOUBLE, matrix_col_comm );
    }
    stop_time = MPI_Wtime();
    average_time = (stop_time-start_time)/num_tests;
    bandwidth = 64.*gathered_size/(1.e6*average_time);
    if( torus_rank == 0 )
    {
        ostringstream msg;
        msg << "matrix_col Allgather: "
            << average_time << " seconds, "
            << Mb << " Megabits,"
            << bandwidth << " Mbps." << endl;
        cout << msg.str() << endl;
    }
    
//     // Force only one property at a time to be nonzero
//     for( unsigned i=0; i<avail_props.size(); ++i )
//     {
//         for( unsigned j=0; j<allgather_codes.size(); ++j )
//         {
//             if( avail_props[i] == j )
//                 MPIX_Set_property( matrix_col_comm, allgather_codes[j], 1 );
//             else
//                 MPIX_Set_property( matrix_col_comm, allgather_codes[j], 0 );
//         }

//         MPI_Barrier( torus_comm );
//         start_time = MPI_Wtime();
//         for( unsigned i=0; i<num_tests; ++i )
//         {
//             MPI_Allgather
//             ( &local_data[0], local_size, MPI_DOUBLE,
//               &gathered_data[0], local_size, MPI_DOUBLE, matrix_col_comm );
//         }
//         stop_time = MPI_Wtime();
//         average_time = (stop_time-start_time)/num_tests;
//         bandwidth = 64.*gathered_size/(1.e6*average_time);
//         if( torus_rank == 0 )
//         {
//             ostringstream msg;
//             msg << "matrix_col Allgather (property " 
//                 << allgather_codes[avail_props[i]] << "): "
//                 << average_time << " seconds, "
//                 << Mb << " Megabits, "
//                 << bandwidth << " Mbps." << endl;
//             cout << msg.str() << endl;
//         }
//     }

//     // Force the MPIDO_USE_RECT_BCAST flag to 1
//     for( unsigned j=0; j<allgather_codes.size(); ++j )
//         MPIX_Set_property( matrix_col_comm, allgather_codes[j], 0 );
//     MPIX_Set_property( matrix_col_comm, 34, 1 );
//     MPI_Barrier( torus_comm );
//     start_time = MPI_Wtime();
//     for( unsigned i=0; i<num_tests; ++i )
//     {
//         MPI_Allgather
//         ( &local_data[0], local_size, MPI_DOUBLE,
//           &gathered_data[0], local_size, MPI_DOUBLE, matrix_col_comm );
//     }
//     stop_time = MPI_Wtime();
//     average_time = (stop_time-start_time)/num_tests;
//     bandwidth = 64.*gathered_size/(1.e6*average_time);
//     if( torus_rank == 0 )
//     {
//         ostringstream msg;
//         msg << "matrix_col Allgather (property 34): " 
//             << average_time << " seconds, "
//             << Mb << " Megabits, "
//             << bandwidth << " Mbps." << endl;
//         cout << msg.str() << endl;
//     }

//     // Get and print the matrix_row_comm properties
//     for( unsigned i=0; i<allgather_codes.size(); ++i )
//     {
//         MPIX_Get_property
//         ( matrix_row_comm, allgather_codes[i], &allgather_props[i] );
//     }

//     if( torus_rank == 0 )
//     {
//         ostringstream msg;
//         msg << "matrix_row_comm properties: " << endl;
//         for( unsigned j=0; j<allgather_codes.size(); ++j )
//             msg << "  property " << allgather_codes[j] << ": " 
//                 << allgather_props[j] << endl;
//         cout << msg.str() << endl;
//     }

//     // Create a list of all of the nonzero properties
//     avail_props.resize(0);
//     for( unsigned i=0; i<allgather_codes.size(); ++i )    {
//         if( allgather_props[i] )
//             avail_props.push_back(i);
//     }

    // Benchmark default Allgather on matrix_row_comm. [VC,* ] <- [MC,* ]
    local_size = (m+(p-1))/p*b;
    gathered_size = local_size*c;
    Mb = 64.*gathered_size/1.e6;
    local_data.resize( local_size );
    gathered_data.resize( gathered_size );
    MPI_Barrier( matrix_row_comm );
    start_time = MPI_Wtime();
    for( unsigned i=0; i<num_tests; ++i )
    {
        MPI_Allgather
        ( &local_data[0], local_size, MPI_DOUBLE,
          &gathered_data[0], local_size, MPI_DOUBLE, matrix_row_comm );
    }
    stop_time = MPI_Wtime();
    average_time = (stop_time-start_time)/num_tests;
    bandwidth = 64.*gathered_size/(1.e6*average_time);
    if( torus_rank == 0 )
    {
        ostringstream msg;
        msg << "matrix_row Allgather: "
            << average_time << " seconds, "
            << Mb << " Megabits,"
            << bandwidth << " Mbps." << endl;
        cout << msg.str() << endl;
    }

//     // Force only one property at a time to be nonzero
//     for( unsigned i=0; i<avail_props.size(); ++i )
//     {
//         for( unsigned j=0; j<allgather_codes.size(); ++j )
//         {
//             if( avail_props[i] == j )
//                 MPIX_Set_property( matrix_row_comm, allgather_codes[j], 1 );
//             else
//                 MPIX_Set_property( matrix_row_comm, allgather_codes[j], 0 );
//         }

//         MPI_Barrier( torus_comm );
//         start_time = MPI_Wtime();
//         for( unsigned i=0; i<num_tests; ++i )
//         {
//             MPI_Allgather
//             ( &local_data[0], local_size, MPI_DOUBLE,
//               &gathered_data[0], local_size, MPI_DOUBLE, matrix_row_comm );
//         }
//         stop_time = MPI_Wtime();
//         average_time = (stop_time-start_time)/num_tests;
//         bandwidth = 64.*gathered_size/(1.e6*average_time);
//         if( torus_rank == 0 )
//         {
//             ostringstream msg;
//             msg << "matrix_row Allgather (property " 
//                 << allgather_codes[avail_props[i]] << "): "
//                 << average_time << " seconds, "
//                 << Mb << " Megabits, "
//                 << bandwidth << " Mbps." << endl;
//             cout << msg.str() << endl;
//         }
//     }

//     // Force the MPIDO_USE_RECT_BCAST flag to 1
//     for( unsigned j=0; j<allgather_codes.size(); ++j )
//         MPIX_Set_property( matrix_row_comm, allgather_codes[j], 0 );
//     MPIX_Set_property( matrix_row_comm, 34, 1 );
//     MPI_Barrier( torus_comm );
//     start_time = MPI_Wtime();
//     for( unsigned i=0; i<num_tests; ++i )
//     {
//         MPI_Allgather
//         ( &local_data[0], local_size, MPI_DOUBLE,
//           &gathered_data[0], local_size, MPI_DOUBLE, matrix_row_comm );
//     }
//     stop_time = MPI_Wtime();
//     average_time = (stop_time-start_time)/num_tests;
//     bandwidth = 64.*gathered_size/(1.e6*average_time);
//     if( torus_rank == 0 )
//     {
//         ostringstream msg;
//         msg << "matrix_row Allgather (property 34): " 
//             << average_time << " seconds, "
//             << Mb << " Megabits, "
//             << bandwidth << " Mbps." << endl;
//         cout << msg.str() << endl;
//     }

    MPI_Finalize();
    return 0;
}