/*
 * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: LicenseRef-NvidiaProprietary
 *
 * NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
 * property and proprietary rights in and to this material, related
 * documentation and any modifications thereto. Any use, reproduction,
 * disclosure or distribution of this material and related documentation
 * without an express license agreement from NVIDIA CORPORATION or
 * its affiliates is strictly prohibited.
 */

#pragma once

#include <complex>
#include <string_view>
#include <string>
#include "cuda_runtime.h"
#include "cublas_v2.h"
#include "cuComplex.h"

// Alignment requirement for scalar variables. Particularly, cuDoubleComplex variables need to be 16-byte aligned
#define __SCALAR_ALIGNAS(__data_type) alignas(sizeof(__data_type)) __data_type

namespace __cublas_std {

inline thread_local cublasHandle_t __handle = nullptr;
inline thread_local cudaStream_t   __stream = nullptr;
inline thread_local float        * __dummy  = nullptr;

struct __handle_cleanup
{
    bool __handle_created;

    ~__handle_cleanup()
    {
        if ( __handle != nullptr )
        {
            cudaFreeAsync( __dummy, __stream );

            cublasDestroy( __handle );
            __handle = nullptr;
        }
    }
};

inline thread_local __handle_cleanup __cleaner;

inline cublasHandle_t __get_cublas_handle()
{
    if ( __handle == nullptr )
    {
        cublasStatus_t __stat = cublasCreate( &__handle );

        if ( __stat != CUBLAS_STATUS_SUCCESS )
        {
            throw std::system_error( std::make_error_code( __stat ), "cublasCreate failed" );
        }

        __cleaner.__handle_created = true;

        cublasGetStream( __handle, &__stream );

        // Pre-allocate some space so that the pool allocator doesn't return memory to the OS
        cudaMallocAsync( (void**)&__dummy, 8, __stream );
    }

    return __handle;
}

template< class _Exec >
inline void __synchronize( _Exec ) {}

inline void __synchronize( __nvhpc_std::__nvhpc_exec< __nvhpc_std::__cublas_exec_space< __nvhpc_std::__nvhpc_sync > > )
{
    cudaError_t __error( cudaStreamSynchronize( __stream ) );

    if ( __error != cudaSuccess )
    {
        throw std::system_error( std::make_error_code( __error ), "CUDA call to cudaStreamSynchronize failed" );
    }
}

inline void __check_cublas_status( cublasStatus_t __stat, std::string_view __function_name, std::string_view __cublas_api_name )
{
    if ( __stat != CUBLAS_STATUS_SUCCESS )
    {
        throw std::system_error( std::make_error_code( __stat )
                               ,   std::string( "cuBLAS call to '" )
                                 + std::string( __cublas_api_name )
                                 + std::string( "' in function '" )
                                 + std::string( __function_name )
                                 + std::string( "' failed" )
                               );
    }
}

template <typename> constexpr cudaDataType_t __cuda_data_type = cudaDataType_t(-1);

#define __DEFINE_CUDA_TYPE_MAPPING(__type, __value) \
  template <> constexpr cudaDataType_t __cuda_data_type<__type> = CUDA_R_##__value; \
  template <> constexpr cudaDataType_t __cuda_data_type<std::complex<__type>> = CUDA_C_##__value

__DEFINE_CUDA_TYPE_MAPPING(float      , 32F);
__DEFINE_CUDA_TYPE_MAPPING(signed char, 8I );
#if __SIGNED_CHARS__
__DEFINE_CUDA_TYPE_MAPPING(char       , 8I );
#else
__DEFINE_CUDA_TYPE_MAPPING(char       , 8U );
#endif

#undef __DEFINE_CUDA_TYPE_MAPPING

template <typename> constexpr cublasComputeType_t __cublas_compute_type = cublasComputeType_t(-1);

template <> constexpr cublasComputeType_t __cublas_compute_type<float> = CUBLAS_COMPUTE_32F;

// BLAS Level 1

// Copy

inline cublasStatus_t __cublas_copy( cublasHandle_t               __handle
                                   , int                          __n
                                   , float                const * __x
                                   , int                          __incx
                                   , float                      * __y
                                   , int                          __incy
                                   )
{
    return cublasScopy( __handle, __n, __x, __incx, __y, __incy );
}

inline cublasStatus_t __cublas_copy( cublasHandle_t               __handle
                                   , int                          __n
                                   , double               const * __x
                                   , int                          __incx
                                   , double                     * __y
                                   , int                          __incy
                                   )
{
    return cublasDcopy( __handle, __n, __x, __incx, __y, __incy );
}

inline cublasStatus_t __cublas_copy( cublasHandle_t               __handle
                                   , int                          __n
                                   , std::complex<float > const * __x
                                   , int                          __incx
                                   , std::complex<float >       * __y
                                   , int                          __incy
                                   )
{
    return cublasCcopy( __handle, __n, (cuComplex const *)__x, __incx, (cuComplex *)__y, __incy );
}

inline cublasStatus_t __cublas_copy( cublasHandle_t               __handle
                                   , int                          __n
                                   , std::complex<double> const * __x
                                   , int                          __incx
                                   , std::complex<double>       * __y
                                   , int                          __incy
                                   )
{
    return cublasZcopy( __handle, __n, (cuDoubleComplex const *)__x, __incx, (cuDoubleComplex *)__y, __incy );
}

// Axpy

inline cublasStatus_t __cublas_axpy( cublasHandle_t               __handle
                                   , int                          __n
                                   , float                const * __alpha
                                   , float                const * __x
                                   , int                          __incx
                                   , float                      * __y
                                   , int                          __incy
                                   )
{
    return cublasSaxpy( __handle, __n, __alpha, __x, __incx, __y, __incy );
}

inline cublasStatus_t __cublas_axpy( cublasHandle_t               __handle
                                   , int                          __n
                                   , double               const * __alpha
                                   , double               const * __x
                                   , int                          __incx
                                   , double                     * __y
                                   , int                          __incy
                                   )
{
    return cublasDaxpy( __handle, __n, __alpha, __x, __incx, __y, __incy );
}

inline cublasStatus_t __cublas_axpy( cublasHandle_t               __handle
                                   , int                          __n
                                   , std::complex<float > const * __alpha
                                   , std::complex<float > const * __x
                                   , int                          __incx
                                   , std::complex<float >       * __y
                                   , int                          __incy
                                   )
{
    return cublasCaxpy( __handle, __n, (cuComplex const *)__alpha, (cuComplex const *)__x, __incx, (cuComplex *)__y, __incy );
}

inline cublasStatus_t __cublas_axpy( cublasHandle_t               __handle
                                   , int                          __n
                                   , std::complex<double> const * __alpha
                                   , std::complex<double> const * __x
                                   , int                          __incx
                                   , std::complex<double>       * __y
                                   , int                          __incy
                                   )
{
    return cublasZaxpy( __handle, __n, (cuDoubleComplex const *)__alpha, (cuDoubleComplex const *)__x, __incx, (cuDoubleComplex *)__y, __incy );
}

// Dot

inline cublasStatus_t __cublas_dot( cublasHandle_t               __handle
                                  , int                          __n
                                  , float                const * __x
                                  , int                          __incx
                                  , float                const * __y
                                  , int                          __incy
                                  , float                      * __result
                                  )
{
    return cublasSdot( __handle, __n, __x, __incx, __y, __incy, __result );
}

inline cublasStatus_t __cublas_dot( cublasHandle_t               __handle
                                  , int                          __n
                                  , double               const * __x
                                  , int                          __incx
                                  , double               const * __y
                                  , int                          __incy
                                  , double                     * __result
                                  )
{
    return cublasDdot( __handle, __n, __x, __incx, __y, __incy, __result );
}

inline cublasStatus_t __cublas_dot( cublasHandle_t               __handle
                                  , int                          __n
                                  , std::complex<float > const * __x
                                  , int                          __incx
                                  , std::complex<float > const * __y
                                  , int                          __incy
                                  , std::complex<float >       * __result
                                  )
{
    return cublasCdotu( __handle, __n, (cuComplex*)__x, __incx, (cuComplex*)__y, __incy, (cuComplex*)__result );
}

inline cublasStatus_t __cublas_dot( cublasHandle_t               __handle
                                  , int                          __n
                                  , std::complex<double> const * __x
                                  , int                          __incx
                                  , std::complex<double> const * __y
                                  , int                          __incy
                                  , std::complex<double>       * __result
                                  )
{
    return cublasZdotu( __handle, __n, (cuDoubleComplex*)__x, __incx, (cuDoubleComplex*)__y, __incy, (cuDoubleComplex*)__result );
}

inline cublasStatus_t __cublas_dotc( cublasHandle_t               __handle
                                   , int                          __n
                                   , std::complex<float > const * __x
                                   , int                          __incx
                                   , std::complex<float > const * __y
                                   , int                          __incy
                                   , std::complex<float >       * __result
                                   )
{
    return cublasCdotc( __handle, __n, (cuComplex*)__x, __incx, (cuComplex*)__y, __incy, (cuComplex*)__result );
}

inline cublasStatus_t __cublas_dotc( cublasHandle_t               __handle
                                   , int                          __n
                                   , std::complex<double> const * __x
                                   , int                          __incx
                                   , std::complex<double> const * __y
                                   , int                          __incy
                                   , std::complex<double>       * __result
                                   )
{
    return cublasZdotc( __handle, __n, (cuDoubleComplex*)__x, __incx, (cuDoubleComplex*)__y, __incy, (cuDoubleComplex*)__result );
}

// Scale

inline cublasStatus_t __cublas_scal( cublasHandle_t         __handle
                                   , int                    __n
                                   , float                  __alpha
                                   , float                * __data
                                   , int                    __incx
                                   )
{
    return cublasSscal( __handle, __n, &__alpha, __data, __incx );
}

inline cublasStatus_t __cublas_scal( cublasHandle_t         __handle
                                   , int                    __n
                                   , double                 __alpha
                                   , double               * __data
                                   , int                    __incx
                                   )
{
    return cublasDscal( __handle, __n, &__alpha, __data, __incx );
}

inline cublasStatus_t __cublas_scal( cublasHandle_t         __handle
                                   , int                    __n
                                   , std::complex<float>    __alpha
                                   , std::complex<float>  * __data
                                   , int                    __incx
                                   )
{
    return cublasCscal( __handle, __n, (cuComplex*)&__alpha, (cuComplex*)__data, __incx );
}

inline cublasStatus_t __cublas_scal( cublasHandle_t         __handle
                                   , int                    __n
                                   , std::complex<double>   __alpha
                                   , std::complex<double> * __data
                                   , int                    __incx
                                   )
{
    __SCALAR_ALIGNAS(std::complex<double>) __aligned_alpha{ __alpha };

    return cublasZscal( __handle, __n, (cuDoubleComplex*)&__aligned_alpha, (cuDoubleComplex*)__data, __incx );
}

// Lacpy - matrix copy
inline cublasStatus_t __cublas_lacpy( cublasHandle_t         __handle
                                    , char                   __uplo
                                    , int                    __m
                                    , int                    __n
                                    , float const          * __A
                                    , int                    __lda
                                    , float                * __B
                                    , int                    __ldb
                                    )
{
    cublasStatus_t __status = CUBLAS_STATUS_SUCCESS;

    if (__uplo == 'A' || __uplo == 'a') {
        // Copy all elements
        for (int __j = 0; __j < __n; ++__j) {
            __status = cublasScopy(__handle, __m, __A + __j * __lda, 1, __B + __j * __ldb, 1);
            if (__status != CUBLAS_STATUS_SUCCESS) return __status;
        }
    } else if (__uplo == 'U' || __uplo == 'u') {
        // Copy upper triangular part
        for (int __j = 0; __j < __n; ++__j) {
            int __len = (__j + 1 < __m) ? __j + 1 : __m;
            __status = cublasScopy(__handle, __len, __A + __j * __lda, 1, __B + __j * __ldb, 1);
            if (__status != CUBLAS_STATUS_SUCCESS) return __status;
        }
    } else if (__uplo == 'L' || __uplo == 'l') {
        // Copy lower triangular part
        for (int __j = 0; __j < __n; ++__j) {
            int __len = __m - __j;
            if (__len > 0) {
                __status = cublasScopy(__handle, __len, __A + __j * __lda + __j, 1, __B + __j * __ldb + __j, 1);
                if (__status != CUBLAS_STATUS_SUCCESS) return __status;
            }
        }
    }

    return __status;
}

inline cublasStatus_t __cublas_lacpy( cublasHandle_t         __handle
                                    , char                   __uplo
                                    , int                    __m
                                    , int                    __n
                                    , double const         * __A
                                    , int                    __lda
                                    , double               * __B
                                    , int                    __ldb
                                    )
{
    cublasStatus_t __status = CUBLAS_STATUS_SUCCESS;

    if (__uplo == 'A' || __uplo == 'a') {
        // Copy all elements
        for (int __j = 0; __j < __n; ++__j) {
            __status = cublasDcopy(__handle, __m, __A + __j * __lda, 1, __B + __j * __ldb, 1);
            if (__status != CUBLAS_STATUS_SUCCESS) return __status;
        }
    } else if (__uplo == 'U' || __uplo == 'u') {
        // Copy upper triangular part
        for (int __j = 0; __j < __n; ++__j) {
            int __len = (__j + 1 < __m) ? __j + 1 : __m;
            __status = cublasDcopy(__handle, __len, __A + __j * __lda, 1, __B + __j * __ldb, 1);
            if (__status != CUBLAS_STATUS_SUCCESS) return __status;
        }
    } else if (__uplo == 'L' || __uplo == 'l') {
        // Copy lower triangular part
        for (int __j = 0; __j < __n; ++__j) {
            int __len = __m - __j;
            if (__len > 0) {
                __status = cublasDcopy(__handle, __len, __A + __j * __lda + __j, 1, __B + __j * __ldb + __j, 1);
                if (__status != CUBLAS_STATUS_SUCCESS) return __status;
            }
        }
    }

    return __status;
}

inline cublasStatus_t __cublas_lacpy( cublasHandle_t              __handle
                                    , char                        __uplo
                                    , int                         __m
                                    , int                         __n
                                    , std::complex<float> const * __A
                                    , int                         __lda
                                    , std::complex<float>       * __B
                                    , int                         __ldb
                                    )
{
    cublasStatus_t __status = CUBLAS_STATUS_SUCCESS;

    if (__uplo == 'A' || __uplo == 'a') {
        // Copy all elements
        for (int __j = 0; __j < __n; ++__j) {
            __status = cublasCcopy(__handle, __m, (cuComplex const *)(__A + __j * __lda), 1, (cuComplex *)(__B + __j * __ldb), 1);
            if (__status != CUBLAS_STATUS_SUCCESS) return __status;
        }
    } else if (__uplo == 'U' || __uplo == 'u') {
        // Copy upper triangular part
        for (int __j = 0; __j < __n; ++__j) {
            int __len = (__j + 1 < __m) ? __j + 1 : __m;
            __status = cublasCcopy(__handle, __len, (cuComplex const *)(__A + __j * __lda), 1, (cuComplex *)(__B + __j * __ldb), 1);
            if (__status != CUBLAS_STATUS_SUCCESS) return __status;
        }
    } else if (__uplo == 'L' || __uplo == 'l') {
        // Copy lower triangular part
        for (int __j = 0; __j < __n; ++__j) {
            int __len = __m - __j;
            if (__len > 0) {
                __status = cublasCcopy(__handle, __len, (cuComplex const *)(__A + __j * __lda + __j), 1, (cuComplex *)(__B + __j * __ldb + __j), 1);
                if (__status != CUBLAS_STATUS_SUCCESS) return __status;
            }
        }
    }

    return __status;
}

inline cublasStatus_t __cublas_lacpy( cublasHandle_t               __handle
                                    , char                         __uplo
                                    , int                          __m
                                    , int                          __n
                                    , std::complex<double> const * __A
                                    , int                          __lda
                                    , std::complex<double>       * __B
                                    , int                          __ldb
                                    )
{
    cublasStatus_t __status = CUBLAS_STATUS_SUCCESS;

    if (__uplo == 'A' || __uplo == 'a') {
        // Copy all elements
        for (int __j = 0; __j < __n; ++__j) {
            __status = cublasZcopy(__handle, __m, (cuDoubleComplex const *)(__A + __j * __lda), 1, (cuDoubleComplex *)(__B + __j * __ldb), 1);
            if (__status != CUBLAS_STATUS_SUCCESS) return __status;
        }
    } else if (__uplo == 'U' || __uplo == 'u') {
        // Copy upper triangular part
        for (int __j = 0; __j < __n; ++__j) {
            int __len = (__j + 1 < __m) ? __j + 1 : __m;
            __status = cublasZcopy(__handle, __len, (cuDoubleComplex const *)(__A + __j * __lda), 1, (cuDoubleComplex *)(__B + __j * __ldb), 1);
            if (__status != CUBLAS_STATUS_SUCCESS) return __status;
        }
    } else if (__uplo == 'L' || __uplo == 'l') {
        // Copy lower triangular part
        for (int __j = 0; __j < __n; ++__j) {
            int __len = __m - __j;
            if (__len > 0) {
                __status = cublasZcopy(__handle, __len, (cuDoubleComplex const *)(__A + __j * __lda + __j), 1, (cuDoubleComplex *)(__B + __j * __ldb + __j), 1);
                if (__status != CUBLAS_STATUS_SUCCESS) return __status;
            }
        }
    }

    return __status;
}

// Calls cublas<T>scal to perform element-wise conjugate

//TODO remove these when all functions are updated to do out-of-place conj
template< typename _T >
cublasStatus_t __cublas_conj( cublasHandle_t            __handle
                            , int                       __n
                            , _T                      * __data
                            )
{
    return CUBLAS_STATUS_SUCCESS;
}

inline cublasStatus_t __cublas_conj( cublasHandle_t            __handle
                                   , int                       __n
                                   , std::complex<float >    * __data
                                   )
{
    float const __alpha = -1.0f;
    cublasStatus_t __stat = cublasSscal( __handle, __n, &__alpha, reinterpret_cast<float *>(__data) + 1, 2 );
    return __stat;
}

inline cublasStatus_t __cublas_conj( cublasHandle_t            __handle
                                   , int                       __n
                                   , std::complex<double>    * __data
                                   )
{
    double const __alpha = -1.0;
    cublasStatus_t __stat = cublasDscal( __handle, __n, &__alpha, reinterpret_cast<double *>(__data) + 1, 2 );
    return __stat;
}

template< typename _T >
cublasStatus_t __cublas_conj( cublasHandle_t            __handle
                            , int                       __n
                            , _T                const * __in
                            , _T                      * __out
                            )
{
    cublasStatus_t __stat = __cublas_copy( __handle, __n, __in, __out );
    return __stat;
}

inline cublasStatus_t __cublas_conj( cublasHandle_t               __handle
                                   , int                          __n
                                   , std::complex<float > const * __in
                                   , std::complex<float >       * __out
                                   )
{
    cublasStatus_t __stat = cublasCcopy( __handle, __n, (cuComplex const *)__in, 1, (cuComplex *)__out, 1 );
    if ( __stat != CUBLAS_STATUS_SUCCESS ) { return __stat; }

    float const __alpha = -1.0f;
    __stat = cublasSscal( __handle, __n, &__alpha, reinterpret_cast<float *>(__out) + 1, 2 );
    return __stat;
}

inline cublasStatus_t __cublas_conj( cublasHandle_t               __handle
                                   , int                          __n
                                   , std::complex<double> const * __in
                                   , std::complex<double>       * __out
                                   )
{
    cublasStatus_t __stat = cublasZcopy( __handle, __n, (cuDoubleComplex const *)__in, 1, (cuDoubleComplex *)__out, 1 );
    if ( __stat != CUBLAS_STATUS_SUCCESS ) { return __stat; }

    double const __alpha = -1.0;
    __stat = cublasDscal( __handle, __n, &__alpha, reinterpret_cast<double *>(__out) + 1, 2 );
    return __stat;
}

// nrm2

inline cublasStatus_t __cublas_nrm2( cublasHandle_t               __handle
                                   , int                          __n
                                   , float                const * __x
                                   , int                          __incx
                                   , float                      * __result
                                   )
{
    return cublasSnrm2( __handle, __n, __x, __incx, __result );
}

inline cublasStatus_t __cublas_nrm2( cublasHandle_t               __handle
                                   , int                          __n
                                   , double               const * __x
                                   , int                          __incx
                                   , double                     * __result
                                   )
{
    return cublasDnrm2( __handle, __n, __x, __incx, __result );
}

inline cublasStatus_t __cublas_nrm2( cublasHandle_t               __handle
                                   , int                          __n
                                   , std::complex<float > const * __x
                                   , int                          __incx
                                   , float                      * __result
                                   )
{
    return cublasScnrm2( __handle, __n, (cuComplex*)__x, __incx, __result );
}

inline cublasStatus_t __cublas_nrm2( cublasHandle_t               __handle
                                   , int                          __n
                                   , std::complex<double> const * __x
                                   , int                          __incx
                                   , double                     * __result
                                   )
{
    return cublasDznrm2( __handle, __n, (cuDoubleComplex*)__x, __incx, __result );
}

// BLAS Level 2

// GEMV

inline cublasStatus_t __cublas_gemv( cublasHandle_t               __handle
                                   , cublasOperation_t            __trans
                                   , int                          __m
                                   , int                          __n
                                   , float                const * __alpha
                                   , float                const * __A
                                   , int                          __lda
                                   , float                const * __x
                                   , float                const * __beta
                                   , float                      * __y
                                   )
{
    return cublasSgemv( __handle, __trans, __m, __n, __alpha, __A, __lda, __x, 1, __beta, __y, 1 );
}

inline cublasStatus_t __cublas_gemv( cublasHandle_t               __handle
                                   , cublasOperation_t            __trans
                                   , int                          __m
                                   , int                          __n
                                   , double               const * __alpha
                                   , double               const * __A
                                   , int                          __lda
                                   , double               const * __x
                                   , double               const * __beta
                                   , double                     * __y
                                   )
{
    return cublasDgemv( __handle, __trans, __m, __n, __alpha, __A, __lda, __x, 1, __beta, __y, 1 );
}

inline cublasStatus_t __cublas_gemv( cublasHandle_t               __handle
                                   , cublasOperation_t            __trans
                                   , int                          __m
                                   , int                          __n
                                   , std::complex<float>  const * __alpha
                                   , std::complex<float>  const * __A
                                   , int                          __lda
                                   , std::complex<float>  const * __x
                                   , std::complex<float>  const * __beta
                                   , std::complex<float>        * __y
                                   )
{
    return cublasCgemv( __handle
                      , __trans
                      , __m
                      , __n
                      , (cuComplex const *)__alpha
                      , (cuComplex const *)__A
                      , __lda
                      , (cuComplex const *)__x
                      , 1
                      , (cuComplex const *)__beta
                      , (cuComplex       *)__y
                      , 1
                      );
}

inline cublasStatus_t __cublas_gemv( cublasHandle_t               __handle
                                   , cublasOperation_t            __trans
                                   , int                          __m
                                   , int                          __n
                                   , std::complex<double> const * __alpha
                                   , std::complex<double> const * __A
                                   , int                          __lda
                                   , std::complex<double> const * __x
                                   , std::complex<double> const * __beta
                                   , std::complex<double>       * __y
                                   )
{
    __SCALAR_ALIGNAS(std::complex<double>) const __aligned_beta{ *__beta };

    return cublasZgemv( __handle
                      , __trans
                      , __m
                      , __n
                      , (cuDoubleComplex const *)__alpha
                      , (cuDoubleComplex const *)__A
                      , __lda
                      , (cuDoubleComplex const *)__x
                      , 1
                      , (cuDoubleComplex const *)&__aligned_beta
                      , (cuDoubleComplex       *)__y
                      , 1
                      );
}

// TRSV

inline cublasStatus_t __cublas_trsv( cublasHandle_t               __handle
                                   , cublasFillMode_t             __uplo
                                   , cublasOperation_t            __trans
                                   , cublasDiagType_t             __diag
                                   , int                          __n
                                   , float                const * __A
                                   , int                          __lda
                                   , float                      * __x
                                   )
{
    return cublasStrsv( __handle, __uplo, __trans, __diag, __n, __A, __lda, __x, 1 );
}

inline cublasStatus_t __cublas_trsv( cublasHandle_t               __handle
                                   , cublasFillMode_t             __uplo
                                   , cublasOperation_t            __trans
                                   , cublasDiagType_t             __diag
                                   , int                          __n
                                   , double               const * __A
                                   , int                          __lda
                                   , double                     * __x
                                   )
{
    return cublasDtrsv( __handle, __uplo, __trans, __diag, __n, __A, __lda, __x, 1 );
}

inline cublasStatus_t __cublas_trsv( cublasHandle_t               __handle
                                   , cublasFillMode_t             __uplo
                                   , cublasOperation_t            __trans
                                   , cublasDiagType_t             __diag
                                   , int                          __n
                                   , std::complex<float > const * __A
                                   , int                          __lda
                                   , std::complex<float >       * __x
                                   )
{
    return cublasCtrsv( __handle, __uplo, __trans, __diag, __n, (cuComplex*)__A, __lda, (cuComplex*)__x, 1 );
}

inline cublasStatus_t __cublas_trsv( cublasHandle_t               __handle
                                   , cublasFillMode_t             __uplo
                                   , cublasOperation_t            __trans
                                   , cublasDiagType_t             __diag
                                   , int                          __n
                                   , std::complex<double> const * __A
                                   , int                          __lda
                                   , std::complex<double>       * __x
                                   )
{
    return cublasZtrsv( __handle, __uplo, __trans, __diag, __n, (cuDoubleComplex*)__A, __lda, (cuDoubleComplex*)__x, 1 );
}

// BLAS Level 3

// GEMM

template< class _Scalar, class _A_t, class _B_t, class _C_t >
cublasStatus_t __cublas_gemm( cublasHandle_t               __handle
                            , cublasOperation_t            __trans_A
                            , cublasOperation_t            __trans_B
                            , int                          __m
                            , int                          __n
                            , int                          __k
                            , _Scalar              const * __alpha
                            , _A_t                 const * __A
                            , int                          __lda
                            , _B_t                 const * __B
                            , int                          __ldb
                            , _Scalar              const * __beta
                            , _C_t                       * __C
                            , int                          __ldc
                            )
{
    return cublasGemmEx( __handle
                       , __trans_A
                       , __trans_B
                       , __m
                       , __n
                       , __k
                       , (void const *)__alpha
                       , (void const *)__A
                       , __cuda_data_type< _A_t >
                       , __lda
                       , (void const *)__B
                       , __cuda_data_type< _B_t >
                       , __ldb
                       , (void const *)__beta
                       , (void       *)__C
                       , __cuda_data_type< _C_t >
                       , __ldc
                       , __cublas_compute_type< _C_t >
                       , CUBLAS_GEMM_DEFAULT
                       );
}

inline cublasStatus_t __cublas_gemm( cublasHandle_t               __handle
                                   , cublasOperation_t            __trans_A
                                   , cublasOperation_t            __trans_B
                                   , int                          __m
                                   , int                          __n
                                   , int                          __k
                                   , float                const * __alpha
                                   , float                const * __A
                                   , int                          __lda
                                   , float                const * __B
                                   , int                          __ldb
                                   , float                const * __beta
                                   , float                      * __C
                                   , int                          __ldc
                                   )
{
    return cublasSgemm( __handle, __trans_A, __trans_B, __m, __n, __k, __alpha, __A, __lda, __B, __ldb, __beta, __C, __ldc );
}

inline cublasStatus_t __cublas_gemm( cublasHandle_t               __handle
                                   , cublasOperation_t            __trans_A
                                   , cublasOperation_t            __trans_B
                                   , int                          __m
                                   , int                          __n
                                   , int                          __k
                                   , double               const * __alpha
                                   , double               const * __A
                                   , int                          __lda
                                   , double               const * __B
                                   , int                          __ldb
                                   , double               const * __beta
                                   , double                     * __C
                                   , int                          __ldc
                                   )
{
    return cublasDgemm( __handle, __trans_A, __trans_B, __m, __n, __k, __alpha, __A, __lda, __B, __ldb, __beta, __C, __ldc );
}

inline cublasStatus_t __cublas_gemm( cublasHandle_t               __handle
                                   , cublasOperation_t            __trans_A
                                   , cublasOperation_t            __trans_B
                                   , int                          __m
                                   , int                          __n
                                   , int                          __k
                                   , std::complex<float>  const * __alpha
                                   , std::complex<float>  const * __A
                                   , int                          __lda
                                   , std::complex<float>  const * __B
                                   , int                          __ldb
                                   , std::complex<float>  const * __beta
                                   , std::complex<float>        * __C
                                   , int                          __ldc
                                   )
{
    return cublasCgemm( __handle
                      , __trans_A
                      , __trans_B
                      , __m
                      , __n
                      , __k
                      , (cuComplex const *)__alpha
                      , (cuComplex const *)__A
                      , __lda
                      , (cuComplex const *)__B
                      , __ldb
                      , (cuComplex const *)__beta
                      , (cuComplex       *)__C
                      , __ldc
                      );
}

inline cublasStatus_t __cublas_gemm( cublasHandle_t               __handle
                                   , cublasOperation_t            __trans_A
                                   , cublasOperation_t            __trans_B
                                   , int                          __m
                                   , int                          __n
                                   , int                          __k
                                   , std::complex<double> const * __alpha
                                   , std::complex<double> const * __A
                                   , int                          __lda
                                   , std::complex<double> const * __B
                                   , int                          __ldb
                                   , std::complex<double> const * __beta
                                   , std::complex<double>       * __C
                                   , int                          __ldc
                                   )
{
    __SCALAR_ALIGNAS(std::complex<double>) const __aligned_alpha{ *__alpha };
    __SCALAR_ALIGNAS(std::complex<double>) const __aligned_beta { *__beta  };

    return cublasZgemm( __handle
                      , __trans_A
                      , __trans_B
                      , __m
                      , __n
                      , __k
                      , (cuDoubleComplex const *)&__aligned_alpha
                      , (cuDoubleComplex const *)__A
                      , __lda
                      , (cuDoubleComplex const *)__B
                      , __ldb
                      , (cuDoubleComplex const *)&__aligned_beta
                      , (cuDoubleComplex       *)__C
                      , __ldc
                      );
}

// SYRK

inline cublasStatus_t __cublas_syrk( cublasHandle_t               __handle
                                   , cublasFillMode_t             __uplo
                                   , cublasOperation_t            __trans
                                   , int                          __n
                                   , int                          __k
                                   , float                const * __alpha
                                   , float                const * __A
                                   , int                          __lda
                                   , float                const * __beta
                                   , float                      * __C
                                   , int                          __ldc
                                   )
{
    return cublasSsyrk( __handle, __uplo, __trans, __n, __k, __alpha, __A, __lda, __beta, __C, __ldc );
}

inline cublasStatus_t __cublas_syrk( cublasHandle_t               __handle
                                   , cublasFillMode_t             __uplo
                                   , cublasOperation_t            __trans
                                   , int                          __n
                                   , int                          __k
                                   , double               const * __alpha
                                   , double               const * __A
                                   , int                          __lda
                                   , double               const * __beta
                                   , double                     * __C
                                   , int                          __ldc
                                   )
{
    return cublasDsyrk( __handle, __uplo, __trans, __n, __k, __alpha, __A, __lda, __beta, __C, __ldc );
}

inline cublasStatus_t __cublas_syrk( cublasHandle_t               __handle
                                   , cublasFillMode_t             __uplo
                                   , cublasOperation_t            __trans
                                   , int                          __n
                                   , int                          __k
                                   , std::complex<float > const * __alpha
                                   , std::complex<float > const * __A
                                   , int                          __lda
                                   , std::complex<float > const * __beta
                                   , std::complex<float >       * __C
                                   , int                          __ldc
                                   )
{
    return cublasCsyrk( __handle
                      , __uplo
                      , __trans
                      , __n
                      , __k
                      , (cuComplex const *)__alpha
                      , (cuComplex const *)__A
                      , __lda
                      , (cuComplex const *)__beta
                      , (cuComplex       *)__C
                      , __ldc
                      );
}

inline cublasStatus_t __cublas_syrk( cublasHandle_t               __handle
                                   , cublasFillMode_t             __uplo
                                   , cublasOperation_t            __trans
                                   , int                          __n
                                   , int                          __k
                                   , std::complex<double> const * __alpha
                                   , std::complex<double> const * __A
                                   , int                          __lda
                                   , std::complex<double> const * __beta
                                   , std::complex<double>       * __C
                                   , int                          __ldc
                                   )
{
    __SCALAR_ALIGNAS(std::complex<double>) const __aligned_alpha{ *__alpha };
    __SCALAR_ALIGNAS(std::complex<double>) const __aligned_beta { *__beta  };

    return cublasZsyrk( __handle
                      , __uplo
                      , __trans
                      , __n
                      , __k
                      , (cuDoubleComplex const *)&__aligned_alpha
                      , (cuDoubleComplex const *)__A
                      , __lda
                      , (cuDoubleComplex const *)&__aligned_beta
                      , (cuDoubleComplex       *)__C
                      , __ldc
                      );
}

// SYR2K

inline cublasStatus_t __cublas_syr2k( cublasHandle_t               __handle
                                    , cublasFillMode_t             __uplo
                                    , cublasOperation_t            __trans
                                    , int                          __n
                                    , int                          __k
                                    , float                const * __alpha
                                    , float                const * __A
                                    , int                          __lda
                                    , float                const * __B
                                    , int                          __ldb
                                    , float                const * __beta
                                    , float                      * __C
                                    , int                          __ldc
                                    )
{
    return cublasSsyr2k( __handle, __uplo, __trans, __n, __k, __alpha, __A, __lda, __B, __ldb, __beta, __C, __ldc );
}

inline cublasStatus_t __cublas_syr2k( cublasHandle_t               __handle
                                    , cublasFillMode_t             __uplo
                                    , cublasOperation_t            __trans
                                    , int                          __n
                                    , int                          __k
                                    , double               const * __alpha
                                    , double               const * __A
                                    , int                          __lda
                                    , double               const * __B
                                    , int                          __ldb
                                    , double               const * __beta
                                    , double                     * __C
                                    , int                          __ldc
                                    )
{
    return cublasDsyr2k( __handle, __uplo, __trans, __n, __k, __alpha, __A, __lda, __B, __ldb, __beta, __C, __ldc );
}

inline cublasStatus_t __cublas_syr2k( cublasHandle_t               __handle
                                    , cublasFillMode_t             __uplo
                                    , cublasOperation_t            __trans
                                    , int                          __n
                                    , int                          __k
                                    , std::complex<float > const * __alpha
                                    , std::complex<float > const * __A
                                    , int                          __lda
                                    , std::complex<float > const * __B
                                    , int                          __ldb
                                    , std::complex<float > const * __beta
                                    , std::complex<float >       * __C
                                    , int                          __ldc
                                    )
{
    return cublasCsyr2k( __handle
                       , __uplo
                       , __trans
                       , __n
                       , __k
                       , (cuComplex const *)__alpha
                       , (cuComplex const *)__A
                       , __lda
                       , (cuComplex const *)__B
                       , __ldb
                       , (cuComplex const *)__beta
                       , (cuComplex       *)__C
                       , __ldc
                       );
}

inline cublasStatus_t __cublas_syr2k( cublasHandle_t               __handle
                                    , cublasFillMode_t             __uplo
                                    , cublasOperation_t            __trans
                                    , int                          __n
                                    , int                          __k
                                    , std::complex<double> const * __alpha
                                    , std::complex<double> const * __A
                                    , int                          __lda
                                    , std::complex<double> const * __B
                                    , int                          __ldb
                                    , std::complex<double> const * __beta
                                    , std::complex<double>       * __C
                                    , int                          __ldc
                                    )
{
    __SCALAR_ALIGNAS(std::complex<double>) const __aligned_alpha{ *__alpha };
    __SCALAR_ALIGNAS(std::complex<double>) const __aligned_beta { *__beta  };

    return cublasZsyr2k( __handle
                       , __uplo
                       , __trans
                       , __n
                       , __k
                       , (cuDoubleComplex const *)&__aligned_alpha
                       , (cuDoubleComplex const *)__A
                       , __lda
                       , (cuDoubleComplex const *)__B
                       , __ldb
                       , (cuDoubleComplex const *)&__aligned_beta
                       , (cuDoubleComplex       *)__C
                       , __ldc
                       );
}

// TRSM

inline cublasStatus_t __cublas_trsm( cublasHandle_t               __handle
                                   , cublasSideMode_t             __side
                                   , cublasFillMode_t             __uplo
                                   , cublasOperation_t            __trans
                                   , cublasDiagType_t             __diag
                                   , int                          __m
                                   , int                          __n
                                   , float                const * __alpha
                                   , float                const * __A
                                   , int                          __lda
                                   , float                      * __B
                                   , int                          __ldb
                                   )
{
    return cublasStrsm( __handle, __side, __uplo, __trans, __diag, __m, __n, __alpha, __A, __lda, __B, __ldb );
}

inline cublasStatus_t __cublas_trsm( cublasHandle_t               __handle
                                   , cublasSideMode_t             __side
                                   , cublasFillMode_t             __uplo
                                   , cublasOperation_t            __trans
                                   , cublasDiagType_t             __diag
                                   , int                          __m
                                   , int                          __n
                                   , double               const * __alpha
                                   , double               const * __A
                                   , int                          __lda
                                   , double                     * __B
                                   , int                          __ldb
                                   )
{
    return cublasDtrsm( __handle, __side, __uplo, __trans, __diag, __m, __n, __alpha, __A, __lda, __B, __ldb );
}

inline cublasStatus_t __cublas_trsm( cublasHandle_t               __handle
                                   , cublasSideMode_t             __side
                                   , cublasFillMode_t             __uplo
                                   , cublasOperation_t            __trans
                                   , cublasDiagType_t             __diag
                                   , int                          __m
                                   , int                          __n
                                   , std::complex<float > const * __alpha
                                   , std::complex<float > const * __A
                                   , int                          __lda
                                   , std::complex<float >       * __B
                                   , int                          __ldb
                                   )
{
    return cublasCtrsm( __handle
                      , __side
                      , __uplo
                      , __trans
                      , __diag
                      , __m
                      , __n
                      , (cuComplex const *)__alpha
                      , (cuComplex const *)__A
                      , __lda
                      , (cuComplex       *)__B
                      , __ldb
                      );
}

inline cublasStatus_t __cublas_trsm( cublasHandle_t               __handle
                                   , cublasSideMode_t             __side
                                   , cublasFillMode_t             __uplo
                                   , cublasOperation_t            __trans
                                   , cublasDiagType_t             __diag
                                   , int                          __m
                                   , int                          __n
                                   , std::complex<double> const * __alpha
                                   , std::complex<double> const * __A
                                   , int                          __lda
                                   , std::complex<double>       * __B
                                   , int                          __ldb
                                   )
{
    __SCALAR_ALIGNAS(std::complex<double>) const __aligned_alpha{ *__alpha };

    return cublasZtrsm( __handle
                      , __side
                      , __uplo
                      , __trans
                      , __diag
                      , __m
                      , __n
                      , (cuDoubleComplex const *)&__aligned_alpha
                      , (cuDoubleComplex const *)__A
                      , __lda
                      , (cuDoubleComplex       *)__B
                      , __ldb
                      );
}

// BLAS-like Extensions

// GEAM

inline cublasStatus_t __cublas_geam( cublasHandle_t               __handle
                                   , cublasOperation_t            __trans_A
                                   , cublasOperation_t            __trans_B
                                   , int                          __m
                                   , int                          __n
                                   , float                const * __alpha
                                   , float                const * __A
                                   , int                          __lda
                                   , float                const * __beta
                                   , float                const * __B
                                   , int                          __ldb
                                   , float                      * __C
                                   , int                          __ldc
                                   )
{
    return cublasSgeam( __handle, __trans_A, __trans_B, __m, __n, __alpha, __A, __lda, __beta, __B, __ldb, __C, __ldc );
}

inline cublasStatus_t __cublas_geam( cublasHandle_t               __handle
                                   , cublasOperation_t            __trans_A
                                   , cublasOperation_t            __trans_B
                                   , int                          __m
                                   , int                          __n
                                   , double               const * __alpha
                                   , double               const * __A
                                   , int                          __lda
                                   , double               const * __beta
                                   , double               const * __B
                                   , int                          __ldb
                                   , double                     * __C
                                   , int                          __ldc
                                   )
{
    return cublasDgeam( __handle, __trans_A, __trans_B, __m, __n, __alpha, __A, __lda, __beta, __B, __ldb, __C, __ldc );
}

inline cublasStatus_t __cublas_geam( cublasHandle_t               __handle
                                   , cublasOperation_t            __trans_A
                                   , cublasOperation_t            __trans_B
                                   , int                          __m
                                   , int                          __n
                                   , std::complex<float > const * __alpha
                                   , std::complex<float > const * __A
                                   , int                          __lda
                                   , std::complex<float > const * __beta
                                   , std::complex<float > const * __B
                                   , int                          __ldb
                                   , std::complex<float >       * __C
                                   , int                          __ldc
                                   )
{
    return cublasCgeam( __handle
                      , __trans_A
                      , __trans_B
                      , __m
                      , __n
                      , (cuComplex const *)__alpha
                      , (cuComplex const *)__A
                      , __lda
                      , (cuComplex const *)__beta
                      , (cuComplex const *)__B
                      , __ldb
                      , (cuComplex       *)__C
                      , __ldc
                      );
}

inline cublasStatus_t __cublas_geam( cublasHandle_t               __handle
                                   , cublasOperation_t            __trans_A
                                   , cublasOperation_t            __trans_B
                                   , int                          __m
                                   , int                          __n
                                   , std::complex<double> const * __alpha
                                   , std::complex<double> const * __A
                                   , int                          __lda
                                   , std::complex<double> const * __beta
                                   , std::complex<double> const * __B
                                   , int                          __ldb
                                   , std::complex<double>       * __C
                                   , int                          __ldc
                                   )
{
    return cublasZgeam( __handle
                      , __trans_A
                      , __trans_B
                      , __m
                      , __n
                      , (cuDoubleComplex const *)__alpha
                      , (cuDoubleComplex const *)__A
                      , __lda
                      , (cuDoubleComplex const *)__beta
                      , (cuDoubleComplex const *)__B
                      , __ldb
                      , (cuDoubleComplex       *)__C
                      , __ldc
                      );
}

cublasOperation_t __op_to_cublas_op( char op );

// Transpose using geam
template< class _Scalar >
cublasStatus_t __cublas_transpose( cublasHandle_t __handle, char __op, int __m, int __n,
                                   _Scalar const * __in, int __ldin, _Scalar * __out, int __ldout )
{
    __SCALAR_ALIGNAS(_Scalar) __alpha = _Scalar{ 1 };
    __SCALAR_ALIGNAS(_Scalar) const __beta = 0;

    // geam output is n x m when transposing m x n
    return __cublas_geam( __handle
                        , __op_to_cublas_op( __op )
                        , CUBLAS_OP_N
                        , __n // rows of output
                        , __m // cols of output
                        , &__alpha
                        , __in
                        , __ldin
                        , &__beta
                        , __out
                        , __ldout
                        , __out
                        , __ldout
                        );
}

} // namespace __cublas_std
