/* ******************************************************************************
 *
 *
 * This program and the accompanying materials are made available under the
 * terms of the Apache License, Version 2.0 which is available at
 * https://www.apache.org/licenses/LICENSE-2.0.
 *
 *  See the NOTICE file distributed with this work for additional
 *  information regarding copyright ownership.
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 * License for the specific language governing permissions and limitations
 * under the License.
 *
 * SPDX-License-Identifier: Apache-2.0
 ******************************************************************************/

//
//  @author raver119@gmail.com
//

#ifndef LIBND4J_BLAS_HELPER_H
#define LIBND4J_BLAS_HELPER_H

// work around conflict with OpenBLAS
struct bfloat16;
#define BFLOAT16 BFLOAT16
#if !SD_CUDA
#include <blas/cblas.h>
#endif
#include <helpers/logger.h>
#include <types/float16.h>
#include <mutex>
#include <atomic>

#ifdef _WIN32
#define CUBLASWINAPI __stdcall
#define CUSOLVERAPI __stdcall
#else
#define CUBLASWINAPI
#define CUSOLVERAPI
#endif

namespace sd {
typedef enum {
  CUBLAS_STATUS_SUCCESS = 0,
  CUBLAS_STATUS_NOT_INITIALIZED = 1,
  CUBLAS_STATUS_ALLOC_FAILED = 3,
  CUBLAS_STATUS_INVALID_VALUE = 7,
  CUBLAS_STATUS_ARCH_MISMATCH = 8,
  CUBLAS_STATUS_MAPPING_ERROR = 11,
  CUBLAS_STATUS_EXECUTION_FAILED = 13,
  CUBLAS_STATUS_INTERNAL_ERROR = 14,
  CUBLAS_STATUS_NOT_SUPPORTED = 15,
  CUBLAS_STATUS_LICENSE_ERROR = 16
} cublasStatus_t;

typedef enum { CUBLAS_OP_N = 0, CUBLAS_OP_T = 1, CUBLAS_OP_C = 2 } cublasOperation_t;

struct cublasContext;
typedef struct cublasContext *cublasHandle_t;

typedef enum {
  CUDA_R_16F = 2,  /* real as a half */
  CUDA_C_16F = 6,  /* complex as a pair of half numbers */
  CUDA_R_32F = 0,  /* real as a float */
  CUDA_C_32F = 4,  /* complex as a pair of float numbers */
  CUDA_R_64F = 1,  /* real as a double */
  CUDA_C_64F = 5,  /* complex as a pair of double numbers */
  CUDA_R_8I = 3,   /* real as a signed char */
  CUDA_C_8I = 7,   /* complex as a pair of signed char numbers */
  CUDA_R_8U = 8,   /* real as a unsigned char */
  CUDA_C_8U = 9,   /* complex as a pair of unsigned char numbers */
  CUDA_R_32I = 10, /* real as a signed int */
  CUDA_C_32I = 11, /* complex as a pair of signed int numbers */
  CUDA_R_32U = 12, /* real as a unsigned int */
  CUDA_C_32U = 13  /* complex as a pair of unsigned int numbers */
} cublasDataType_t;

#if !defined(SD_CUDA)
typedef void (*CblasSgemv)(CBLAS_ORDER Layout, CBLAS_TRANSPOSE TransA, int M, int N, float alpha, float *A, int lda,
                           float *X, int incX, float beta, float *Y, int incY);

typedef void (*CblasDgemv)(CBLAS_ORDER Layout, CBLAS_TRANSPOSE TransA, int M, int N, double alpha, double *A, int lda,
                           double *X, int incX, double beta, double *Y, int incY);

typedef void (*CblasSgemm)(CBLAS_ORDER Layout, CBLAS_TRANSPOSE TransA, CBLAS_TRANSPOSE TransB, int M, int N, int K,
                           float alpha, float *A, int lda, float *B, int ldb, float beta, float *C, int ldc);

typedef void (*CblasDgemm)(CBLAS_ORDER Layout, CBLAS_TRANSPOSE TransA, CBLAS_TRANSPOSE TransB, int M, int N, int K,
                           double alpha, double *A, int lda, double *B, int ldb, double beta, double *C, int ldc);

typedef void (*CblasSgemmBatch)(CBLAS_ORDER Layout, CBLAS_TRANSPOSE *TransA_Array, CBLAS_TRANSPOSE *TransB_Array,
                                int *M_Array, int *N_Array, int *K_Array, float *alpha_Array, float **A_Array,
                                int *lda_Array, float **B_Array, int *ldb_Array, float *beta_Array, float **C_Array,
                                int *ldc_Array, int group_count, int *group_size);

typedef void (*CblasDgemmBatch)(CBLAS_ORDER Layout, CBLAS_TRANSPOSE *TransA_Array, CBLAS_TRANSPOSE *TransB_Array,
                                int *M_Array, int *N_Array, int *K_Array, double *alpha_Array, double **A_Array,
                                int *lda_Array, double **B_Array, int *ldb_Array, double *beta_Array, double **C_Array,
                                int *ldc_Array, int group_count, int *group_size);
#endif
#ifdef LAPACK_ROW_MAJOR
#undef LAPACK_ROW_MAJOR
#endif

#ifdef LAPACK_COL_MAJOR
#undef LAPACK_COL_MAJOR
#endif
enum LAPACK_LAYOUT { LAPACK_ROW_MAJOR = 101, LAPACK_COL_MAJOR = 102 };

typedef int (*LapackeSgesvd)(LAPACK_LAYOUT matrix_layout, char jobu, char jobvt, int m, int n, float *a, int lda,
                             float *s, float *u, int ldu, float *vt, int ldvt, float *superb);

typedef int (*LapackeDgesvd)(LAPACK_LAYOUT matrix_layout, char jobu, char jobvt, int m, int n, double *a, int lda,
                             double *s, double *u, int ldu, double *vt, int ldvt, double *superb);

typedef int (*LapackeSgesdd)(LAPACK_LAYOUT matrix_layout, char jobz, int m, int n, float *a, int lda, float *s,
                             float *u, int ldu, float *vt, int ldvt);
typedef int (*LapackeDgesdd)(LAPACK_LAYOUT matrix_layout, char jobz, int m, int n, double *a, int lda, double *s,
                             double *u, int ldu, double *vt, int ldvt);

#if defined(SD_CUDA)
typedef cublasStatus_t(CUBLASWINAPI *CublasSgemv)(cublasHandle_t handle, cublasOperation_t trans, int m, int n,
                                                  float *alpha, /* host or device pointer */
                                                  float *A, int lda, float *x, int incx,
                                                  float *beta, /* host or device pointer */
                                                  float *y, int incy);

typedef cublasStatus_t(CUBLASWINAPI *CublasDgemv)(cublasHandle_t handle, cublasOperation_t trans, int m, int n,
                                                  double *alpha, /* host or device pointer */
                                                  double *A, int lda, double *x, int incx,
                                                  double *beta, /* host or device pointer */
                                                  double *y, int incy);

typedef cublasStatus_t(CUBLASWINAPI *CublasHgemm)(cublasHandle_t handle, cublasOperation_t transa,
                                                  cublasOperation_t transb, int m, int n, int k,
                                                  __half *alpha, /* host or device pointer */
                                                  __half *A, int lda, __half *B, int ldb,
                                                  __half *beta, /* host or device pointer */
                                                  __half *C, int ldc);

typedef cublasStatus_t(CUBLASWINAPI *CublasSgemm)(cublasHandle_t handle, cublasOperation_t transa,
                                                  cublasOperation_t transb, int m, int n, int k,
                                                  float *alpha, /* host or device pointer */
                                                  float *A, int lda, float *B, int ldb,
                                                  float *beta, /* host or device pointer */
                                                  float *C, int ldc);

typedef cublasStatus_t(CUBLASWINAPI *CublasDgemm)(cublasHandle_t handle, cublasOperation_t transa,
                                                  cublasOperation_t transb, int m, int n, int k,
                                                  double *alpha, /* host or device pointer */
                                                  double *A, int lda, double *B, int ldb,
                                                  double *beta, /* host or device pointer */
                                                  double *C, int ldc);

typedef cublasStatus_t(CUBLASWINAPI *CublasSgemmEx)(cublasHandle_t handle, cublasOperation_t transa,
                                                    cublasOperation_t transb, int m, int n, int k,
                                                    float *alpha, /* host or device pointer */
                                                    void *A, cublasDataType_t Atype, int lda, void *B,
                                                    cublasDataType_t Btype, int ldb,
                                                    float *beta, /* host or device pointer */
                                                    void *C, cublasDataType_t Ctype, int ldc);

typedef cublasStatus_t(CUBLASWINAPI *CublasHgemmBatched)(cublasHandle_t handle, cublasOperation_t transa,
                                                         cublasOperation_t transb, int m, int n, int k,
                                                         __half *alpha, /* host or device pointer */
                                                         __half *Aarray[], int lda, __half *Barray[], int ldb,
                                                         __half *beta, /* host or device pointer */
                                                         __half *Carray[], int ldc, int batchCount);

typedef cublasStatus_t(CUBLASWINAPI *CublasSgemmBatched)(cublasHandle_t handle, cublasOperation_t transa,
                                                         cublasOperation_t transb, int m, int n, int k,
                                                         float *alpha, /* host or device pointer */
                                                         float *Aarray[], int lda, float *Barray[], int ldb,
                                                         float *beta, /* host or device pointer */
                                                         float *Carray[], int ldc, int batchCount);

typedef cublasStatus_t(CUBLASWINAPI *CublasDgemmBatched)(cublasHandle_t handle, cublasOperation_t transa,
                                                         cublasOperation_t transb, int m, int n, int k,
                                                         double *alpha, /* host or device pointer */
                                                         double *Aarray[], int lda, double *Barray[], int ldb,
                                                         double *beta, /* host or device pointer */
                                                         double *Carray[], int ldc, int batchCount);

typedef enum {
  CUSOLVER_STATUS_SUCCESS = 0,
  CUSOLVER_STATUS_NOT_INITIALIZED = 1,
  CUSOLVER_STATUS_ALLOC_FAILED = 2,
  CUSOLVER_STATUS_INVALID_VALUE = 3,
  CUSOLVER_STATUS_ARCH_MISMATCH = 4,
  CUSOLVER_STATUS_MAPPING_ERROR = 5,
  CUSOLVER_STATUS_EXECUTION_FAILED = 6,
  CUSOLVER_STATUS_INTERNAL_ERROR = 7,
  CUSOLVER_STATUS_MATRIX_TYPE_NOT_SUPPORTED = 8,
  CUSOLVER_STATUS_NOT_SUPPORTED = 9,
  CUSOLVER_STATUS_ZERO_PIVOT = 10,
  CUSOLVER_STATUS_INVALID_LICENSE = 11
} cusolverStatus_t;

typedef enum { CUSOLVER_EIG_TYPE_1 = 1, CUSOLVER_EIG_TYPE_2 = 2, CUSOLVER_EIG_TYPE_3 = 3 } cusolverEigType_t;

typedef enum { CUSOLVER_EIG_MODE_NOVECTOR = 0, CUSOLVER_EIG_MODE_VECTOR = 1 } cusolverEigMode_t;

struct cusolverDnContext;
typedef struct cusolverDnContext *cusolverDnHandle_t;

typedef cusolverStatus_t(CUSOLVERAPI *CusolverDnSgesvdBufferSize)(cusolverDnHandle_t handle, int m, int n, int *lwork);

typedef cusolverStatus_t(CUSOLVERAPI *CusolverDnDgesvdBufferSize)(cusolverDnHandle_t handle, int m, int n, int *lwork);

typedef cusolverStatus_t(CUSOLVERAPI *CusolverDnSgesvd)(cusolverDnHandle_t handle, signed char jobu, signed char jobvt,
                                                        int m, int n, float *A, int lda, float *S, float *U, int ldu,
                                                        float *VT, int ldvt, float *work, int lwork, float *rwork,
                                                        int *info);

typedef cusolverStatus_t(CUSOLVERAPI *CusolverDnDgesvd)(cusolverDnHandle_t handle, signed char jobu, signed char jobvt,
                                                        int m, int n, double *A, int lda, double *S, double *U, int ldu,
                                                        double *VT, int ldvt, double *work, int lwork, double *rwork,
                                                        int *info);
#endif
enum BlasFunctions {
  GEMV = 0,
  GEMM = 1,
};

class BlasHelper {
 private:
  // Private constructor for singleton pattern
  BlasHelper();

