#include "BufferExchange.h"

template <class T> int BufferExchange<T>::_wait( std::vector<MPI_Request> &reqs )
{
	int result, index;

	START_ROUTINE_TIMER

	result = MPI_Waitany( reqs.size(), &reqs[0], &index, MPI_STATUS_IGNORE );
	if( result != MPI_SUCCESS ) ERROR_MACRO( "MPI_Waitany()\n" );

	END_ROUTINE_TIMER

	if( index == MPI_UNDEFINED ) return -1;
	else return index;
}
template <class T> void BufferExchange<T>::_waitall( std::vector<MPI_Request> &reqs )
{
	int result;
	
	START_ROUTINE_TIMER

	if( 1 )
	{
		// default - use MPI_Waitall()
		result = MPI_Waitall( reqs.size(), &reqs[0], MPI_STATUSES_IGNORE );
		if( result != MPI_SUCCESS ) ERROR_MACRO( "MPI_Waitall()" );
	}
	else
	{
		// print time intervals between each receive:
		double s, s2, e;
		
		Log( "%s:\n", __func__ );
		s = GetTime();
		s2 = s;
		while( (result = _wait(reqs) ) != -1 )
		{
			e = GetTime();
			Log( "\t peer index %d : %g, %g\n", result, e-s2, e-s );
			s2 = e;
		}
		e = GetTime();
		Log( "\t total: %g\n", result, e-s );
	}

	END_ROUTINE_TIMER
}
template <class T> BufferExchange<T>::BufferExchange()
{
	full_exchange_length_tag = 1;
	full_exchange_buffer_tag = 2;
	exchange_tag = 3;
}
template <class T> BufferExchange<T>::~BufferExchange() {}
template <class T> void BufferExchange<T>::Clear( IntVec &_ranks )
{
	size_t N;
	
	ranks = _ranks;
	N = ranks.size();
	
	recv_reqs.resize( N );
	send_reqs.resize( N );
	
	recv_buffers.resize( N );
	send_buffers.resize( N );
	
	for( size_t i=0; i<N; i++ )
	{
		recv_reqs[i] = MPI_REQUEST_NULL;
		send_reqs[i] = MPI_REQUEST_NULL;
		
		recv_buffers[i].clear();
		send_buffers[i].clear();
	}
}
template <class T> void BufferExchange<T>::FullExchange()
{
	const int divisor = 1024; // (1024*1024);
	size_t N;
	int result, index;
	unsigned int vs;
	unsigned int rss0, rss1, rss2, rss3, rss4;
	
	START_ROUTINE_TIMER
	
	//
	// Repurpose existing MPI_Request arrays, rename for clarity.
	//
	std::vector<MPI_Request> &length_reqs = recv_reqs;
	std::vector<MPI_Request> &buffer_reqs = send_reqs;
	
	N = ranks.size();
	
	getmem( &rss0, &vs );

	//
	// Incoming/outgoing buffer lengths IN UNITS OF T!
	//
	recv_lengths.resize( N );
	send_lengths.resize( N );

	//
	// Single flat array of handles for both incoming and outgoing lengths, and incoming/outgoing buffers (hence N*2).
	//
	length_reqs.resize( N*2 );
	buffer_reqs.resize( N*2 );

	for( size_t i=0; i<2*N; i++ )
	{
		length_reqs[i] = MPI_REQUEST_NULL;
		buffer_reqs[i] = MPI_REQUEST_NULL;
	}

		
	getmem( &rss1, &vs );

	//
	// Send & receive buffer lengths - order important, as we assume 0->(N-1) are receives, N->(N*2)-1 are sends.
	//
	for( size_t i=0; i<N; i++ )
	{
		result = MPI_Irecv( &recv_lengths[i], 1, MPI_INT, ranks[i], full_exchange_length_tag, MPI_COMM_WORLD, &length_reqs[0+i] );
		if( result != MPI_SUCCESS ) ERROR_MACRO( "MPI_Irecv" );
	}
	for( size_t i=0; i<N; i++ )
	{
		send_lengths[i] = (int)send_buffers[i].size();
		
		result = MPI_Isend( &send_lengths[i], 1, MPI_INT, ranks[i], full_exchange_length_tag, MPI_COMM_WORLD, &length_reqs[N+i] );
		if( result != MPI_SUCCESS ) ERROR_MACRO( "MPI_Irecv" );
	}

	getmem( &rss2, &vs );

	//
	// Wait on sends/receives of buffer lengths, then send/recv buffer if length nonzero.
	//
	int n_buffer_reqs = 0;
	while( 1 )
	{
		result = MPI_Waitany( N*2, &length_reqs[0], &index, MPI_STATUSES_IGNORE );
		if( result != MPI_SUCCESS ) ERROR_MACRO( "MPI_Waitany" );
		
		if( index == MPI_UNDEFINED ) break; // no more active requests - stop here.
		
		if( index < (int)N )
		{
			//
			// Length received; if length > 0, receive actual data.
			//
			int i = index;
			if( recv_lengths[i] > 0 )
			{
				recv_buffers[i].resize( recv_lengths[i] );
				
				result = MPI_Irecv( (char *)&recv_buffers[i][0], sizeof(T)*recv_lengths[i], MPI_CHAR, ranks[i], full_exchange_buffer_tag, MPI_COMM_WORLD, &buffer_reqs[n_buffer_reqs] );
				if( result != MPI_SUCCESS ) ERROR_MACRO( "MPI_Irecv" );
				
				n_buffer_reqs++;
			}
			else
			{
				// prevent stale previous data in recv_buffers[i]: ensure new size() returns zero where received a zero length!
				recv_buffers[i].clear();
			}
		}
		else
		{
			//
			// Length sent; if length > 0, send actual data.
			//
			int i = index-N;
			if( send_lengths[i] > 0 )
			{
				result = MPI_Isend( (char *)&send_buffers[i][0], sizeof(T)*send_lengths[i], MPI_CHAR, ranks[i], full_exchange_buffer_tag, MPI_COMM_WORLD, &buffer_reqs[n_buffer_reqs] );
				if( result != MPI_SUCCESS ) ERROR_MACRO( "MPI_Isend" );
				
				n_buffer_reqs++;
			}
		}
	}
	
	getmem( &rss3, &vs );

	//
	// Wait individual
	//
	/*
	while( 1 )
	{
		int index;
		
		result = MPI_Waitany( n_buffer_reqs, &buffer_reqs[0], &index, MPI_STATUSES_IGNORE );
		if( result != MPI_SUCCESS ) ERROR_MACRO( "MPI_Waitany" );
		
		if( index == MPI_UNDEFINED ) break;
	}
	*/
	
	//
	// Wait all
	//
	result = MPI_Waitall( n_buffer_reqs, &buffer_reqs[0], MPI_STATUSES_IGNORE );
	if( result != MPI_SUCCESS ) ERROR_MACRO( "MPI_Waitall" );

	getmem( &rss4, &vs );

	Log( "%s() memory deltas:\n", __func__ );
	Log( "\t initial setup: %f\n", (double)(rss1-rss0)/divisor );
	Log( "\t send/recv lengths: %f\n", (double)(rss2-rss1)/divisor );
	Log( "\t send/recv buffers: %f\n", (double)(rss3-rss2)/divisor );
	Log( "\t final waitall: %f\n", (double)(rss4-rss3)/divisor );

	END_ROUTINE_TIMER
}
template <class T> void BufferExchange<T>::PostRecvs()
{
	int result;
	
	START_ROUTINE_TIMER

	for( size_t i=0; i<ranks.size(); i++ )
	{
		size_t len = recv_buffers[i].size()*sizeof(T);
	
		recv_reqs[i] = MPI_REQUEST_NULL;
		if( len > 0 )
		{
			result = MPI_Irecv( (char *)&recv_buffers[i][0], len, MPI_CHAR, ranks[i], exchange_tag, MPI_COMM_WORLD, &recv_reqs[i] );
			if( result != MPI_SUCCESS ) ERROR_MACRO( "MPI_Irecv( %d bytes ) to rank %d (%d)\n", (int)len, ranks[i], (int)i );
		}
	}

	END_ROUTINE_TIMER
}
template <class T> void BufferExchange<T>::PostSends()
{
	int result;

	START_ROUTINE_TIMER

	for( size_t i=0; i<ranks.size(); i++ )
	{
		size_t len = send_buffers[i].size()*sizeof(T);
	
		send_reqs[i] = MPI_REQUEST_NULL;
		if( len > 0 )
		{
			result = MPI_Isend( (char *)&send_buffers[i][0], len, MPI_CHAR, ranks[i], exchange_tag, MPI_COMM_WORLD, &send_reqs[i] );
			if( result != MPI_SUCCESS ) ERROR_MACRO( "MPI_Isend( %d bytes ) to rank %d (%d)\n", (int)len, ranks[i], (int)i );
		}
	}

	END_ROUTINE_TIMER
}
template <class T> int BufferExchange<T>::WaitRecv() { return _wait( recv_reqs ); }
template <class T> int BufferExchange<T>::WaitSend() { return _wait( send_reqs ); }
template <class T> void BufferExchange<T>::WaitRecvs() { _waitall( recv_reqs ); }
template <class T> void BufferExchange<T>::WaitSends() { _waitall( send_reqs ); }
template <class T> void BufferExchange<T>::CancelRecvs()
{
	int result;

	START_ROUTINE_TIMER

	for( size_t i=0; i<recv_reqs.size(); i++ )
	{
		if( recv_reqs[i] != MPI_REQUEST_NULL )
		{
			result = MPI_Cancel( &recv_reqs[i] );
			if( result != MPI_SUCCESS ) ERROR_MACRO( "MPI_Cancel( recv_reqs[%d] )", (int)i );
		}
	}
	MPI_Waitall( recv_reqs.size(), &recv_reqs[0], MPI_STATUSES_IGNORE ); // necessary to free resources

	END_ROUTINE_TIMER
}


//
// NEED TO DECLARE EXPLICIT TYPE-INSTANTIATED TEMPLATE CLASSES HERE!
// We can avoid this by putting the implementation in the header, as is more conventional
// for template classes, but I prefer the simpler header and we know what types we'll use
// in BufferExchange objects in advance.
//

template class BufferExchange<int>;