/**
 * @brief CUDA optimizing strategy.
 *
 * @author Miao Zhang
 *
 * @copyright Copyright (c) 2024, TBPLaS develop team.
 */
#ifdef WITH_CUDA
#ifndef CUDA_MATH_H
#define CUDA_MATH_H
#include <stdio.h>
#include <cuda_runtime.h>
#include <cuComplex.h>
#include "cudaConfig.h"
#include "../base/utils.h"
#include "../base/datatypes.h"
using namespace std;
using namespace tbplas::base;


/**
 * @brief Calculate Y = a * mat * X
 *
 * @tparam sa_t datatype of factor a
 * @tparam vx_t datatype for vector x
 * @tparam vy_t datatype for vector y
 * @param[in] a the factor
 * @param[in] x incoming vector
 * @param[out] y result vector
 */
__global__ void amv_kernel(
    const cuDoubleComplex a,
    const cuDoubleComplex* x, const int x_size,
    const uint64_t* indptr_, const uint64_t* indices_, const cuDoubleComplex* values_,
    cuDoubleComplex* y, const int y_size);

/**
 * @brief Calculate Y = mat * X
 *
 * @tparam vx_t datatype for vector x
 * @tparam vy_t datatype for vector y
 * @param[in] x incoming vector
 * @param[out] y result vector
 */
__global__ void mv_kernel(
    const cuDoubleComplex* x, const int x_size,
    const uint64_t* indptr_, const uint64_t* indices_, const cuDoubleComplex* values_,
    cuDoubleComplex* y, const int y_size);

/**
 * @brief Calculate OUT = a * mat * X - Y
 *
 * @tparam sa_t datatype of factor a
 * @tparam vx_t datatype for vector x
 * @tparam vy_t datatype for vector y
 * @tparam vout_t datatype for vector out
 * @param[in] a the factor
 * @param[in] x incoming vector y
 * @param[in] y incoming vector y
 * @param[out] out result vector
 */
__global__ void amxsy_out_kernel(
    const cuDoubleComplex a,
    const cuDoubleComplex* x, const cuDoubleComplex* y,
    const uint64_t* indptr, const uint64_t* indices, const cuDoubleComplex* values,
    cuDoubleComplex* out, const int dim);

/**
 * @brief Calculate OUT = a * X + b * Y + c * Z
 *
 * @tparam sa_t datatype of scalar a
 * @tparam vx_t datatype of vector x
 * @tparam sb_t datatype of scalar b
 * @tparam vy_y datatype of vector y
 * @tparam sc_t datatype of scalar c
 * @tparam vz_t datatype of vector z
 * @tparam vout_t datatype of vector out
 * @param[in] a scalar a
 * @param[in] x vector x
 * @param[in] b scalar b
 * @param[in] y vector y
 * @param[in] c scalar c
 * @param[in] z vector z
 * @param[out] out vector out
 */
__global__ void axpbypcz_out_kernel(
    const cuDoubleComplex a, const cuDoubleComplex* x, const int x_size,
    const cuDoubleComplex b, const cuDoubleComplex* y, const int y_size,
    const cuDoubleComplex c, const cuDoubleComplex* z, const int z_size,
    cuDoubleComplex* out, int dim);

/**
 * @brief  Calculate Y = a * mat * X - Y
 *
 * @tparam sa_t datatype of factor a
 * @tparam vx_t datatype for vector x
 * @tparam vy_t datatype for vector y
 * @param[in] a the factor
 * @param[in] x incoming vector
 * @param[inout] y result vector
 */
__global__ void amxsy_kernel(
    const double a,
    const cuDoubleComplex* x, const int x_size,
    cuDoubleComplex* y, const int y_size,
    const int dim,
    const uint64_t* indptr, const uint64_t* indices, const cuDoubleComplex* values);

/**
 * @brief Calculate Y = a * X + Y
 *
 * @tparam sa_t datatype of scalar a
 * @tparam vx_t datatype of vector x
 * @tparam vy_t datatype of vector y
 * @param[in] a scalar a
 * @param[in] x vector x
 * @param[inout] y vector y containing the result
 */
