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

#ifndef KOKKOS_SIMD_COMMON_MATH_HPP
#define KOKKOS_SIMD_COMMON_MATH_HPP

#include <Kokkos_Macros.hpp>
#ifdef KOKKOS_ENABLE_EXPERIMENTAL_CXX20_MODULES
import kokkos.core; // Kokkos::min, etc.
#else
#include <Kokkos_Core.hpp>  // Kokkos::min, etc.
#endif

namespace Kokkos {

namespace Experimental {

namespace simd_abi {
class scalar;
}

template <class T, class Abi>
class basic_simd;

template <class T, class Abi>
class basic_simd_mask;

template <typename T, Impl::NonScalarAbi Abi>
KOKKOS_IMPL_HOST_FORCEINLINE_FUNCTION T
reduce_min(basic_simd<T, Abi> const& v,
           typename basic_simd<T, Abi>::mask_type const& m) {
  auto result = Kokkos::reduction_identity<T>::min();
  for (Impl::simd_size_t i = 0; i < v.size(); ++i) {
    if (m[i]) result = Kokkos::min(result, v[i]);
  }
  return result;
}

template <class T, Impl::NonScalarAbi Abi>
KOKKOS_IMPL_HOST_FORCEINLINE_FUNCTION T
reduce_max(basic_simd<T, Abi> const& v,
           typename basic_simd<T, Abi>::mask_type const& m) {
  auto result = Kokkos::reduction_identity<T>::max();
  for (Impl::simd_size_t i = 0; i < v.size(); ++i) {
    if (m[i]) result = Kokkos::max(result, v[i]);
  }
  return result;
}

template <class T, Impl::NonScalarAbi Abi, class BinaryOperation = std::plus<>>
KOKKOS_IMPL_HOST_FORCEINLINE_FUNCTION T
reduce(basic_simd<T, Abi> const& v,
       typename basic_simd<T, Abi>::mask_type const& m, BinaryOperation op = {},
       T identity = Impl::Identity<T, BinaryOperation>()) {
  if (none_of(m)) {
    return identity;
  }
  T result = identity;
  for (Impl::simd_size_t i = 0; i < v.size(); ++i) {
    if (m[i]) result = op(result, v[i]);
  }
  return result;
}

}  // namespace Experimental

template <class T, Experimental::Impl::NonScalarAbi Abi>
KOKKOS_IMPL_HOST_FORCEINLINE_FUNCTION Experimental::basic_simd<T, Abi> min(
    Experimental::basic_simd<T, Abi> const& a,
    Experimental::basic_simd<T, Abi> const& b) {
  using simd_type           = Experimental::basic_simd<T, Abi>;
  T vals[simd_type::size()] = {0};
  for (Experimental::Impl::simd_size_t i = 0; i < simd_type::size(); ++i) {
    vals[i] = Kokkos::min(a[i], b[i]);
  }
  return Experimental::simd_unchecked_load<simd_type>(
      vals, Experimental::simd_flag_default);
}

template <class T, Experimental::Impl::NonScalarAbi Abi>
KOKKOS_IMPL_HOST_FORCEINLINE_FUNCTION Experimental::basic_simd<T, Abi> max(
    Experimental::basic_simd<T, Abi> const& a,
    Experimental::basic_simd<T, Abi> const& b) {
  using simd_type           = Experimental::basic_simd<T, Abi>;
  T vals[simd_type::size()] = {0};
  for (Experimental::Impl::simd_size_t i = 0; i < simd_type::size(); ++i) {
    vals[i] = Kokkos::max(a[i], b[i]);
  }
  return Experimental::simd_unchecked_load<simd_type>(
      vals, Experimental::simd_flag_default);
}

// fallback implementations of <cmath> functions.
// individual Abi types may provide overloads with more efficient
// implementations.

#define KOKKOS_IMPL_SIMD_UNARY_FUNCTION(FUNC)                                  \
  template <class T, Experimental::Impl::NonScalarAbi Abi>                     \
  KOKKOS_IMPL_HOST_FORCEINLINE_FUNCTION Experimental::basic_simd<T, Abi> FUNC( \
      Experimental::basic_simd<T, Abi> const& a) {                             \
    using simd_type           = Experimental::basic_simd<T, Abi>;              \
    T vals[simd_type::size()] = {0};                                           \
    for (Experimental::Impl::simd_size_t i = 0; i < simd_type::size(); ++i) {  \
      vals[i] = Kokkos::FUNC(a[i]);                                            \
    }                                                                          \
    return Experimental::simd_unchecked_load<simd_type>(                       \
        vals, Experimental::simd_flag_default);                                \
  }                                                                            \
  template <class T, Experimental::Impl::ScalarAbi Abi>                        \
  KOKKOS_FORCEINLINE_FUNCTION constexpr Experimental::basic_simd<T, Abi> FUNC( \
      Experimental::basic_simd<T, Abi> const& a) {                             \
    return Kokkos::FUNC(a[0]);                                                 \
  }

KOKKOS_IMPL_SIMD_UNARY_FUNCTION(abs)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(exp)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(exp2)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(log)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(log10)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(log2)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(sqrt)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(cbrt)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(sin)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(cos)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(tan)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(asin)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(acos)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(atan)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(sinh)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(cosh)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(tanh)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(asinh)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(acosh)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(atanh)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(erf)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(erfc)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(tgamma)
KOKKOS_IMPL_SIMD_UNARY_FUNCTION(lgamma)

#define KOKKOS_IMPL_SIMD_BINARY_FUNCTION(FUNC)                                 \
  template <class T, Experimental::Impl::NonScalarAbi Abi>                     \
  KOKKOS_IMPL_HOST_FORCEINLINE_FUNCTION Experimental::basic_simd<T, Abi> FUNC( \
      Experimental::basic_simd<T, Abi> const& a,                               \
      Experimental::basic_simd<T, Abi> const& b) {                             \
    using simd_type           = Experimental::basic_simd<T, Abi>;              \
    T vals[simd_type::size()] = {0};                                           \
    for (Experimental::Impl::simd_size_t i = 0; i < simd_type::size(); ++i) {  \
      vals[i] = Kokkos::FUNC(a[i], b[i]);                                      \
    }                                                                          \
    return Experimental::simd_unchecked_load<simd_type>(                       \
        vals, Experimental::simd_flag_default);                                \
  }                                                                            \
  template <class T, Experimental::Impl::ScalarAbi Abi>                        \
  KOKKOS_FORCEINLINE_FUNCTION constexpr Experimental::basic_simd<T, Abi> FUNC( \
      Experimental::basic_simd<T, Abi> const& a,                               \
      Experimental::basic_simd<T, Abi> const& b) {                             \
    return Kokkos::FUNC(a[0], b[0]);                                           \
  }

KOKKOS_IMPL_SIMD_BINARY_FUNCTION(pow)
KOKKOS_IMPL_SIMD_BINARY_FUNCTION(hypot)
KOKKOS_IMPL_SIMD_BINARY_FUNCTION(atan2)
KOKKOS_IMPL_SIMD_BINARY_FUNCTION(copysign)

#define KOKKOS_IMPL_SIMD_TERNARY_FUNCTION(FUNC)                                \
  template <class T, Experimental::Impl::NonScalarAbi Abi>                     \
  KOKKOS_IMPL_HOST_FORCEINLINE_FUNCTION Experimental::basic_simd<T, Abi> FUNC( \
      Experimental::basic_simd<T, Abi> const& a,                               \
      Experimental::basic_simd<T, Abi> const& b,                               \
      Experimental::basic_simd<T, Abi> const& c) {                             \
    using simd_type           = Experimental::basic_simd<T, Abi>;              \
    T vals[simd_type::size()] = {0};                                           \
    for (Experimental::Impl::simd_size_t i = 0; i < simd_type::size(); ++i) {  \
      vals[i] = Kokkos::FUNC(a[i], b[i], c[i]);                                \
    }                                                                          \
    return Experimental::simd_unchecked_load<simd_type>(                       \
        vals, Experimental::simd_flag_default);                                \
  }                                                                            \
  template <class T, Experimental::Impl::ScalarAbi Abi>                        \
  KOKKOS_FORCEINLINE_FUNCTION constexpr Experimental::basic_simd<T, Abi> FUNC( \
      Experimental::basic_simd<T, Abi> const& a,                               \
      Experimental::basic_simd<T, Abi> const& b,                               \
      Experimental::basic_simd<T, Abi> const& c) {                             \
    return Kokkos::FUNC(a[0], b[0], c[0]);                                     \
  }

KOKKOS_IMPL_SIMD_TERNARY_FUNCTION(fma)
KOKKOS_IMPL_SIMD_TERNARY_FUNCTION(hypot)

}  // namespace Kokkos

#endif
