/*
 * Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
 *
 * Permission is hereby granted, free of charge, to any person obtaining a
 * copy of this software and associated documentation files (the "Software"),
 * to deal in the Software without restriction, including without limitation
 * the rights to use, copy, modify, merge, publish, distribute, sublicense,
 * and/or sell copies of the Software, and to permit persons to whom the
 * Software is furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
 * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 * DEALINGS IN THE SOFTWARE.
 */

#pragma once

#include <algorithm>
#include <array>
#include <functional>
#include <memory>
#include <sstream>
#include <utility>

#include <cudnn.h>
#include <cudnn_backend.h>

#include "cudnn_frontend_utils.h"

namespace cudnn_frontend {

///
/// Rng Descriptor Class
/// This class tells the properties of the Rng operation
/// Properties:
///
/// Use RngDescBuilder_v8 to build this class.
/// Describe returns a string describing the Rng operation
///
class RngDesc_v8 : public BackendDescriptor {
   public:
    friend class RngDescBuilder_v8;
    std::string
    describe() const override {
        std::stringstream ss;
#if (CUDNN_VERSION >= 8700)
        ss  << "CUDNN_BACKEND_RNG_DESCRIPTOR: "
            << "Distribution Type: " << to_string(distribution)
            << ", Normal Distribution Mean: " << normal_dist_mean 
            << ", Normal Distribution Standard Deviation: " << normal_dist_std_dev 
            << ", Uniform Distribution Maximum: " << uniform_dist_max 
            << ", Uniform Distribution Minimum: " << uniform_dist_min 
            << ", Bernoulli Distribution Probability: " << bernoulli_dist_probability;
#endif
        return ss.str();
    }

    RngDesc_v8(RngDesc_v8 &&from) = default;
    RngDesc_v8 &
    operator=(RngDesc_v8 &&) = default;

    ~RngDesc_v8() = default;

    /** @defgroup RngDescBuilder_v8
     *  Get individual property of RngDesc_v8 class
     *  @{
     */
    
    double
    getNormalDistMean() const {
        return normal_dist_mean;
    }

    double
    getNormalDistStdDev() const {
        return normal_dist_std_dev;
    }

    double
    getUniformDistMax() const {
        return uniform_dist_max;
    }

    double
    getUniformDistMin() const {
        return normal_dist_std_dev;
    }

    double
    getBernoulliDistProbability() const {
        return bernoulli_dist_probability;
    }

#if (CUDNN_VERSION >= 8700)
    cudnnRngDistribution_t
    getDistribution() const {
        return distribution;
    }
#endif
    /** @} */

   private:

    RngDesc_v8()                    = default;
    RngDesc_v8(RngDesc_v8 const &) = delete;
    RngDesc_v8 &
    operator=(RngDesc_v8 const &) = delete;

    // default values for attributes 
    double normal_dist_mean = -1;
    double normal_dist_std_dev = -1;
    double uniform_dist_max = -1;
    double uniform_dist_min = -1;
    double bernoulli_dist_probability = -1;

#if (CUDNN_VERSION >= 8700)
    cudnnRngDistribution_t distribution = CUDNN_RNG_DISTRIBUTION_BERNOULLI;
#endif
    };

///
/// RngDescBuilder_v8 Class
/// Helper class used to build RngDesc_v8 class
class RngDescBuilder_v8 {
   public:
    /** @defgroup RngDescBuilder_v8
     *  Set individual property of RngDesc_v8 class
     *  @{
     */

#if (CUDNN_VERSION >= 8700)
    //! Set Rng distribution for the Rng Operation
    auto
    setRngDistribution(cudnnRngDistribution_t distribution_) -> RngDescBuilder_v8 & {
        m_RngDesc.distribution = distribution_;
        return *this;
    }
#endif

    //! Set normal distribution params (mean and std dev) for the Rng Operation
    auto
    setNormalDistParams(double normal_dist_mean_, double normal_dist_std_dev_) -> RngDescBuilder_v8 & {
        m_RngDesc.normal_dist_mean = normal_dist_mean_;
        m_RngDesc.normal_dist_std_dev = normal_dist_std_dev_;
        return *this;
    }
    