__global__ void axpy_kernel(
    const cuDoubleComplex a, const cuDoubleComplex* x,
    const int dim,
    cuDoubleComplex* y
);

__global__ void axpy_kernel(
    const double a, const cuDoubleComplex* x,
    const int dim,
    cuDoubleComplex* y
);
/**
 * @brief Calculate prod = x.conj() * mat * y
 *
 * @tparam vx_t datatype of vector x
 * @tparam vy_t datatype of vector y
 * @param[in] x vector x
 * @param[in] y vector y
 * @return the product
 */
__global__ void xmy_kernel(const cuDoubleComplex* x, const int x_size,
    const cuDoubleComplex* y, const int y_size, const int dim,
    const uint64_t* indptr, const uint64_t* indices, const cuDoubleComplex* values,
    cuDoubleComplex* product);

/**
 * @brief Calculate the inner product between complex vectors
 *
 * @param[in]  x left operand
 * @param[in]  y right operand
 * @param[out] out the inner product
 *
 * @note The conjugate of left operand is taken automatically.
 */
__global__ void vec_dot_kernel(const cuDoubleComplex* x,
    const cuDoubleComplex* y, const int size, cuDoubleComplex* out);

__global__ void vec_dot_kernel2(
    const cuDoubleComplex* buffer_D,
    cuDoubleComplex** wf_vb_tn_D,
    const size_t  row,
    const size_t  buffer_size,
    const size_t  num_kernel,
    const bool isRowMajor, cuDoubleComplex* mu_sample_D);

/**
 * @brief Calculate the norm^2 of vector
 *
 * @param[in] x the vector
 * @return out of vector
 */
__global__ void vec_norm2_kernel(const cuDoubleComplex* x, 
    const int size, cuDoubleComplex* out);

/**
 * @brief Calculate X = X / a
 *
 * @tparam vx_t datatype of vector x
 * @tparam sa_t datatype of scalar a
 * @param[inout] x left operand
 * @param[in] a right operand
 */
__global__ void vec_div_sca_kernel(
    cuDoubleComplex* x,
    const int size,
    const double a);

/**
 * @brief Calculate prod = x.conj() * m * y for diagonal matrix m
 *
 * @tparam vx_t datatype of vector x
 * @tparam vm_t datatype of diagonal matrix m
 * @tparam vy_t datatype of vector y
 * @param[in] x vector x
 * @param[in] m diagonal matrix m
 * @param[in] y vector y
 * @return the product
 */
__global__ void xmy_diag_kernel(
    const cuDoubleComplex* x, const cuDoubleComplex* y,
    const cuDoubleComplex* m, const int dim,cuDoubleComplex* product);

/**
 * @brief Calculate Z = a * X + b * Y + Z
 *
 * @tparam sa_t datatype of scalar a
 * @tparam vx_t datatype of vector x
 * @tparam sb_t datatype of scalar b
 * @tparam vy_y datatype of vector y
 * @tparam vz_t datatype of vector z
 * @param[in] a scalar a
 * @param[in] x vector x
 * @param[in] b scalar b
 * @param[in] y vector y
 * @param[inout] z vector z
 */
__global__ void axpbypz_kernel(
    const cuDoubleComplex a, const cuDoubleComplex* x,
    const cuDoubleComplex b, const cuDoubleComplex* y, 
    const int dim,
    cuDoubleComplex* z);

__global__ void amxpy_out_gpu(
    const cuDoubleComplex a,
    const cuDoubleComplex* x, const int x_size,
    const cuDoubleComplex* y, const int y_size,
    const uint64_t* indptr, const uint64_t* indices, const cuDoubleComplex* values,
    cuDoubleComplex* out, const int out_size, const int dim);


__global__ void amxpy_gpu(
    const cuDoubleComplex a,
    const cuDoubleComplex* x, const int x_size,
    cuDoubleComplex* y, const int y_size,
    const int dim,
    const uint64_t* indptr, const uint64_t* indices, const cuDoubleComplex* values);


#endif //CUDA_MATH_H
#endif //WITH_CUDA