  bool _hasHgemv = false;
  bool _hasHgemm = false;
  bool _hasHgemmBatch = false;

  bool _hasSgemv = false;
  bool _hasSgemm = false;
  bool _hasSgemmBatch = false;

  bool _hasDgemv = false;
  bool _hasDgemm = false;
  bool _hasDgemmBatch = false;

  // Mutex for serializing BLAS calls to prevent OpenBLAS TLS corruption
  // and race conditions when called from multiple threads concurrently
  mutable std::mutex _blasMutex;

  // Controls whether BLAS calls are serialized (default: true for safety)
  // Can be disabled via SD_BLAS_SERIALIZE=0 for testing or if using a thread-safe BLAS
  std::atomic<bool> _serializeBlasCalls{true};

  // Number of threads OpenBLAS should use (0 = use OpenBLAS default)
  std::atomic<int> _openblasThreads{0};
#if !defined(SD_CUDA)
  CblasSgemv cblasSgemv;
  CblasDgemv cblasDgemv;
  CblasSgemm cblasSgemm;
  CblasDgemm cblasDgemm;
  CblasSgemmBatch cblasSgemmBatch;
  CblasDgemmBatch cblasDgemmBatch;
  LapackeSgesvd lapackeSgesvd;
  LapackeDgesvd lapackeDgesvd;
  LapackeSgesdd lapackeSgesdd;
  LapackeDgesdd lapackeDgesdd;
#endif

#if defined(SD_CUDA)
  CublasSgemv cublasSgemv;
  CublasDgemv cublasDgemv;
  CublasHgemm cublasHgemm;
  CublasSgemm cublasSgemm;
  CublasDgemm cublasDgemm;
  CublasSgemmEx cublasSgemmEx;
  CublasHgemmBatched cublasHgemmBatched;
  CublasSgemmBatched cublasSgemmBatched;
  CublasDgemmBatched cublasDgemmBatched;
  CusolverDnSgesvdBufferSize cusolverDnSgesvdBufferSize;
  CusolverDnDgesvdBufferSize cusolverDnDgesvdBufferSize;
  CusolverDnSgesvd cusolverDnSgesvd;
  CusolverDnDgesvd cusolverDnDgesvd;
#endif

