// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// SPDX-FileCopyrightText: Copyright Contributors to the Kokkos project

/// \author Kyungjoo Kim (kyukim@sandia.gov)

#include "gtest/gtest.h"
#include "Kokkos_Core.hpp"
#include "Kokkos_Random.hpp"

#include "KokkosBlas1_set.hpp"
#include "KokkosBlas1_scal.hpp"

#include "KokkosKernels_TestUtils.hpp"

using namespace KokkosBlas;

namespace Test {

enum : int { BlasSet = 0, BlasScale = 1 };

struct KokkosKernelTag {};
struct NaiveTag {};

template <typename DeviceType, typename ViewType, typename ScalarType, typename AlgoTagType, int TestID>
struct Functor_TestBlasSerialMatUtil {
  ScalarType _alpha;
  ViewType _a;

  KOKKOS_INLINE_FUNCTION
  Functor_TestBlasSerialMatUtil(const ScalarType alpha, const ViewType &a) : _alpha(alpha), _a(a) {}

  KOKKOS_INLINE_FUNCTION
  void operator()(const KokkosKernelTag &, const int i) const {
    auto A = Kokkos::subview(_a, i, Kokkos::ALL(), Kokkos::ALL());
    switch (TestID) {
      case BlasSet: KokkosBlas::SerialSet::invoke(_alpha, A); break;
      case BlasScale: KokkosBlas::SerialScale::invoke(_alpha, A); break;
    }
  }

  KOKKOS_INLINE_FUNCTION
  void operator()(const NaiveTag &, const int k) const {
    // MD Note: changing because of the error with -werror
    auto A      = Kokkos::subview(_a, k, Kokkos::ALL(), Kokkos::ALL());
    const int m = A.extent(0), n = A.extent(1);
    switch (TestID) {
      case BlasSet: {
        for (int i = 0; i < m; ++i)
          for (int j = 0; j < n; ++j) A(i, j) = _alpha;
        break;
      }
      case BlasScale: {
        for (int i = 0; i < m; ++i)
          for (int j = 0; j < n; ++j) A(i, j) *= _alpha;
        break;
      }
    }
  }

  inline int run() {
    typedef typename ViewType::value_type value_type;
    std::string name_region("KokkosBlas::Test::SerialMatUtil");
    const std::string name_value_type = Test::value_type_name<value_type>();
    std::string name_work_tag         = (std::is_same<AlgoTagType, KokkosKernelTag>::value ? "::KokkosBlas"
                                         : std::is_same<AlgoTagType, NaiveTag>::value      ? "::Naive"
                                                                                           : "::UnknownWorkTag");
    std::string name_test_id          = (TestID == BlasSet ? "Set" : TestID == BlasScale ? "Scale" : "UnknownTest");
    std::string name                  = name_region + name_value_type + name_work_tag + name_test_id;
    Kokkos::Profiling::pushRegion(name.c_str());
    Kokkos::RangePolicy<typename DeviceType::execution_space, AlgoTagType> policy(0, _a.extent(0));
    Kokkos::parallel_for(name.c_str(), policy, *this);
    Kokkos::Profiling::popRegion();
    return 0;
  }
};

template <typename DeviceType, typename ViewType, typename ScalarType, int TestID>
void impl_test_blas_matutil(const int N, const int BlkSize) {
  /// typedefs
  typedef typename ViewType::value_type value_type;
  typedef KokkosKernels::ArithTraits<value_type> ats;

  /// radomized input testing views
  const ScalarType alpha = 11.1;
  ViewType a("a", N, BlkSize, BlkSize);
  ViewType b("b", N, BlkSize, BlkSize);

  Kokkos::Random_XorShift64_Pool<typename DeviceType::execution_space> random(13718);
  Kokkos::fill_random(a, random, value_type(1.0));

  Kokkos::fence();

  Kokkos::deep_copy(b, a);

  /// test body
  Functor_TestBlasSerialMatUtil<DeviceType, ViewType, ScalarType, NaiveTag, TestID>(alpha, a).run();
  Functor_TestBlasSerialMatUtil<DeviceType, ViewType, ScalarType, KokkosKernelTag, TestID>(alpha, b).run();

  Kokkos::fence();

  /// for comparison send it to host
  typename ViewType::host_mirror_type a_host = Kokkos::create_mirror_view(a);
  typename ViewType::host_mirror_type b_host = Kokkos::create_mirror_view(b);

  Kokkos::deep_copy(a_host, a);
  Kokkos::deep_copy(b_host, b);

  /// check a = b
  typename ats::mag_type eps = 100 * std::numeric_limits<typename ats::mag_type>::epsilon();
  for (int k = 0; k < N; ++k)
    for (int i = 0; i < BlkSize; ++i)
      for (int j = 0; j < BlkSize; ++j) EXPECT_NEAR_KK(b_host(k, i, j), a_host(k, i, j), eps);
}
}  // namespace Test

