#include <TNL/Experimental/Multimaps/EllpackIndexMultimap.h>

using namespace TNL;

using IndexType = int;
using Device = Devices::Host;
using LocalIndexType = short;

#ifdef HAVE_GTEST 
#include "gtest/gtest.h"

TEST( MultimapTest, TestTypedefs )
{
   using MultimapType = TNL::EllpackIndexMultimap< IndexType, Device, LocalIndexType >;
   const bool same_index = std::is_same< typename MultimapType::IndexType, IndexType >::value;
   ASSERT_TRUE( same_index );
   const bool same_device = std::is_same< typename MultimapType::DeviceType, Device >::value;
   ASSERT_TRUE( same_device );
   const bool same_localindex = std::is_same< typename MultimapType::LocalIndexType, LocalIndexType >::value;
   ASSERT_TRUE( same_localindex );
}

TEST( MultimapTest, TestSettingSizes )
{
   using MultimapType = TNL::EllpackIndexMultimap< IndexType, Device, LocalIndexType >;

   IndexType inputs = 10;
   LocalIndexType valuesGlobalMax = 3;
   LocalIndexType valuesLocalMax = 2;

   MultimapType map;
   map.setKeysRange( inputs );
   ASSERT_EQ( map.getKeysRange(), inputs );

   typename MultimapType::ValuesAllocationVectorType allocationRanges;
   ASSERT_TRUE( allocationRanges.setSize( inputs ) );
   allocationRanges.setValue( valuesGlobalMax );
   allocationRanges[ 0 ] = 0;
   allocationRanges[ 1 ] = 1;
   ASSERT_TRUE( map.allocate( allocationRanges ) );

   for( IndexType i = 0; i < inputs; i++ ) {
      auto values = map.getValues( i );
      const auto constValues = ( (const MultimapType) map ).getValues( i );

      // uninitialized should be equal to the value from the allocation vector
      ASSERT_EQ( values.getSize(), allocationRanges[ i ] );
      ASSERT_EQ( constValues.getSize(), allocationRanges[ i ] );

      // setting lower sizes
      ASSERT_TRUE( values.setSize( valuesLocalMax ) );
      ASSERT_EQ( values.getSize(), valuesLocalMax );
      ASSERT_EQ( constValues.getSize(), valuesLocalMax );

      // setting wrong local sizes should not be allowed
      ASSERT_FALSE( values.setSize( valuesGlobalMax + 1 ) );
      ASSERT_EQ( values.getSize(), valuesLocalMax );
      ASSERT_EQ( constValues.getSize(), valuesLocalMax );

      // setting global max should succeed
      ASSERT_TRUE( values.setSize( valuesGlobalMax ) );
      ASSERT_EQ( values.getSize(), valuesGlobalMax );
      ASSERT_EQ( constValues.getSize(), valuesGlobalMax );
   }
}

TEST( MultimapTest, TestSettingValues )
{
   using MultimapType = TNL::EllpackIndexMultimap< IndexType, Device, LocalIndexType >;

   IndexType inputs = 10;
   LocalIndexType allocatedValues = 2;

   MultimapType map;
   map.setKeysRange( inputs );
   ASSERT_EQ( map.getKeysRange(), inputs );

   typename MultimapType::ValuesAllocationVectorType allocationRanges;
   ASSERT_TRUE( allocationRanges.setSize( inputs ) );
   allocationRanges.setValue( allocatedValues );
   ASSERT_TRUE( map.allocate( allocationRanges ) );

   for( IndexType i = 0; i < inputs; i++ ) {
      auto values = map.getValues( i );
      const auto constValues = ( (const MultimapType) map ).getValues( i );

      ASSERT_TRUE( values.setSize( allocatedValues ) );

      for( LocalIndexType o = 0; o < allocatedValues; o++ )
         values.setValue( o, i + o );

      for( LocalIndexType o = 0; o < allocatedValues; o++ ) {
         ASSERT_EQ( values.getValue( o ), i + o );
         ASSERT_EQ( values[ o ], i + o );
         ASSERT_EQ( constValues.getValue( o ), i + o );
         ASSERT_EQ( constValues[ o ], i + o );
      }

      for( LocalIndexType o = 0; o < allocatedValues; o++ )
         values[ o ] = i * o;

      for( LocalIndexType o = 0; o < allocatedValues; o++ ) {
         ASSERT_EQ( values.getValue( o ), i * o );
         ASSERT_EQ( values[ o ], i * o );
         ASSERT_EQ( constValues.getValue( o ), i * o );
         ASSERT_EQ( constValues[ o ], i * o );
      }
   }
}
#endif

int main( int argc, char* argv[] )
{
#ifdef HAVE_GTEST
   ::testing::InitGoogleTest( &argc, argv );
   return RUN_ALL_TESTS();
#else
   return EXIT_FAILURE;
#endif
}
