/*
//@HEADER
// ************************************************************************
//
//                        Kokkos v. 3.0
//       Copyright (2020) National Technology & Engineering
//               Solutions of Sandia, LLC (NTESS).
//
// Under the terms of Contract DE-NA0003525 with NTESS,
// 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 NTESS "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 NTESS 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>

namespace TestCXX11 {

template <class DeviceType>
struct FunctorAddTest {
  typedef Kokkos::View<double**, DeviceType> view_type;
  typedef DeviceType execution_space;
  typedef typename Kokkos::TeamPolicy<execution_space>::member_type team_member;

  view_type a_, b_;

  FunctorAddTest(view_type& a, view_type& b) : a_(a), b_(b) {}

  KOKKOS_INLINE_FUNCTION
  void operator()(const int& i) const {
    b_(i, 0) = a_(i, 1) + a_(i, 2);
    b_(i, 1) = a_(i, 0) - a_(i, 3);
    b_(i, 2) = a_(i, 4) + a_(i, 0);
    b_(i, 3) = a_(i, 2) - a_(i, 1);
    b_(i, 4) = a_(i, 3) + a_(i, 4);
  }

  KOKKOS_INLINE_FUNCTION
  void operator()(const team_member& dev) const {
    const int begin = dev.league_rank() * 4;
    const int end   = begin + 4;
    for (int i = begin + dev.team_rank(); i < end; i += dev.team_size()) {
      b_(i, 0) = a_(i, 1) + a_(i, 2);
      b_(i, 1) = a_(i, 0) - a_(i, 3);
      b_(i, 2) = a_(i, 4) + a_(i, 0);
      b_(i, 3) = a_(i, 2) - a_(i, 1);
      b_(i, 4) = a_(i, 3) + a_(i, 4);
    }
  }
};

template <class DeviceType, bool PWRTest>
double AddTestFunctor() {
  typedef Kokkos::TeamPolicy<DeviceType> policy_type;

  Kokkos::View<double**, DeviceType> a("A", 100, 5);
  Kokkos::View<double**, DeviceType> b("B", 100, 5);
  typename Kokkos::View<double**, DeviceType>::HostMirror h_a =
      Kokkos::create_mirror_view(a);
  typename Kokkos::View<double**, DeviceType>::HostMirror h_b =
      Kokkos::create_mirror_view(b);

  for (int i = 0; i < 100; i++) {
    for (int j = 0; j < 5; j++) {
      h_a(i, j) = 0.1 * i / (1.1 * j + 1.0) + 0.5 * j;
    }
  }
  Kokkos::deep_copy(a, h_a);

  if (PWRTest == false) {
    Kokkos::parallel_for(100, FunctorAddTest<DeviceType>(a, b));
  } else {
    Kokkos::parallel_for(policy_type(25, Kokkos::AUTO),
                         FunctorAddTest<DeviceType>(a, b));
  }
  Kokkos::deep_copy(h_b, b);

  double result = 0;
  for (int i = 0; i < 100; i++) {
    for (int j = 0; j < 5; j++) {
      result += h_b(i, j);
    }
  }

  return result;
}

#if defined(KOKKOS_ENABLE_CXX11_DISPATCH_LAMBDA)
template <class DeviceType, bool PWRTest>
double AddTestLambda() {
  Kokkos::View<double**, DeviceType> a("A", 100, 5);
  Kokkos::View<double**, DeviceType> b("B", 100, 5);
  typename Kokkos::View<double**, DeviceType>::HostMirror h_a =
      Kokkos::create_mirror_view(a);
  typename Kokkos::View<double**, DeviceType>::HostMirror h_b =
      Kokkos::create_mirror_view(b);

  for (int i = 0; i < 100; i++) {
    for (int j = 0; j < 5; j++) {
      h_a(i, j) = 0.1 * i / (1.1 * j + 1.0) + 0.5 * j;
    }
  }
  Kokkos::deep_copy(a, h_a);

  if (PWRTest == false) {
    Kokkos::parallel_for(
        100, KOKKOS_LAMBDA(const int& i) {
          b(i, 0) = a(i, 1) + a(i, 2);
          b(i, 1) = a(i, 0) - a(i, 3);
          b(i, 2) = a(i, 4) + a(i, 0);
          b(i, 3) = a(i, 2) - a(i, 1);
          b(i, 4) = a(i, 3) + a(i, 4);
        });
  } else {
    typedef Kokkos::TeamPolicy<DeviceType> policy_type;
    typedef typename policy_type::member_type team_member;

    policy_type policy(25, Kokkos::AUTO);

    Kokkos::parallel_for(
        policy, KOKKOS_LAMBDA(const team_member& dev) {
          const int begin = dev.league_rank() * 4;
          const int end   = begin + 4;
          for (int i = begin + dev.team_rank(); i < end; i += dev.team_size()) {
            b(i, 0) = a(i, 1) + a(i, 2);
            b(i, 1) = a(i, 0) - a(i, 3);
            b(i, 2) = a(i, 4) + a(i, 0);
            b(i, 3) = a(i, 2) - a(i, 1);
            b(i, 4) = a(i, 3) + a(i, 4);
          }
        });
  }
  Kokkos::deep_copy(h_b, b);

  double result = 0;
  for (int i = 0; i < 100; i++) {
    for (int j = 0; j < 5; j++) {
      result += h_b(i, j);
    }
  }

  return result;
}
#else
template <class DeviceType, bool PWRTest>
double AddTestLambda() {
  return AddTestFunctor<DeviceType, PWRTest>();
}
#endif

template <class DeviceType>
struct FunctorReduceTest {
  typedef Kokkos::View<double**, DeviceType> view_type;
  typedef DeviceType execution_space;
  typedef double value_type;
  typedef typename Kokkos::TeamPolicy<execution_space>::member_type team_member;

  view_type a_;

  FunctorReduceTest(view_type& a) : a_(a) {}

  KOKKOS_INLINE_FUNCTION
  void operator()(const int& i, value_type& sum) const {
    sum += a_(i, 1) + a_(i, 2);
    sum += a_(i, 0) - a_(i, 3);
    sum += a_(i, 4) + a_(i, 0);
    sum += a_(i, 2) - a_(i, 1);
    sum += a_(i, 3) + a_(i, 4);
  }

  KOKKOS_INLINE_FUNCTION
  void operator()(const team_member& dev, value_type& sum) const {
    const int begin = dev.league_rank() * 4;
    const int end   = begin + 4;
    for (int i = begin + dev.team_rank(); i < end; i += dev.team_size()) {
      sum += a_(i, 1) + a_(i, 2);
      sum += a_(i, 0) - a_(i, 3);
      sum += a_(i, 4) + a_(i, 0);
      sum += a_(i, 2) - a_(i, 1);
      sum += a_(i, 3) + a_(i, 4);
    }
  }

  KOKKOS_INLINE_FUNCTION
  void init(value_type& update) const { update = 0.0; }

  KOKKOS_INLINE_FUNCTION
  void join(volatile value_type& update,
            volatile value_type const& input) const {
    update += input;
  }
};

template <class DeviceType, bool PWRTest>
double ReduceTestFunctor() {
  typedef Kokkos::TeamPolicy<DeviceType> policy_type;
  typedef Kokkos::View<double**, DeviceType> view_type;
  typedef Kokkos::View<double, Kokkos::HostSpace, Kokkos::MemoryUnmanaged>
      unmanaged_result;

  view_type a("A", 100, 5);
  typename view_type::HostMirror h_a = Kokkos::create_mirror_view(a);

  for (int i = 0; i < 100; i++) {
    for (int j = 0; j < 5; j++) {
      h_a(i, j) = 0.1 * i / (1.1 * j + 1.0) + 0.5 * j;
    }
  }
  Kokkos::deep_copy(a, h_a);

  double result = 0.0;
  if (PWRTest == false) {
    Kokkos::parallel_reduce(100, FunctorReduceTest<DeviceType>(a),
                            unmanaged_result(&result));
  } else {
    Kokkos::parallel_reduce(policy_type(25, Kokkos::AUTO),
                            FunctorReduceTest<DeviceType>(a),
                            unmanaged_result(&result));
  }
  Kokkos::fence();

  return result;
}

#if defined(KOKKOS_ENABLE_CXX11_DISPATCH_LAMBDA)
template <class DeviceType, bool PWRTest>
double ReduceTestLambda() {
  typedef Kokkos::TeamPolicy<DeviceType> policy_type;
  typedef Kokkos::View<double**, DeviceType> view_type;
  typedef Kokkos::View<double, Kokkos::HostSpace, Kokkos::MemoryUnmanaged>
      unmanaged_result;

  view_type a("A", 100, 5);
  typename view_type::HostMirror h_a = Kokkos::create_mirror_view(a);

  for (int i = 0; i < 100; i++) {
    for (int j = 0; j < 5; j++) {
      h_a(i, j) = 0.1 * i / (1.1 * j + 1.0) + 0.5 * j;
    }
  }
  Kokkos::deep_copy(a, h_a);

  double result = 0.0;

  if (PWRTest == false) {
    Kokkos::parallel_reduce(
        100,
        KOKKOS_LAMBDA(const int& i, double& sum) {
          sum += a(i, 1) + a(i, 2);
          sum += a(i, 0) - a(i, 3);
          sum += a(i, 4) + a(i, 0);
          sum += a(i, 2) - a(i, 1);
          sum += a(i, 3) + a(i, 4);
        },
        unmanaged_result(&result));
  } else {
    typedef typename policy_type::member_type team_member;
    Kokkos::parallel_reduce(
        policy_type(25, Kokkos::AUTO),
        KOKKOS_LAMBDA(const team_member& dev, double& sum) {
          const int begin = dev.league_rank() * 4;
          const int end   = begin + 4;
          for (int i = begin + dev.team_rank(); i < end; i += dev.team_size()) {
            sum += a(i, 1) + a(i, 2);
            sum += a(i, 0) - a(i, 3);
            sum += a(i, 4) + a(i, 0);
            sum += a(i, 2) - a(i, 1);
            sum += a(i, 3) + a(i, 4);
          }
        },
        unmanaged_result(&result));
  }
  Kokkos::fence();

  return result;
}
#else
template <class DeviceType, bool PWRTest>
double ReduceTestLambda() {
  return ReduceTestFunctor<DeviceType, PWRTest>();
}
#endif

template <class DeviceType>
double TestVariantLambda(int test) {
  switch (test) {
    case 1: return AddTestLambda<DeviceType, false>();
    case 2: return AddTestLambda<DeviceType, true>();
    case 3: return ReduceTestLambda<DeviceType, false>();
    case 4: return ReduceTestLambda<DeviceType, true>();
  }

  return 0;
}

template <class DeviceType>
double TestVariantFunctor(int test) {
  switch (test) {
    case 1: return AddTestFunctor<DeviceType, false>();
    case 2: return AddTestFunctor<DeviceType, true>();
    case 3: return ReduceTestFunctor<DeviceType, false>();
    case 4: return ReduceTestFunctor<DeviceType, true>();
  }

  return 0;
}

template <class DeviceType>
bool Test(int test) {
#ifdef KOKKOS_ENABLE_CXX11_DISPATCH_LAMBDA
  double res_functor = TestVariantFunctor<DeviceType>(test);
  double res_lambda  = TestVariantLambda<DeviceType>(test);

  char testnames[5][256] = {" ", "AddTest", "AddTest TeamPolicy", "ReduceTest",
                            "ReduceTest TeamPolicy"};
  bool passed            = true;

  auto a = res_functor;
  auto b = res_lambda;
  // use a tolerant comparison because functors and lambdas vectorize
  // differently https://github.com/trilinos/Trilinos/issues/3233
  auto rel_err = (std::abs(b - a) / std::max(std::abs(a), std::abs(b)));
  auto tol     = 1e-14;
  if (rel_err > tol) {
    passed = false;

    std::cout << "CXX11 ( test = '" << testnames[test]
              << "' FAILED : relative error " << rel_err << " > tolerance "
              << tol << std::endl;
  }

  return passed;
#else
  (void)test;
  return true;
#endif
}

}  // namespace TestCXX11

namespace Test {
TEST(TEST_CATEGORY, cxx11) {
  if (std::is_same<Kokkos::DefaultExecutionSpace, TEST_EXECSPACE>::value) {
    ASSERT_TRUE((TestCXX11::Test<TEST_EXECSPACE>(1)));
    ASSERT_TRUE((TestCXX11::Test<TEST_EXECSPACE>(2)));
    ASSERT_TRUE((TestCXX11::Test<TEST_EXECSPACE>(3)));
    ASSERT_TRUE((TestCXX11::Test<TEST_EXECSPACE>(4)));
  }
}

}  // namespace Test
