/*
//@HEADER
// ************************************************************************
//
//                        Kokkos v. 2.0
//              Copyright (2014) Sandia Corporation
//
// Under the terms of Contract DE-AC04-94AL85000 with Sandia Corporation,
// the U.S. Government retains certain rights in this software.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
// 1. Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright
// notice, this list of conditions and the following disclaimer in the
// documentation and/or other materials provided with the distribution.
//
// 3. Neither the name of the Corporation nor the names of the
// contributors may be used to endorse or promote products derived from
// this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY SANDIA CORPORATION "AS IS" AND ANY
// EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL SANDIA CORPORATION OR THE
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
// LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
//
// Questions? Contact Christian R. Trott (crtrott@sandia.gov)
//
// ************************************************************************
//@HEADER
*/

#include <Kokkos_Core.hpp>

#if ! defined( KOKKOS_ENABLE_TASKDAG ) || \
    defined( KOKKOS_ENABLE_DEFAULT_DEVICE_TYPE_THREADS )

int main()
{
  return 0 ;
}

#else

#include <cstdio>
#include <cstring>
#include <cstdlib>
#include <limits>

#include <impl/Kokkos_Timer.hpp>

using ExecSpace = Kokkos::DefaultExecutionSpace ;

inline
long eval_fib( long n )
{
  constexpr long mask = 0x03;

  long fib[4] = { 0, 1, 0, 0 };

  for ( long i = 2; i <= n; ++i ) {
    fib[ i & mask ] = fib[ ( i - 1 ) & mask ] + fib[ ( i - 2 ) & mask ];
  }

  return fib[ n & mask ];
}

inline
long fib_alloc_count( long n )
{
  constexpr long mask = 0x03;

  long count[4] = { 1, 1, 0, 0 };

  for ( long i = 2; i <= n; ++i ) {
    count[ i & mask ] = 2 // this task plus the 'when_all' task
                      + count[ ( i - 1 ) & mask ]
                      + count[ ( i - 2 ) & mask ];
  }

  return count[ n & mask ];
}

template< class Space >
struct TestFib {

  using Scheduler   = Kokkos::TaskScheduler< Space > ;
  using MemorySpace = typename Scheduler::memory_space ;
  using MemberType  = typename Scheduler::member_type ;
  using FutureType  = Kokkos::Future< long , Space > ;

  typedef long value_type ;

  Scheduler  sched ;
  FutureType dep[2] ;
  const value_type n ;

  KOKKOS_INLINE_FUNCTION
  TestFib( const Scheduler & arg_sched , const value_type arg_n )
    : sched( arg_sched ), dep{} , n( arg_n ) {}

  KOKKOS_INLINE_FUNCTION
  void operator()( const MemberType & , value_type & result ) noexcept
    {
      if ( n < 2 ) {
        result = n ;
      }
      else if ( ! dep[0].is_null() && ! dep[1].is_null() ) {
        result = dep[0].get() + dep[1].get();
      }
      else {
        // Spawn new children and respawn myself to sum their results.
        // Spawn lower value at higher priority as it has a shorter
        // path to completion.

        dep[1] = Kokkos::task_spawn
          ( Kokkos::TaskSingle( sched, Kokkos::TaskPriority::High )
          , TestFib( sched, n - 2 ) );

        dep[0] = Kokkos::task_spawn
          ( Kokkos::TaskSingle( sched )
          , TestFib( sched, n - 1 ) );

        Kokkos::Future< ExecSpace > fib_all = Kokkos::when_all( dep, 2 );

        if ( ! dep[0].is_null() && ! dep[1].is_null() && ! fib_all.is_null() ) {
          // High priority to retire this branch.
          Kokkos::respawn( this, fib_all, Kokkos::TaskPriority::High );
        }
        else {
          Kokkos::abort("Failed nested task spawn (allocation)");
        }
      }
    }
};