    //! Set normal distribution mean for the Rng Operation
    auto
    setNormalDistMean(double normal_dist_mean_) -> RngDescBuilder_v8 & {
        m_RngDesc.normal_dist_mean = normal_dist_mean_;
        return *this;
    }

    //! Set normal distribution std dev for the Rng Operation
    auto
    setNormalDistStdDev(double normal_dist_std_dev_) -> RngDescBuilder_v8 & {
        m_RngDesc.normal_dist_std_dev = normal_dist_std_dev_;
        return *this;
    }

    //! Set uniform distribution params (min and max) for the Rng Operation
    auto
    setUniformDistParams(double uniform_dist_max_, double uniform_dist_min_) -> RngDescBuilder_v8 & {
        m_RngDesc.uniform_dist_max = uniform_dist_max_;
        m_RngDesc.uniform_dist_min = uniform_dist_min_;
        return *this;
    }

    //! Set uniform distribution max for the Rng Operation
    auto
    setUniformDistMax(double uniform_dist_max_) -> RngDescBuilder_v8 & {
        m_RngDesc.uniform_dist_max = uniform_dist_max_;
        return *this;
    }

    //! Set uniform distribution min for the Rng Operation
    auto
    setUniformDistMin(double uniform_dist_min_) -> RngDescBuilder_v8 & {
        m_RngDesc.uniform_dist_min = uniform_dist_min_;
        return *this;
    }

    //! Set bernoulli distribution probability for the Rng Operation
    auto
    setBernoulliDistProbability(double bernoulli_dist_probability_) -> RngDescBuilder_v8 & {
        m_RngDesc.bernoulli_dist_probability = bernoulli_dist_probability_;
        return *this;
    }

    /** @} */