template <typename DeviceType, typename ValueType, typename ScalarType, int TestID>
int test_blas_matutil() {
#if defined(KOKKOSKERNELS_INST_LAYOUTLEFT)
  {
    typedef Kokkos::View<ValueType ***, Kokkos::LayoutLeft, DeviceType> ViewType;
    Test::impl_test_blas_matutil<DeviceType, ViewType, ScalarType, TestID>(0, 10);
    Test::impl_test_blas_matutil<DeviceType, ViewType, ScalarType, TestID>(10, 15);
    Test::impl_test_blas_matutil<DeviceType, ViewType, ScalarType, TestID>(1024, 9);
    Test::impl_test_blas_matutil<DeviceType, ViewType, ScalarType, TestID>(132231, 3);
  }
#endif
#if defined(KOKKOSKERNELS_INST_LAYOUTRIGHT)
  {
    typedef Kokkos::View<ValueType ***, Kokkos::LayoutRight, DeviceType> ViewType;
    Test::impl_test_blas_matutil<DeviceType, ViewType, ScalarType, TestID>(0, 10);
    Test::impl_test_blas_matutil<DeviceType, ViewType, ScalarType, TestID>(10, 15);
    Test::impl_test_blas_matutil<DeviceType, ViewType, ScalarType, TestID>(1024, 9);
    Test::impl_test_blas_matutil<DeviceType, ViewType, ScalarType, TestID>(132231, 3);
  }
#endif

  return 0;
}

// Real test cases

#if defined(KOKKOSKERNELS_INST_FLOAT)
TEST_F(TestCategory, blas_scalar_serial_set_float_float) {
  test_blas_matutil<TestDevice, float, float, ::Test::BlasSet>();
}
TEST_F(TestCategory, blas_scalar_serial_scale_float_float) {
  test_blas_matutil<TestDevice, float, float, ::Test::BlasScale>();
}
#endif

#if defined(KOKKOSKERNELS_INST_DOUBLE)
TEST_F(TestCategory, blas_scalar_serial_set_double_double) {
  test_blas_matutil<TestDevice, double, double, ::Test::BlasSet>();
}
TEST_F(TestCategory, blas_scalar_serial_scale_double_double) {
  test_blas_matutil<TestDevice, double, double, ::Test::BlasScale>();
}
#endif

// Complex test cases

#if defined(KOKKOSKERNELS_INST_COMPLEX_DOUBLE)
TEST_F(TestCategory, blas_scalar_serial_set_dcomplex_dcomplex) {
  test_blas_matutil<TestDevice, Kokkos::complex<double>, Kokkos::complex<double>, ::Test::BlasSet>();
}
TEST_F(TestCategory, blas_scalar_serial_scale_dcomplex_dcomplex) {
  test_blas_matutil<TestDevice, Kokkos::complex<double>, Kokkos::complex<double>, ::Test::BlasScale>();
}
TEST_F(TestCategory, blas_scalar_serial_set_dcomplex_double) {
  test_blas_matutil<TestDevice, Kokkos::complex<double>, double, ::Test::BlasSet>();
}
TEST_F(TestCategory, blas_scalar_serial_scale_dcomplex_double) {
  test_blas_matutil<TestDevice, Kokkos::complex<double>, double, ::Test::BlasScale>();
}
#endif