int main( int argc , char* argv[] )
{
  static const char help[] = "--help" ;
  static const char alloc_size[]   = "--alloc_size=" ;
  static const char super_size[]   = "--super_size=" ;
  static const char repeat_outer[] = "--repeat_outer=" ;
  static const char input_value[]  = "--input=" ;

  long total_alloc_size    = 1000000 ;
  int  min_superblock_size =   10000 ;
  int  test_repeat_outer   =       1 ;
  int  fib_input           =       4 ;

  int  ask_help = 0 ;

  for(int i=1;i<argc;i++)
  {
     const char * const a = argv[i];

     if ( ! strncmp(a,help,strlen(help) ) ) ask_help = 1 ;

     if ( ! strncmp(a,alloc_size,strlen(alloc_size) ) )
       total_alloc_size = atol( a + strlen(alloc_size) );

     if ( ! strncmp(a,super_size,strlen(super_size) ) )
       min_superblock_size = atoi( a + strlen(super_size) );

     if ( ! strncmp(a,repeat_outer,strlen(repeat_outer) ) )
       test_repeat_outer = atoi( a + strlen(repeat_outer) );

     if ( ! strncmp(a,input_value,strlen(input_value) ) )
       fib_input = atoi( a + strlen(input_value) );
  }

  const long fib_output   = eval_fib( fib_input );
  const long number_alloc = fib_alloc_count( fib_input );

  const unsigned min_block_size =  32 ;
  const unsigned max_block_size = 128 ;

  long task_count_max = 0 ;
  long task_count_accum = 0 ;
  long test_result = 0 ;

  if ( ask_help ) {
    std::cout << "command line options:"
              << " " << help
              << " " << alloc_size << "##"
              << " " << super_size << "##"
              << " " << input_value << "##"
              << " " << repeat_outer << "##"
              << std::endl ;
    return -1;
  }

  typedef TestFib< ExecSpace >  Functor ;

  Kokkos::initialize(argc,argv);

  Functor::Scheduler sched( Functor::MemorySpace()
                          , total_alloc_size
                          , min_block_size
                          , max_block_size
                          , min_superblock_size
                          );

  Functor::FutureType f =
    Kokkos::host_spawn( Kokkos::TaskSingle( sched )
                      , Functor( sched , fib_input )
                      );

  Kokkos::wait( sched );

  test_result = f.get();

  task_count_max   = sched.allocated_task_count_max();
  task_count_accum = sched.allocated_task_count_accum();

  if ( number_alloc != task_count_accum ) {
    std::cout << " number_alloc( " << number_alloc << " )"
              << " != task_count_accum( " << task_count_accum << " )"
              << std::endl ;
  }

  if ( fib_output != test_result ) {
    std::cout << " answer( " << fib_output << " )"
              << " != result( " << test_result << " )"
              << std::endl ;
  }

  if ( fib_output != test_result || number_alloc != task_count_accum ) {
    printf("  TEST FAILED\n");
    return -1;
  }

  double min_time = std::numeric_limits<double>::max();
  double time_sum = 0;

  for ( int i = 0 ; i < test_repeat_outer ; ++i ) {
    Kokkos::Impl::Timer timer ;

    Functor::FutureType ftmp =
      Kokkos::host_spawn( Kokkos::TaskSingle( sched )
                        , Functor( sched , fib_input )
                        );

    Kokkos::wait( sched );
    auto this_time = timer.seconds();
    min_time = std::min(min_time, this_time);
    time_sum += this_time;
  }

  auto avg_time = time_sum / test_repeat_outer;

  Kokkos::finalize();

  printf( "\"taskdag: alloc super repeat input output task-accum task-max\" %ld %d %d %d %ld %ld %ld\n"
        , total_alloc_size
        , min_superblock_size
        , test_repeat_outer
        , fib_input
        , fib_output
        , task_count_accum
        , task_count_max );

  printf( "\"taskdag: time (min, avg)\" %g %g\n", min_time, avg_time);
  printf( "\"taskdag: tasks per second (max, avg)\" %g %g\n"
        , number_alloc / min_time
        , number_alloc / avg_time );

  return 0 ;
}

#endif

