// SPDX-FileComment: This file is part of TNL - Template Numerical Library (https://tnl-project.org/)
// SPDX-License-Identifier: MIT

#pragma once

#include <TNL/Devices/Sequential.h>
#include <TNL/Devices/Host.h>
#include <TNL/Devices/Cuda.h>
#include <TNL/Backend/Macros.h>
#include <TNL/Algorithms/reduce.h>

namespace TNL::Algorithms::detail {

template< typename Device >
struct Contains;

template< typename Device >
struct ContainsOnlyValue;

template<>
struct Contains< Devices::Sequential >
{
   template< typename Element, typename Index >
   __cuda_callable__
   bool
   operator()( const Element* data, const Index size, const Element& value )
   {
      if( data == nullptr || size == 0 )
         return false;

      for( Index i = 0; i < size; i++ )
         if( data[ i ] == value )
            return true;
      return false;
   }
};

template<>
struct ContainsOnlyValue< Devices::Sequential >
{
   template< typename Element, typename Index >
   __cuda_callable__
   bool
   operator()( const Element* data, const Index size, const Element& value )
   {
      if( data == nullptr || size == 0 )
         return true;

      for( Index i = 0; i < size; i++ )
         if( ! ( data[ i ] == value ) )
            return false;
      return true;
   }
};

template<>
struct Contains< Devices::Host >
{
   template< typename Element, typename Index >
   bool
   operator()( const Element* data, const Index size, const Element& value )
   {
      if( data == nullptr || size == 0 )
         return false;

      if( Devices::Host::isOMPEnabled() && Devices::Host::getMaxThreadsCount() > 1 ) {
         auto fetch = [ = ]( Index i ) -> bool
         {
            return data[ i ] == value;
         };
         return reduce< Devices::Host >( static_cast< Index >( 0 ), size, fetch, std::logical_or<>{}, false );
      }
      else {
         // sequential algorithm can return as soon as it finds a match
         return Contains< Devices::Sequential >{}( data, size, value );
      }
   }
};

template<>
struct ContainsOnlyValue< Devices::Host >
{
   template< typename Element, typename Index >
   bool
   operator()( const Element* data, const Index size, const Element& value )
   {
      if( data == nullptr || size == 0 )
         return true;

      if( Devices::Host::isOMPEnabled() && Devices::Host::getMaxThreadsCount() > 1 ) {
         auto fetch = [ data, value ]( Index i ) -> bool
         {
            return data[ i ] == value;
         };
         return reduce< Devices::Host >( static_cast< Index >( 0 ), size, fetch, std::logical_and<>{}, true );
      }
      else {
         // sequential algorithm can return as soon as it finds a mismatch
         return ContainsOnlyValue< Devices::Sequential >{}( data, size, value );
      }
   }
};

template<>
struct Contains< Devices::Cuda >
{
   template< typename Element, typename Index >
   bool
   operator()( const Element* data, const Index size, const Element& value )
   {
      if( data == nullptr || size == 0 )
         return false;

      auto fetch = [ = ] __cuda_callable__( Index i ) -> bool
      {
         return data[ i ] == value;
      };
      return reduce< Devices::Cuda >( static_cast< Index >( 0 ), size, fetch, std::logical_or<>{}, false );
   }
};

template<>
struct ContainsOnlyValue< Devices::Cuda >
{
   template< typename Element, typename Index >
   bool
   operator()( const Element* data, const Index size, const Element& value )
   {
      if( data == nullptr || size == 0 )
         return true;

      auto fetch = [ = ] __cuda_callable__( Index i ) -> bool
      {
         return data[ i ] == value;
      };
      return reduce< Devices::Cuda >( static_cast< Index >( 0 ), size, fetch, std::logical_and<>{}, true );
   }
};

}  // namespace TNL::Algorithms::detail