 public:
  static BlasHelper &getInstance();

  void initializeFunctions(Pointer *functions);
  void initializeDeviceFunctions(Pointer *functions);

  template <typename T>
  bool hasGEMV();

  template <typename T>
  bool hasGEMM();

  bool hasGEMM(const DataType dtype);
  bool hasGEMV(const DataType dtype);

  template <typename T>
  bool hasBatchedGEMM();
#if !defined(SD_CUDA)
  CblasSgemv sgemv();
  CblasDgemv dgemv();

  CblasSgemm sgemm();
  CblasDgemm dgemm();

  CblasSgemmBatch sgemmBatched();
  CblasDgemmBatch dgemmBatched();

  LapackeSgesvd sgesvd();
  LapackeDgesvd dgesvd();

  LapackeSgesdd sgesdd();
  LapackeDgesdd dgesdd();
#endif

  // BLAS call serialization methods to prevent OpenBLAS TLS corruption
  // and race conditions in multi-threaded environments

  /**
   * Returns a lock guard that serializes BLAS calls.
   * Use this before making any BLAS call to prevent concurrent access.
   * Example:
   *   auto lock = BlasHelper::getInstance().lockBlas();
   *   sgemm()(...);
   */
  std::unique_lock<std::mutex> lockBlas() const;

  /**
   * Check if BLAS call serialization is enabled.
   * Default is true for safety with OpenBLAS.
   */
  bool isSerializeBlasCalls() const;

  /**
   * Enable or disable BLAS call serialization.
   * Disable only if using a thread-safe BLAS implementation (e.g., MKL).
   */
  void setSerializeBlasCalls(bool serialize);

  /**
   * Get the configured number of OpenBLAS threads.
   * Returns 0 if using OpenBLAS default.
   */
  int getOpenblasThreads() const;

  /**
   * Set the number of threads OpenBLAS should use.
   * Set to 0 to use OpenBLAS default, or a specific number for control.
   * This is applied when BLAS is initialized.
   */
  void setOpenblasThreads(int threads);

  /**
   * Initialize OpenBLAS thread settings from environment variables.
   * Checks SD_OPENBLAS_THREADS and SD_BLAS_SERIALIZE.
   * Called automatically during getInstance() initialization.
   */
  void initializeBlasThreading();

  // destructor
  ~BlasHelper() noexcept;
};
}  // namespace sd

#endif