/*******************************************************************************
* Copyright 2020-2021 Intel Corporation.
*
* This software and the related documents are Intel copyrighted  materials,  and
* your use of  them is  governed by the  express license  under which  they were
* provided to you (License).  Unless the License provides otherwise, you may not
* use, modify, copy, publish, distribute,  disclose or transmit this software or
* the related documents without Intel's prior written permission.
*
* This software and the related documents  are provided as  is,  with no express
* or implied  warranties,  other  than those  that are  expressly stated  in the
* License.
*******************************************************************************/

#ifndef _MKL_RNG_DEVICE_DISTRIBUTIONS_HPP_
#define _MKL_RNG_DEVICE_DISTRIBUTIONS_HPP_

#include <limits>

#include "oneapi/mkl/rng/device/detail/distribution_base.hpp"
#include "oneapi/mkl/rng/device/functions.hpp"

namespace oneapi {
namespace mkl {
namespace rng {
namespace device {

// CONTINUOUS AND DISCRETE RANDOM NUMBER DISTRIBUTIONS

// Class template oneapi::mkl::rng::device::uniform
//
// Represents continuous and discrete uniform random number distribution
//
// Supported types:
//      float
//      double
//      std::int32_t
//      std::uint32_t
//
// Supported methods:
//      oneapi::mkl::rng::device::uniform_method::standard
//      oneapi::mkl::rng::device::uniform_method::accurate
//
// Input arguments:
//      a - left bound. 0.0 by default
//      b - right bound. 1.0 by default (for std::(u)int32_t std::numeric_limits<std::int32_t>::max()
//          is used for accurate method and 2^23 is used for standard method)
//
// Note: using (un)signed integer uniform distribution with uniform_method::standard method may
// cause incorrect statistics of the produced random numbers (due to rounding error) if 
// (abs(b - a) > 2^23) || (abs(b) > 2^23) || (abs(a) > 2^23)
// Please use uniform_method::accurate method instead
//
template <typename Type, typename Method>
class uniform : detail::distribution_base<uniform<Type, Method>> {
public:
    static_assert(std::is_same<Method, uniform_method::standard>::value ||
                      std::is_same<Method, uniform_method::accurate>::value,
                  "oneMKL: rng/uniform: method is incorrect");

    static_assert(std::is_same<Type, float>::value || std::is_same<Type, double>::value ||
                      std::is_same<Type, std::int32_t>::value || std::is_same<Type, std::uint32_t>::value,
                  "oneMKL: rng/uniform: type is not supported");

    using method_type = Method;
    using result_type = Type;
    using param_type = typename detail::distribution_base<uniform<Type, Method>>::param_type;

    uniform()
            : detail::distribution_base<uniform<Type, Method>>(static_cast<Type>(0.0),
                    std::is_integral<Type>::value ? (std::is_same<Method,
                    uniform_method::standard>::value ? (1 << 23) :
                    std::numeric_limits<Type>::max()): static_cast<Type>(1.0)) {}

    explicit uniform(Type a, Type b) : detail::distribution_base<uniform<Type, Method>>(a, b) {}
    explicit uniform(const param_type& pt) : detail::distribution_base<uniform<Type, Method>>(pt.a_, pt.b_) {}

    Type a() const {
        return detail::distribution_base<uniform<Type, Method>>::a();
    }

    Type b() const {
        return detail::distribution_base<uniform<Type, Method>>::b();
    }

    param_type param() const {
        return detail::distribution_base<uniform<Type, Method>>::param();
    }

    void param(const param_type& pt) {
        detail::distribution_base<uniform<Type, Method>>::param(pt);
    }

private:
    template <typename Distr, typename Engine>
    friend auto generate(Distr& distr, Engine& engine) ->
        typename std::conditional<Engine::vec_size == 1, typename Distr::result_type,
                                  sycl::vec<typename Distr::result_type, Engine::vec_size>>::type;

    template <typename Distr, typename Engine>
    friend typename Distr::result_type generate_single(Distr& distr, Engine& engine);
};

// Class template oneapi::mkl::rng::device::gaussian
//
// Represents continuous normal random number distribution
//
// Supported types:
//      float
//      double
//
// Supported methods:
//      oneapi::mkl::rng::device::gaussian_method::box_muller2
//
// Input arguments:
//      mean   - mean. 0 by default
//      stddev - standard deviation. 1.0 by default
//
template <typename RealType, typename Method>
class gaussian : detail::distribution_base<gaussian<RealType, Method>> {
public:
    static_assert(std::is_same<Method, gaussian_method::box_muller2>::value,
                  "oneMKL: rng/gaussian: method is incorrect");

    static_assert(std::is_same<RealType, float>::value || std::is_same<RealType, double>::value,
                  "oneMKL: rng/gaussian: type is not supported");

