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

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

#include "KokkosBatched_Vector.hpp"

#include "KokkosKernels_TestUtils.hpp"

using namespace KokkosBatched;

namespace Test {

  template<typename VectorTagType,int VectorLength>
  void impl_test_batched_vector_math() {
    /// random data initialization
    typedef Vector<VectorTagType,VectorLength> vector_type;
    
    typedef typename vector_type::value_type value_type;    
    const int vector_length = vector_type::vector_length;
    
    typedef Kokkos::Details::ArithTraits<value_type> ats;
    typedef typename ats::mag_type mag_type;

    vector_type a, b, aref, bref;
    const value_type one(1.0);
    const value_type half(0.5);
    const value_type maxmin(1.0e-7);
    const mag_type eps = 1.0e3 * ats::epsilon();

    Random<value_type> random;
    for (int iter=0;iter<100;++iter) {
      for (int k=0;k<vector_length;++k) {
        const auto aval = (random.value() + half);
        const auto bval = (random.value() + half);
        
        aref[k] = max(min(aval, one), maxmin);
        bref[k] = max(min(bval, one), maxmin);
      }

      {

#undef CHECK
#define CHECK(op)                                               \
        {                                                       \
          a = op(aref);                                         \
          for (int i=0;i<vector_length;++i)                     \
            EXPECT_NEAR_KK( a[i], ats::op(aref[i]), eps*a[i]);  \
        }
        
        CHECK(sqrt);
        CHECK(cbrt);
        CHECK(log);
        CHECK(exp);
        CHECK(sin);
        CHECK(cos);
        CHECK(tan);
        CHECK(sinh);
        CHECK(cosh);
        CHECK(tanh);
        CHECK(asin);
        CHECK(acos);
        CHECK(atan);

#undef CHECK
#define CHECK                                                           \
        {                                                               \
          a = pow(aref,bref);                                           \
          for (int i=0;i<vector_length;++i)                             \
            EXPECT_NEAR_KK( a[i], ats::pow(aref[i], bref[i]), eps*a[i] ); \
        }                                                               \
        CHECK;
        
#undef CHECK
#define CHECK(op)                                                       \
        {                                                               \
          mag_type beta = mag_type(3.2);                                \
          a = op(aref,beta);                                            \
          for (int i=0;i<vector_length;++i)                             \
            EXPECT_NEAR_KK( a[i], ats::op(aref[i], beta), eps*a[i] );   \
        }
        
        CHECK(pow);
        
#undef CHECK
#define CHECK(op)                                                       \
        {                                                               \
          value_type alpha = random.value() + 2.0;                      \
          a = op(alpha,bref);                                           \
          for (int i=0;i<vector_length;++i)                             \
            EXPECT_NEAR_KK( a[i], ats::op(alpha, bref[i]), eps*a[i] );  \
        }
        
        CHECK(pow);
#undef CHECK
      } // end test body
    } // end for
  } // impl
} // namespace

template<typename DeviceType,typename VectorTagType,int VectorLength>
int test_batched_vector_math() {
  static_assert(Kokkos::Impl::SpaceAccessibility<DeviceType,Kokkos::HostSpace >::accessible,
                "vector datatype is only tested on host space");
  Test::impl_test_batched_vector_math<VectorTagType,VectorLength>();
  
  return 0;
}

// template<typename ValueType>
// int test_complex_pow() {
//   typedef Kokkos::Details::ArithTraits<Kokkos::complex<ValueType> > ats;
//   typedef typename ats::mag_type mag_type;

//   const mag_type eps = 1.0e3 * ats::epsilon();
  
//   mag_type a2 = 4.5;
//   Kokkos::complex<mag_type> a0(3.2, -1.4), a1(1.2, 2.3);
//   std::complex<mag_type> b0(3.2, -1.4), b1(1.2, 2.3);
  
//   Test::EXPECT_NEAR_KK( ats::pow(a0,a1), std::pow(b0,b1), eps );
//   Test::EXPECT_NEAR_KK( ats::pow(a0,a2), std::pow(b0,a2), eps );

//   return 0;
// }

///
/// SIMD
///

#if defined(KOKKOSKERNELS_INST_FLOAT)
TEST_F( TestCategory, batched_vector_math_simd_float3 ) {
  test_batched_vector_math<TestExecSpace,SIMD<float>,3>();
}
TEST_F( TestCategory, batched_vector_math_simd_float8 ) {
  test_batched_vector_math<TestExecSpace,SIMD<float>,8>();
}
#endif

#if defined(KOKKOSKERNELS_INST_DOUBLE)
TEST_F( TestCategory, batched_vector_math_simd_double3 ) {
  test_batched_vector_math<TestExecSpace,SIMD<double>,3>();
}
TEST_F( TestCategory, batched_vector_math_simd_double4 ) {
  test_batched_vector_math<TestExecSpace,SIMD<double>,4>();
}
#endif


// using namespace Test;

// #if defined(KOKKOSKERNELS_INST_COMPLEX_FLOAT)
// TEST_F( TestCategory, batched_vector_math_simd_scomplex3 ) {
//   test_complex_pow<float>();
//   test_batched_vector_math<TestExecSpace,SIMD<Kokkos::complex<float> >,3>();
// }
// TEST_F( TestCategory, batched_vector_math_simd_scomplex4 ) {
//   test_batched_vector_math<TestExecSpace,SIMD<Kokkos::complex<float> >,4>();
// }
// #endif

// #if defined(KOKKOSKERNELS_INST_COMPLEX_DOUBLE)
// TEST_F( TestCategory, batched_vector_math_simd_dcomplex3 ) {
//   test_complex_pow<double>();
//   test_batched_vector_math<TestExecSpace,SIMD<Kokkos::complex<double> >,3>();
// }
// TEST_F( TestCategory, batched_vector_math_simd_dcomplex2 ) {
//   test_batched_vector_math<TestExecSpace,SIMD<Kokkos::complex<double> >,2>();
// }
// #endif