    //! constructs the RngDesc_v8 by calling the cudnn API
    //! Throws the appropriate error message
    RngDesc_v8 &&
    build() {
#if (CUDNN_VERSION >= 8700)
        // Create a descriptor. Memory allocation happens here.
        auto status = m_RngDesc.initialize_managed_backend_pointer(CUDNN_BACKEND_RNG_DESCRIPTOR);
        if (status != CUDNN_STATUS_SUCCESS) {
            set_error_and_throw_exception(
                &m_RngDesc, status, "CUDNN_BACKEND_RNG_DESCRIPTOR: cudnnCreate Failed");
            return std::move(m_RngDesc);
        }

        // Once Created lets set the descriptor parameters.
        status = cudnnBackendSetAttribute(m_RngDesc.pointer->get_backend_descriptor(), 
                                          CUDNN_ATTR_RNG_DISTRIBUTION, 
                                          CUDNN_TYPE_RNG_DISTRIBUTION, 
                                          1,
                                          &(m_RngDesc.distribution));
        if (status != CUDNN_STATUS_SUCCESS) {
            set_error_and_throw_exception(
                &m_RngDesc,
                status,
                "CUDNN_BACKEND_RNG_DESCRIPTOR: SetAttribute CUDNN_ATTR_RNG_DISTRIBUTION Failed");
            return std::move(m_RngDesc);
        }

        status = cudnnBackendSetAttribute(m_RngDesc.pointer->get_backend_descriptor(), 
                                          CUDNN_ATTR_RNG_NORMAL_DIST_MEAN, 
                                          CUDNN_TYPE_DOUBLE, 
                                          1,
                                          &(m_RngDesc.normal_dist_mean));
        if (status != CUDNN_STATUS_SUCCESS) {
            set_error_and_throw_exception(
                &m_RngDesc,
                status,
                "CUDNN_BACKEND_RNG_DESCRIPTOR: SetAttribute CUDNN_ATTR_RNG_NORMAL_DIST_MEAN Failed");
            return std::move(m_RngDesc);
        }

        status = cudnnBackendSetAttribute(m_RngDesc.pointer->get_backend_descriptor(), 
                                          CUDNN_ATTR_RNG_NORMAL_DIST_STANDARD_DEVIATION, 
                                          CUDNN_TYPE_DOUBLE, 
                                          1,
                                          &(m_RngDesc.normal_dist_std_dev));
        if (status != CUDNN_STATUS_SUCCESS) {
            set_error_and_throw_exception(
                &m_RngDesc,
                status,
                "CUDNN_BACKEND_RNG_DESCRIPTOR: SetAttribute CUDNN_ATTR_RNG_NORMAL_DIST_STANDARD_DEVIATION Failed");
            return std::move(m_RngDesc);
        }

        status = cudnnBackendSetAttribute(m_RngDesc.pointer->get_backend_descriptor(), 
                                          CUDNN_ATTR_RNG_UNIFORM_DIST_MAXIMUM, 
                                          CUDNN_TYPE_DOUBLE, 
                                          1,
                                          &(m_RngDesc.uniform_dist_max));
        if (status != CUDNN_STATUS_SUCCESS) {
            set_error_and_throw_exception(
                &m_RngDesc,
                status,
                "CUDNN_BACKEND_RNG_DESCRIPTOR: SetAttribute CUDNN_ATTR_RNG_UNIFORM_DIST_MAXIMUM Failed");
            return std::move(m_RngDesc);
        }

        status = cudnnBackendSetAttribute(m_RngDesc.pointer->get_backend_descriptor(), 
                                          CUDNN_ATTR_RNG_UNIFORM_DIST_MINIMUM, 
                                          CUDNN_TYPE_DOUBLE, 
                                          1,
                                          &(m_RngDesc.uniform_dist_min));
        if (status != CUDNN_STATUS_SUCCESS) {
            set_error_and_throw_exception(
                &m_RngDesc,
                status,
                "CUDNN_BACKEND_RNG_DESCRIPTOR: SetAttribute CUDNN_ATTR_RNG_UNIFORM_DIST_MINIMUM Failed");
            return std::move(m_RngDesc);
        }

        status = cudnnBackendSetAttribute(m_RngDesc.pointer->get_backend_descriptor(), 
                                          CUDNN_ATTR_RNG_BERNOULLI_DIST_PROBABILITY, 
                                          CUDNN_TYPE_DOUBLE, 
                                          1,
                                          &(m_RngDesc.bernoulli_dist_probability));
        if (status != CUDNN_STATUS_SUCCESS) {
            set_error_and_throw_exception(
                &m_RngDesc,
                status,
                "CUDNN_BACKEND_RNG_DESCRIPTOR: SetAttribute CUDNN_ATTR_RNG_BERNOULLI_DIST_PROBABILITY Failed");
            return std::move(m_RngDesc);
        }

        // Finalizing the descriptor
        status = cudnnBackendFinalize(m_RngDesc.pointer->get_backend_descriptor());
        if (status != CUDNN_STATUS_SUCCESS) {
            set_error_and_throw_exception(
                &m_RngDesc, status, "CUDNN_BACKEND_RNG_DESCRIPTOR: cudnnFinalize Failed");
            return std::move(m_RngDesc);
        }
        getLogger() << "[cudnn_frontend] " << m_RngDesc << std::endl;
        return std::move(m_RngDesc);
#else 
    set_error_and_throw_exception(&m_RngDesc,
                                    CUDNN_STATUS_NOT_SUPPORTED,
                                    "CUDNN_RNG_DESCRIPTOR: Rng only supported in cuDNN v8.7 or later");
    return std::move(m_RngDesc);
#endif
    }

    explicit RngDescBuilder_v8()                  = default;
    ~RngDescBuilder_v8()                          = default;
    RngDescBuilder_v8(RngDescBuilder_v8 &&)      = delete;
    RngDescBuilder_v8(RngDescBuilder_v8 const &) = delete;
    RngDescBuilder_v8 &
    operator=(RngDescBuilder_v8 const &) = delete;

   private:
    RngDesc_v8 m_RngDesc;
};
}
