#include "Time.h"

///////////////////////////
// MultiSocketBackedGrid //
///////////////////////////
template< int Channels , class Real >
MultiSocketBackedGrid< Channels , Real >::MultiSocketBackedGrid( Socket* socks , int* rowSizes , int sockCount , int rows , bool blockingSend , bool separateColors )
{
	_read = true;
	_separateColors = separateColors;
	_blockingSend = blockingSend;
	_sockCount = sockCount;
	_rows = rows;
	_rowSize = 0;
	for( int i=0 ; i<_sockCount ; i++ ) _rowSize += rowSizes[i];

	_rowSizes = ( int* )malloc( sizeof(int) * _sockCount );
	_socks = ( Socket* )malloc( sizeof(Socket) * _sockCount );
	_subData = AllocArray< Pointer( Real ) >( _sockCount , "MultiSocketBackedGrid::subData" );
	if( !_rowSizes || !_socks || !_subData ) fprintf( stderr , "Failed to allocate\n" ) ,  exit(0);
	for( int i=0 ; i<_sockCount ; i++ )
	{
		_rowSizes[i] = rowSizes[i];
		_socks[i] = socks[i];
		_subData[i] = AllocArray< Real >( Channels * rowSizes[i] , "MultiSocketBackedGrid::subData[]" );
		if( !_subData[i] ) fprintf( stderr , "Failed to allocate\n" ) , exit(0);
	}

	_data = AllocArray< Real >( _rowSize * Channels , "MultiSocketBackedGrid::data" );
	if( !_data ) fprintf( stderr , "Failed to allocate\n" ) , exit(0);
}
template< int Channels , class Real >
MultiSocketBackedGrid< Channels , Real >::~MultiSocketBackedGrid(void)
{
	for( int i=0 ; i<_sockCount ; i++ ) FreeArray( _subData[i] );
	FreeArray( _data );
	FreeArray( _subData );
	free( _socks );
	free( _rowSizes );

	_rows = _rowSize = 0;
	_rowSizes = NULL;
	_socks = NULL;
}
template< int Channels , class Real >
int	MultiSocketBackedGrid< Channels , Real >::rows( void ) const { return _rows; }
template< int Channels , class Real >
int MultiSocketBackedGrid< Channels , Real >::rowSize( void ) const { return _rowSize * Channels * sizeof(Real); }
template< int Channels , class Real >
Pointer( byte ) MultiSocketBackedGrid< Channels , Real >::operator[]( int idx )
{
	if( _read && !_readComplete )
	{
		int offset = 0;
		for( int i=0 ; i<_sockCount ; i++ )
		{
			ReceiveOnSocket( _socks[i] , _subData[i] , _rowSizes[i]*Channels*sizeof(Real) , _blockingSend , "MultiSocketBackedGrid<Channels,Real>::[]" );
			if( _separateColors ) for( int c=0 ; c<Channels ; c++ ) memcpy( _data + offset + c*_rowSize , _subData[i]+c*_rowSizes[i] , sizeof(Real)*_rowSizes[i] );
			else memcpy( _data+offset*Channels , _subData[i] , _rowSizes[i]*Channels*sizeof(Real) );
			offset += _rowSizes[i];
		}
		_readComplete = true;
	}
#if ASSERT_MEMORY_ACCESS
	if( idx<0 || idx>=_rows || idx!=_current) fprintf( stderr , "[MultiSocketBackedGrid] Index out of bounds: %d != %d || [ %d , %d )\n" , idx , _current , 0 , _rows ) , exit(0);
#endif // ASSERT_MEMORY_ACCESS
	return ( Pointer( byte ) )_data;
}
template< int Channels , class Real >
void MultiSocketBackedGrid< Channels , Real >::advance( void )
{
	if( !_read )
	{
		int offset = 0;
		for( int i=0 ; i<_sockCount ; i++ )
		{
			if( _separateColors ) for( int c=0 ; c<Channels ; c++ ) memcpy( _subData[i]+c*_rowSizes[i] , _data + offset + c*_rowSize , sizeof(Real)*_rowSizes[i] );
			else memcpy( _subData[i] , _data+offset*Channels , _rowSizes[i]*Channels*sizeof(Real) );
			SendOnSocket( _socks[i] , _subData[i] , _rowSizes[i]*Channels*sizeof(Real) , _blockingSend  , "MultiSocketBackedGrid< Channels,Real>::advance" );
			offset += _rowSizes[i];
		}
	}
	else
	{
		if( !_readComplete )
		{
			int offset = 0;
			for( int i=0 ; i<_sockCount ; i++ )
			{
				if( _separateColors ) for( int c=0 ; c<Channels ; c++ ) memcpy( _subData[i]+c*_rowSizes[i] , _data + offset + c*_rowSize , sizeof(Real)*_rowSizes[i] );
				else memcpy( _subData[i] , _data+offset*Channels , _rowSizes[i]*Channels*sizeof(Real) );
				ReceiveOnSocket( _socks[i] , _subData[i] , _rowSizes[i]*Channels*sizeof(Real) , _blockingSend , "MultiSocketBackedGrid<Channels,Real>::advance"  );
				offset += _rowSizes[i];
			}
		}
		_readComplete = false;
	}
	_current++;
}
template< int Channels , class Real >
void MultiSocketBackedGrid< Channels , Real >::reset( bool read , int minWindowSize )
{
	if(!_read ) for( int i=0 ; i<_sockCount ; i++ ) EndSendOnSocket( _socks[i] , _blockingSend , "MultiSocketBackedGrid<Channels,Real>::reset" );
	_read = read;
	_readComplete = !_read;
	_current = 0;
	if( _read ) for( int i=0 ; i<_sockCount ; i++ ) StartReceiveOnSocket( _socks[i] , _blockingSend , "MultiSocketBackedGrid<Channels,Real >::reset" );
}