    using method_type = Method;
    using result_type = RealType;
    using param_type = typename detail::distribution_base<gaussian<RealType, Method>>::param_type;

    gaussian()
            : detail::distribution_base<gaussian<RealType, Method>>(static_cast<RealType>(0.0),
                                                                    static_cast<RealType>(1.0)) {}

    explicit gaussian(RealType mean, RealType stddev)
            : detail::distribution_base<gaussian<RealType, Method>>(mean, stddev) {}
    explicit gaussian(const param_type& pt)
            : detail::distribution_base<gaussian<RealType, Method>>(pt.mean_, pt.stddev_) {}

    RealType mean() const {
        return detail::distribution_base<gaussian<RealType, Method>>::mean();
    }

    RealType stddev() const {
        return detail::distribution_base<gaussian<RealType, Method>>::stddev();
    }

    param_type param() const {
        return detail::distribution_base<gaussian<RealType, Method>>::param();
    }

    void param(const param_type& pt) {
        detail::distribution_base<gaussian<RealType, Method>>::param(pt);
    }

    template <typename Distr, typename Engine>
    friend auto generate(Distr& distr, Engine& engine) ->
        typename std::conditional<Engine::vec_size == 1, typename Distr::result_type,
                                  sycl::vec<typename Distr::result_type, Engine::vec_size>>::type;
    template <typename Distr, typename Engine>
    friend typename Distr::result_type generate_single(Distr& distr, Engine& engine);
};

// Class template oneapi::mkl::rng::device::lognormal
//
// Represents continuous lognormal random number distribution
//
// Supported types:
//      float
//      double
//
// Supported methods:
//      oneapi::mkl::rng::device::lognormal_method::box_muller2
//
// Input arguments:
//      m     - mean of the subject normal distribution. 0.0 by default
//      s     - standard deviation of the subject normal distribution. 1.0 by default
//      displ - displacement. 0.0 by default
//      scale - scalefactor. 1.0 by default
//
template <typename RealType, typename Method>
class lognormal : detail::distribution_base<lognormal<RealType, Method>> {
public:
    static_assert(std::is_same<Method, lognormal_method::box_muller2>::value,
                  "oneMKL: rng/lognormal: method is incorrect");

    static_assert(std::is_same<RealType, float>::value || std::is_same<RealType, double>::value,
                  "oneMKL: rng/lognormal: type is not supported");

    using method_type = Method;
    using result_type = RealType;
    using param_type = typename detail::distribution_base<lognormal<RealType, Method>>::param_type;

    lognormal()
            : detail::distribution_base<lognormal<RealType, Method>>(
                  static_cast<RealType>(0.0), static_cast<RealType>(1.0),
                  static_cast<RealType>(0.0), static_cast<RealType>(1.0)) {}

    explicit lognormal(RealType m, RealType s, RealType displ = static_cast<RealType>(0.0),
                       RealType scale = static_cast<RealType>(1.0))
            : detail::distribution_base<lognormal<RealType, Method>>(m, s, displ, scale) {}
    explicit lognormal(const param_type& pt)
            : detail::distribution_base<lognormal<RealType, Method>>(pt.m_, pt.s_, pt.displ_, pt.scale_) {}

    RealType m() const {
        return detail::distribution_base<lognormal<RealType, Method>>::m();
    }

    RealType s() const {
        return detail::distribution_base<lognormal<RealType, Method>>::s();
    }

    RealType displ() const {
        return detail::distribution_base<lognormal<RealType, Method>>::displ();
    }

    RealType scale() const {
        return detail::distribution_base<lognormal<RealType, Method>>::scale();
    }

    param_type param() const {
        return detail::distribution_base<lognormal<RealType, Method>>::param();
    }

    void param(const param_type& pt) {
        detail::distribution_base<lognormal<RealType, Method>>::param(pt);
    }

    template <typename Distr, typename Engine>
    friend auto generate(Distr& distr, Engine& engine) ->
        typename std::conditional<Engine::vec_size == 1, typename Distr::result_type,
                                  sycl::vec<typename Distr::result_type, Engine::vec_size>>::type;
    template <typename Distr, typename Engine>
    friend typename Distr::result_type generate_single(Distr& distr, Engine& engine);
};

// Class template oneapi::mkl::rng::device::bits
//
// Represents bits of underlying random number engine
//
// Supported types:
//      std::uint32_t
//
template <typename UIntType>
class bits : detail::distribution_base<bits<UIntType>> {
public:
    static_assert(std::is_same<UIntType, std::uint32_t>::value,
                  "oneMKL: rng/bits: type is not supported");
    using result_type = UIntType;

private:
    template <typename Distr, typename Engine>
    friend auto generate(Distr& distr, Engine& engine) ->
        typename std::conditional<Engine::vec_size == 1, typename Distr::result_type,
                                  sycl::vec<typename Distr::result_type, Engine::vec_size>>::type;

