// SPDX-FileCopyrightText: Copyright (c) 2008-2021, NVIDIA Corporation. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

#pragma once

#include <thrust/detail/config.h>

#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
#  pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
#  pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
#  pragma system_header
#endif // no system header

#include <thrust/detail/type_traits.h>
#include <thrust/random/uniform_int_distribution.h>
#include <thrust/random/uniform_real_distribution.h>

THRUST_NAMESPACE_BEGIN

namespace random
{
template <typename IntType>
_CCCL_HOST_DEVICE uniform_int_distribution<IntType>::uniform_int_distribution(IntType a, IntType b)
    : m_param(a, b)
{} // end uniform_int_distribution::uniform_int_distribution()

template <typename IntType>
_CCCL_HOST_DEVICE uniform_int_distribution<IntType>::uniform_int_distribution(const param_type& parm)
    : m_param(parm)
{} // end uniform_int_distribution::uniform_int_distribution()

template <typename IntType>
_CCCL_HOST_DEVICE void uniform_int_distribution<IntType>::reset()
{} // end uniform_int_distribution::reset()

template <typename IntType>
template <typename UniformRandomNumberGenerator>
_CCCL_HOST_DEVICE typename uniform_int_distribution<IntType>::result_type
uniform_int_distribution<IntType>::operator()(UniformRandomNumberGenerator& urng)
{
  return operator()(urng, m_param);
} // end uniform_int_distribution::operator()()

template <typename IntType>
template <typename UniformRandomNumberGenerator>
_CCCL_HOST_DEVICE typename uniform_int_distribution<IntType>::result_type
uniform_int_distribution<IntType>::operator()(UniformRandomNumberGenerator& urng, const param_type& parm)
{
  // XXX this implementation is somewhat hacky and will skip
  //     values if the range of the RNG is smaller than the range of the distribution
  //     we should improve this implementation in a later version

  using float_type = typename thrust::detail::largest_available_float::type;

  const float_type real_min(static_cast<float_type>(parm.first));
  const float_type real_max(static_cast<float_type>(parm.second));

  // add one to the right end of the interval because it is half-open
  // XXX adding 1.0 to a potentially large floating point number seems like a bad idea
  uniform_real_distribution<float_type> real_dist(real_min, real_max + float_type(1));

  return static_cast<result_type>(real_dist(urng));
} // end uniform_int_distribution::operator()()

template <typename IntType>
_CCCL_HOST_DEVICE typename uniform_int_distribution<IntType>::result_type uniform_int_distribution<IntType>::a() const
{
  return m_param.first;
} // end uniform_int_distribution<IntType>::a()

template <typename IntType>
_CCCL_HOST_DEVICE typename uniform_int_distribution<IntType>::result_type uniform_int_distribution<IntType>::b() const
{
  return m_param.second;
} // end uniform_int_distribution::b()

template <typename IntType>
_CCCL_HOST_DEVICE typename uniform_int_distribution<IntType>::param_type uniform_int_distribution<IntType>::param() const
{
  return m_param;
} // end uniform_int_distribution::param()

template <typename IntType>
_CCCL_HOST_DEVICE void uniform_int_distribution<IntType>::param(const param_type& parm)
{
  m_param = parm;
} // end uniform_int_distribution::param()

template <typename IntType>
_CCCL_HOST_DEVICE typename uniform_int_distribution<IntType>::result_type uniform_int_distribution<IntType>::min
THRUST_PREVENT_MACRO_SUBSTITUTION() const
{
  return a();
} // end uniform_int_distribution::min()

template <typename IntType>
_CCCL_HOST_DEVICE typename uniform_int_distribution<IntType>::result_type uniform_int_distribution<IntType>::max
THRUST_PREVENT_MACRO_SUBSTITUTION() const
{
  return b();
} // end uniform_int_distribution::max()

template <typename IntType>
_CCCL_HOST_DEVICE bool uniform_int_distribution<IntType>::equal(const uniform_int_distribution& rhs) const
{
  return param() == rhs.param();
}

template <typename IntType>
template <typename CharT, typename Traits>
std::basic_ostream<CharT, Traits>&
uniform_int_distribution<IntType>::stream_out(std::basic_ostream<CharT, Traits>& os) const
{
  using ostream_type = std::basic_ostream<CharT, Traits>;
  using ios_base     = typename ostream_type::ios_base;

  // save old flags and fill character
  const typename ios_base::fmtflags flags = os.flags();
  const CharT fill                        = os.fill();

  const CharT space = os.widen(' ');
  os.flags(ios_base::dec | ios_base::fixed | ios_base::left);
  os.fill(space);

  os << a() << space << b();

  // restore old flags and fill character
  os.flags(flags);
  os.fill(fill);
  return os;
}

template <typename IntType>
template <typename CharT, typename Traits>
std::basic_istream<CharT, Traits>& uniform_int_distribution<IntType>::stream_in(std::basic_istream<CharT, Traits>& is)
{
  using istream_type = std::basic_istream<CharT, Traits>;
  using ios_base     = typename istream_type::ios_base;

  // save old flags
  const typename ios_base::fmtflags flags = is.flags();

  is.flags(ios_base::skipws);

  is >> m_param.first >> m_param.second;

  // restore old flags
  is.flags(flags);
  return is;
}

template <typename IntType>
_CCCL_HOST_DEVICE bool
operator==(const uniform_int_distribution<IntType>& lhs, const uniform_int_distribution<IntType>& rhs)
{
  return thrust::random::detail::random_core_access::equal(lhs, rhs);
}

template <typename IntType>
_CCCL_HOST_DEVICE bool
operator!=(const uniform_int_distribution<IntType>& lhs, const uniform_int_distribution<IntType>& rhs)
{
  return !(lhs == rhs);
}

template <typename IntType, typename CharT, typename Traits>
std::basic_ostream<CharT, Traits>&
operator<<(std::basic_ostream<CharT, Traits>& os, const uniform_int_distribution<IntType>& d)
{
  return thrust::random::detail::random_core_access::stream_out(os, d);
}

template <typename IntType, typename CharT, typename Traits>
std::basic_istream<CharT, Traits>&
operator>>(std::basic_istream<CharT, Traits>& is, uniform_int_distribution<IntType>& d)
{
  return thrust::random::detail::random_core_access::stream_in(is, d);
}
} // namespace random

THRUST_NAMESPACE_END
