/***************************************************************************************************
 * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: BSD-3-Clause
 *
 * 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 copyright holder nor the names of its
 * contributors may be used to endorse or promote products derived from
 * this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "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 THE COPYRIGHT HOLDER OR 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.
 *
 **************************************************************************************************/
#pragma once

#include <cute/config.hpp>

#include <cute/util/type_traits.hpp>
#include <cute/numeric/math.hpp>

namespace cute
{

template <class T, T v>
struct constant : CUTE_STL_NAMESPACE::integral_constant<T,v> {
  static constexpr T value = v;
  using value_type = T;
  using type = constant<T,v>;
  CUTE_HOST_DEVICE constexpr operator   value_type() const noexcept { return value; }
  CUTE_HOST_DEVICE constexpr value_type operator()() const noexcept { return value; }
};

template <class T, T v>
using integral_constant = constant<T,v>;

template <bool b>
using bool_constant = constant<bool,b>;

using true_type  = bool_constant<true>;
using false_type = bool_constant<false>;

//
// Traits
//

// Use std::is_integral<T> to match built-in integral types (int, int64_t, unsigned, etc)
// Use cute::is_integral<T> to match both built-in integral types AND constant<T,t>

template <class T>
struct is_integral : bool_constant<CUTE_STL_NAMESPACE::is_integral<T>::value> {};
template <class T, T v>
struct is_integral<constant<T,v>> : true_type {};

// is_static detects if an (abstract) value is defined completely by it's type (no members)

template <class T>
struct is_static : bool_constant<is_empty<T>::value> {};

// is_constant detects if a type is a constant<T,v> and if v is equal to a value

template <auto n, class T>
struct is_constant : false_type {};
template <auto n, class T, T v>
struct is_constant<n, constant<T,v>       > : bool_constant<v == n> {};
template <auto n, class T, T v>
struct is_constant<n, constant<T,v> const > : bool_constant<v == n> {};
template <auto n, class T, T v>
struct is_constant<n, constant<T,v> const&> : bool_constant<v == n> {};
template <auto n, class T, T v>
struct is_constant<n, constant<T,v>      &> : bool_constant<v == n> {};
template <auto n, class T, T v>
struct is_constant<n, constant<T,v>     &&> : bool_constant<v == n> {};

//
// Specializations
//

template <int v>
using Int = constant<int,v>;

using _m32    = Int<-32>;
using _m24    = Int<-24>;
using _m16    = Int<-16>;
using _m12    = Int<-12>;
using _m10    = Int<-10>;
using _m9     = Int<-9>;
using _m8     = Int<-8>;
using _m7     = Int<-7>;
using _m6     = Int<-6>;
using _m5     = Int<-5>;
using _m4     = Int<-4>;
using _m3     = Int<-3>;
using _m2     = Int<-2>;
using _m1     = Int<-1>;
using _0      = Int<0>;
using _1      = Int<1>;
using _2      = Int<2>;
using _3      = Int<3>;
using _4      = Int<4>;
using _5      = Int<5>;
using _6      = Int<6>;
using _7      = Int<7>;
using _8      = Int<8>;
using _9      = Int<9>;
using _10     = Int<10>;
using _12     = Int<12>;
using _16     = Int<16>;
using _24     = Int<24>;
using _32     = Int<32>;
using _64     = Int<64>;
using _96     = Int<96>;
using _128    = Int<128>;
using _192    = Int<192>;
using _256    = Int<256>;
using _512    = Int<512>;
using _1024   = Int<1024>;
using _2048   = Int<2048>;
using _4096   = Int<4096>;
using _8192   = Int<8192>;
using _16384  = Int<16384>;
using _32768  = Int<32768>;
using _65536  = Int<65536>;
using _131072 = Int<131072>;
using _262144 = Int<262144>;
using _524288 = Int<524288>;

/***************/
/** Operators **/
/***************/

#define CUTE_LEFT_UNARY_OP(OP)                                       \
  template <class T, T t>                                            \
  CUTE_HOST_DEVICE constexpr                                         \
  constant<decltype(OP t), (OP t)>                                   \
  operator OP (constant<T,t>) {                                      \
    return {};                                                       \
  }
#define CUTE_RIGHT_UNARY_OP(OP)                                      \
  template <class T, T t>                                            \
  CUTE_HOST_DEVICE constexpr                                         \
  constant<decltype(t OP), (t OP)>                                   \
  operator OP (constant<T,t>) {                                      \
    return {};                                                       \
  }

#define CUTE_BINARY_OP(OP)                                           \
  template <class T, T t, class U, U u>                              \
  CUTE_HOST_DEVICE constexpr                                         \
  constant<decltype(t OP u), (t OP u)>                               \
  operator OP (constant<T,t>, constant<U,u>) {                       \
    return {};                                                       \
  }

CUTE_LEFT_UNARY_OP(+);
CUTE_LEFT_UNARY_OP(-);
CUTE_LEFT_UNARY_OP(~);
CUTE_LEFT_UNARY_OP(!);
CUTE_LEFT_UNARY_OP(*);

CUTE_BINARY_OP( +);
CUTE_BINARY_OP( -);
CUTE_BINARY_OP( *);
CUTE_BINARY_OP( /);
CUTE_BINARY_OP( %);
CUTE_BINARY_OP( &);
CUTE_BINARY_OP( |);
CUTE_BINARY_OP( ^);
CUTE_BINARY_OP(<<);
CUTE_BINARY_OP(>>);

CUTE_BINARY_OP(&&);
CUTE_BINARY_OP(||);

CUTE_BINARY_OP(==);
CUTE_BINARY_OP(!=);
CUTE_BINARY_OP( >);
CUTE_BINARY_OP( <);
CUTE_BINARY_OP(>=);
CUTE_BINARY_OP(<=);

#undef CUTE_BINARY_OP
#undef CUTE_LEFT_UNARY_OP
#undef CUTE_RIGHT_UNARY_OP

//
// Mixed static-dynamic special cases
//

template <class T, class U,
          __CUTE_REQUIRES(CUTE_STL_NAMESPACE::is_integral<U>::value)>
CUTE_HOST_DEVICE constexpr
constant<T, 0>
operator*(constant<T, 0>, U) {
  return {};
}

template <class U, class T,
          __CUTE_REQUIRES(CUTE_STL_NAMESPACE::is_integral<U>::value)>
CUTE_HOST_DEVICE constexpr
constant<T, 0>
operator*(U, constant<T, 0>) {
  return {};
}

template <class T, class U,
          __CUTE_REQUIRES(CUTE_STL_NAMESPACE::is_integral<U>::value)>
CUTE_HOST_DEVICE constexpr
constant<T, 0>
operator/(constant<T, 0>, U) {
  return {};
}

template <class U, class T,
          __CUTE_REQUIRES(CUTE_STL_NAMESPACE::is_integral<U>::value)>
CUTE_HOST_DEVICE constexpr
constant<T, 0>
operator%(U, constant<T, 1>) {
  return {};
}

template <class U, class T,
          __CUTE_REQUIRES(CUTE_STL_NAMESPACE::is_integral<U>::value)>
CUTE_HOST_DEVICE constexpr
constant<T, 0>
operator%(U, constant<T,-1>) {
  return {};
}

template <class T, class U,
          __CUTE_REQUIRES(CUTE_STL_NAMESPACE::is_integral<U>::value)>
CUTE_HOST_DEVICE constexpr
constant<T, 0>
operator%(constant<T, 0>, U) {
  return {};
}

template <class T, class U,
          __CUTE_REQUIRES(CUTE_STL_NAMESPACE::is_integral<U>::value)>
CUTE_HOST_DEVICE constexpr
constant<T, 0>
operator&(constant<T, 0>, U) {
  return {};
}

template <class T, class U,
          __CUTE_REQUIRES(CUTE_STL_NAMESPACE::is_integral<U>::value)>
CUTE_HOST_DEVICE constexpr
constant<T, 0>
operator&(U, constant<T, 0>) {
  return {};
}

template <class T, T t, class U,
          __CUTE_REQUIRES(CUTE_STL_NAMESPACE::is_integral<U>::value && !bool(t))>
CUTE_HOST_DEVICE constexpr
constant<bool, false>
operator&&(constant<T, t>, U) {
  return {};
}

template <class T, T t, class U,
          __CUTE_REQUIRES(CUTE_STL_NAMESPACE::is_integral<U>::value && !bool(t))>
CUTE_HOST_DEVICE constexpr
constant<bool, false>
operator&&(U, constant<T, t>) {
  return {};
}

template <class T, class U, T t,
          __CUTE_REQUIRES(CUTE_STL_NAMESPACE::is_integral<U>::value && bool(t))>
CUTE_HOST_DEVICE constexpr
constant<bool, true>
operator||(constant<T, t>, U) {
  return {};
}

template <class T, class U, T t,
          __CUTE_REQUIRES(CUTE_STL_NAMESPACE::is_integral<U>::value && bool(t))>
CUTE_HOST_DEVICE constexpr
constant<bool, true>
operator||(U, constant<T, t>) {
  return {};
}

//
// Named functions from math.hpp
//

#define CUTE_NAMED_UNARY_FN(OP)                                      \
  template <class T, T t>                                            \
  CUTE_HOST_DEVICE constexpr                                         \
  constant<decltype(OP(t)), OP(t)>                                   \
  OP (constant<T,t>) {                                               \
    return {};                                                       \
  }

#define CUTE_NAMED_BINARY_FN(OP)                                     \
  template <class T, T t, class U, U u>                              \
  CUTE_HOST_DEVICE constexpr                                         \
  constant<decltype(OP(t,u)), OP(t,u)>                               \
  OP (constant<T,t>, constant<U,u>) {                                \
    return {};                                                       \
  }                                                                  \
                                                                     \
  template <class T, T t, class U,                                   \
            __CUTE_REQUIRES(CUTE_STL_NAMESPACE::is_integral<U>::value)>             \
  CUTE_HOST_DEVICE constexpr                                         \
  auto                                                               \
  OP (constant<T,t>, U u) {                                          \
    return OP(t,u);                                                  \
  }                                                                  \
                                                                     \
  template <class T, class U, U u,                                   \
            __CUTE_REQUIRES(CUTE_STL_NAMESPACE::is_integral<T>::value)>             \
  CUTE_HOST_DEVICE constexpr                                         \
  auto                                                               \
  OP (T t, constant<U,u>) {                                          \
    return OP(t,u);                                                  \
  }

CUTE_NAMED_UNARY_FN(abs);
CUTE_NAMED_UNARY_FN(signum);
CUTE_NAMED_UNARY_FN(has_single_bit);

CUTE_NAMED_BINARY_FN(max);
CUTE_NAMED_BINARY_FN(min);
CUTE_NAMED_BINARY_FN(shiftl);
CUTE_NAMED_BINARY_FN(shiftr);
CUTE_NAMED_BINARY_FN(gcd);
CUTE_NAMED_BINARY_FN(lcm);

#undef CUTE_NAMED_UNARY_FN
#undef CUTE_NAMED_BINARY_FN

//
// Other functions
//

template <class T, T t, class U, U u>
CUTE_HOST_DEVICE constexpr
constant<decltype(t / u), t / u>
safe_div(constant<T, t>, constant<U, u>) {
  static_assert(t % u == 0, "Static safe_div requires t % u == 0");
  return {};
}

template <class T, T t, class U,
          __CUTE_REQUIRES(CUTE_STL_NAMESPACE::is_integral<U>::value)>
CUTE_HOST_DEVICE constexpr
auto
safe_div(constant<T, t>, U u) {
  return t / u;
}

template <class T, class U, U u,
          __CUTE_REQUIRES(CUTE_STL_NAMESPACE::is_integral<T>::value)>
CUTE_HOST_DEVICE constexpr
auto
safe_div(T t, constant<U, u>) {
  return t / u;
}

// cute::true_type prefers standard conversion to std::true_type
//   over user-defined conversion to bool
template <class TrueType, class FalseType>
CUTE_HOST_DEVICE constexpr
decltype(auto)
conditional_return(true_type, TrueType&& t, FalseType&&) {
  return static_cast<TrueType&&>(t);
}

// cute::false_type prefers standard conversion to std::false_type
//   over user-defined conversion to bool
template <class TrueType, class FalseType>
CUTE_HOST_DEVICE constexpr
decltype(auto)
conditional_return(false_type, TrueType&&, FalseType&& f) {
  return static_cast<FalseType&&>(f);
}

// TrueType and FalseType must have a common type
template <class TrueType, class FalseType>
CUTE_HOST_DEVICE constexpr
auto
conditional_return(bool b, TrueType const& t, FalseType const& f) {
  return b ? t : f;
}

// TrueType and FalseType don't require a common type
template <bool b, class TrueType, class FalseType>
CUTE_HOST_DEVICE constexpr
auto
conditional_return(TrueType const& t, FalseType const& f) {
  if constexpr (b) {
    return t;
  } else {
    return f;
  }
}

//
// Display utilities
//

template <class T, T N>
CUTE_HOST_DEVICE void print(integral_constant<T,N> const&) {
  printf("_%d", N);
}

#if !defined(__CUDACC_RTC__)
template <class T, T N>
CUTE_HOST std::ostream& operator<<(std::ostream& os, integral_constant<T,N> const&) {
  return os << "_" << N;
}
#endif

} // end namespace cute