    template <typename Distr, typename Engine>
    friend typename Distr::result_type generate_single(Distr& distr, Engine& engine);
};

// Class template oneapi::mkl::rng::device::exponential
//
// Represents continuous exponential random number distribution
//
// Supported types:
//      float
//      double
//
// Supported methods:
//      oneapi::mkl::rng::device::exponential_method::icdf
//      oneapi::mkl::rng::device::exponential_method::icdf_accurate
//
// Input arguments:
//      displ - displacement. 0.0 by default
//      scale - scalefactor. 1.0 by default
//
template <typename RealType, typename Method>
class exponential : detail::distribution_base<exponential<RealType, Method>> {
public:
    static_assert(std::is_same<Method, exponential_method::icdf>::value ||
                  std::is_same<Method, exponential_method::icdf_accurate>::value,
                  "oneMKL: rng/exponential: method is incorrect");

    static_assert(std::is_same<RealType, float>::value || std::is_same<RealType, double>::value,
                  "oneMKL: rng/exponential: type is not supported");

    using method_type = Method;
    using result_type = RealType;
    using param_type = typename detail::distribution_base<exponential<RealType, Method>>::param_type;

    exponential()
            : detail::distribution_base<exponential<RealType, Method>>(
                  static_cast<RealType>(0.0), static_cast<RealType>(1.0)) {}

    explicit exponential(RealType a, RealType beta)
            : detail::distribution_base<exponential<RealType, Method>>(a, beta) {}

    explicit exponential(const param_type& pt)
            : detail::distribution_base<exponential<RealType, Method>>(pt.a_, pt.beta_) {}

    RealType a() const {
        return detail::distribution_base<exponential<RealType, Method>>::a();
    }

    RealType beta() const {
        return detail::distribution_base<exponential<RealType, Method>>::beta();
    }

    param_type param() const {
        return detail::distribution_base<exponential<RealType, Method>>::param();
    }

    void param(const param_type& pt) {
        detail::distribution_base<exponential<RealType, Method>>::param(pt);
    }

    template <typename Distr, typename Engine>
    friend auto generate(Distr& distr, Engine& engine) ->
        typename std::conditional<Engine::vec_size == 1, typename Distr::result_type,
                                  sycl::vec<typename Distr::result_type, Engine::vec_size>>::type;
    template <typename Distr, typename Engine>
    friend typename Distr::result_type generate_single(Distr& distr, Engine& engine);
};

// Class template oneapi::mkl::rng::device::poisson
//
// Represents discrete poisson random number distribution
//
// Supported types:
//      std::int32_t
//      std::uint32_t
//
// Supported methods:
//      oneapi::mkl::rng::device::poisson_method::devroye
//
// Input arguments:
//      lambda - mean value. 1.0 by default
//
template <typename IntType, typename Method>
class poisson : detail::distribution_base<poisson<IntType, Method>> {
public:
    static_assert(std::is_same<Method, poisson_method::devroye>::value,
                  "oneMKL: rng/poisson: method is incorrect");

    static_assert(std::is_same<IntType, std::int32_t>::value || std::is_same<IntType, std::uint32_t>::value,
                  "oneMKL: rng/poisson: type is not supported");

    using method_type = Method;
    using result_type = IntType;
    using param_type = typename detail::distribution_base<poisson<IntType, Method>>::param_type;

    poisson()
            : detail::distribution_base<poisson<IntType, Method>>(0.5) {}

    explicit poisson(double lambda)
            : detail::distribution_base<poisson<IntType, Method>>(lambda) {}
    explicit poisson(const param_type& pt)
            : detail::distribution_base<poisson<IntType, Method>>(pt.lambda_) {}

    double lambda() const {
        return detail::distribution_base<poisson<IntType, Method>>::lambda();
    }

    param_type param() const {
        return detail::distribution_base<poisson<IntType, Method>>::param();
    }

    void param(const param_type& pt) {
        detail::distribution_base<poisson<IntType, Method>>::param(pt);
    }

    template <typename Distr, typename Engine>
    friend auto generate(Distr& distr, Engine& engine) ->
        typename std::conditional<Engine::vec_size == 1, typename Distr::result_type,
                                  sycl::vec<typename Distr::result_type, Engine::vec_size>>::type;
    template <typename Distr, typename Engine>
    friend typename Distr::result_type generate_single(Distr& distr, Engine& engine);
};

} // namespace device
} // namespace rng
} // namespace mkl
} // namespace oneapi

#endif // _MKL_RNG_DEVICE_DISTRIBUTIONS_HPP_
