/*******************************************************************************
 *
 * MIT License
 *
 * Copyright (c) 2022 Advanced Micro Devices, Inc.
 *
 * 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 "d_vector.hpp"

#include "device_vector.hpp"

#include "host_vector.hpp"

#include "hipsparselt_init.hpp"

//!
//! @brief enum to check for NaN initialization of the Input vector/matrix
//!
typedef enum hipsparselt_check_nan_init_
{
    // Alpha sets NaN
    hipsparselt_client_alpha_sets_nan,

    // Beta sets NaN
    hipsparselt_client_beta_sets_nan,

    //  Never set NaN
    hipsparselt_client_never_set_nan

} hipsparselt_check_nan_init;

//!
//! @brief Template for initializing a host (non_batched|batched|strided_batched)vector.
//! @param that That vector.
//! @param rand_gen The random number generator
//! @param seedReset Reset the seed if true, do not reset the seed otherwise.
//!
template <typename U, typename T>
void hipsparselt_init_template(U& that, T rand_gen(), bool seedReset, bool alternating_sign = false)
{
    if(seedReset)
        hipsparselt_seedrand();

    for(int64_t batch_index = 0; batch_index < that.batch_count(); ++batch_index)
    {
        auto*     batched_data = that[batch_index];
        ptrdiff_t inc          = that.inc();
        auto      n            = that.n();

        if(inc < 0)
            batched_data -= (n - 1) * inc;

        if(alternating_sign)
        {
            for(int64_t i = 0; i < n; ++i)
            {
                auto value            = rand_gen();
                batched_data[i * inc] = (i ^ 0) & 1 ? value : negate(value);
            }
        }
        else
        {
            for(int64_t i = 0; i < n; ++i)
                batched_data[i * inc] = rand_gen();
        }
    }
}

//!
//! @brief Initialize a host_vector with NaNs.
//! @param that The host_vector to be initialized.
//! @param seedReset reset he seed if true, do not reset the seed otherwise.
//!
template <typename T>
inline void hipsparselt_init_nan(host_vector<T>& that, bool seedReset = false)
{
    hipsparselt_init_template(that, random_nan_generator<T>, seedReset);
}

//!
//! @brief Initialize a host_vector.
//! @param that The host_vector.
//! @param seedReset reset the seed if true, do not reset the seed otherwise.
//!
template <typename T>
inline void hipsparselt_init(host_vector<T>& that, bool seedReset = false)
{
    if(seedReset)
        hipsparselt_seedrand();
    hipsparselt_init(that, that.size(), 1, 1);
}

//!
//! @brief Initialize a host_vector.
//! @param hx The host_vector.
//! @param arg Specifies the argument class.
//! @param N Length of the host vector.
//! @param incx Increment for the host vector.
//! @param stride_x Incement between the host vector.
//! @param batch_count number of instances in the batch.
//! @param nan_init Initialize vector with Nan's depending upon the hipsparselt_check_nan_init enum value.
//! @param seedReset reset the seed if true, do not reset the seed otherwise. Use init_cos if seedReset is true else use init_sin.
//! @param alternating_sign Initialize vector so adjacent entries have alternating sign.
//!
template <typename T>
inline void hipsparselt_init_vector(host_vector<T>&            hx,
                                    const Arguments&           arg,
                                    int64_t                    N,
                                    int64_t                    incx,
                                    int64_t                    stride_x,
                                    int64_t                    batch_count,
                                    hipsparselt_check_nan_init nan_init,
                                    bool                       seedReset        = false,
                                    bool                       alternating_sign = false)
{
    if(seedReset)
        hipsparselt_seedrand();

    if(nan_init == hipsparselt_client_alpha_sets_nan && hipsparselt_isnan(arg.alpha))
    {
        hipsparselt_init_nan(hx, 1, N, incx, stride_x, batch_count);
    }
    else if(nan_init == hipsparselt_client_beta_sets_nan && hipsparselt_isnan(arg.beta))
    {
        hipsparselt_init_nan(hx, 1, N, incx, stride_x, batch_count);
    }
    else if(arg.initialization == hipsparselt_initialization::hpl)
    {
        if(alternating_sign)
            hipsparselt_init_hpl_alternating_sign(hx, 1, N, incx, stride_x, batch_count);
        else
            hipsparselt_init_hpl(hx, 1, N, incx, stride_x, batch_count);
    }
    else if(arg.initialization == hipsparselt_initialization::rand_int)
    {
        if(alternating_sign)
            hipsparselt_init_alternating_sign(hx, 1, N, incx, stride_x, batch_count);
        else
            hipsparselt_init(hx, 1, N, incx, stride_x, batch_count);
    }
    else if(arg.initialization == hipsparselt_initialization::trig_float)
    {
        if(seedReset)
            hipsparselt_init_cos(hx, 1, N, incx, stride_x, batch_count);
        else
            hipsparselt_init_sin(hx, 1, N, incx, stride_x, batch_count);
    }
}

//!
//! @brief Initialize a host matrix.
//! @param hA The host matrix.
//! @param arg Specifies the argument class.
//! @param M Length of the host matrix.
//! @param N Length of the host matrix.
//! @param lda Leading dimension of the host matrix.
//! @param stride_A Incement between the host matrix.
//! @param batch_count number of instances in the batch.
//! @param nan_init Initialize matrix with Nan's depending upon the hipsparselt_check_nan_init enum value.
//! @param seedReset reset the seed if true, do not reset the seed otherwise. Use init_cos if seedReset is true else use init_sin.
//! @param alternating_sign Initialize matrix so adjacent entries have alternating sign.
//!
template <typename T>
inline void hipsparselt_init_matrix(host_vector<T>&            hA,
                                    const Arguments&           arg,
                                    int64_t                    M,
                                    int64_t                    N,
                                    int64_t                    lda,
                                    int64_t                    stride_A,
                                    int32_t                    batch_count,
                                    hipsparselt_check_nan_init nan_init,
                                    bool                       seedReset        = false,
                                    bool                       alternating_sign = false)
{
    if(seedReset)
        hipsparselt_seedrand();

    if(nan_init == hipsparselt_client_alpha_sets_nan && hipsparselt_isnan(arg.alpha))
    {
        hipsparselt_init_nan(hA, M, N, lda, stride_A, batch_count);
    }
    else if(nan_init == hipsparselt_client_beta_sets_nan && hipsparselt_isnan(arg.beta))
    {
        hipsparselt_init_nan(hA, M, N, lda, stride_A, batch_count);
    }
    else if(arg.initialization == hipsparselt_initialization::hpl)
    {
        if(alternating_sign)
            hipsparselt_init_hpl_alternating_sign(hA, M, N, lda, stride_A, batch_count);
        else
            hipsparselt_init_hpl(hA, M, N, lda, stride_A, batch_count);
    }
    else if(arg.initialization == hipsparselt_initialization::rand_int)
    {
        if(alternating_sign)
            hipsparselt_init_alternating_sign(hA, M, N, lda, stride_A, batch_count);
        else
            hipsparselt_init(hA, M, N, lda, stride_A, batch_count);
    }
    else if(arg.initialization == hipsparselt_initialization::trig_float)
    {
        if(seedReset)
            hipsparselt_init_cos(hA, M, N, lda, stride_A, batch_count);
        else
            hipsparselt_init_sin(hA, M, N, lda, stride_A, batch_count);
    }
}