/*
 * 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.
 */

#include <vector>
#include <iostream>
#include <algorithm> // std::copy_n
#include <cuda_runtime.h>
#include <cublas_v2.h>
#include <linalg>
#include "../utils/stdblas_utils.h"
#include "../utils/cublas_utils.h"

#define N_REPEAT 5

namespace la = std::linalg;

using size_type = int;
using data_type = float;
using vector_type = std::mdspan<data_type, std::dims<1, size_type>>;
using matrix_type = std::mdspan<data_type, std::dims<2, size_type>, std::layout_left>;

data_type random_value() { return static_cast<data_type>( rand() ) / static_cast<data_type>( RAND_MAX ); }

__global__ void sqrt_kernel( float * A )
{
  unsigned int pid = blockIdx.x * blockDim.x + threadIdx.x;

  if ( pid == 0 )
  {
    A[pid] = sqrt( A[pid] );
  }
}

int cholesky_factor( cublasHandle_t handle, cublasFillMode_t uplo, float * A, int n, int lda, bool sync )
{
    cudaStream_t stream;
    cublasGetStream( handle, &stream );

    if ( n == 0 ) { return 0; }
    if ( n == 1 )
    {
        sqrt_kernel<<<1,1,0,stream>>>( A );
        return 0;
    }

    int const n1 = n / 2, n2 = n - n1;

    // Factor A11
    float *A11p = A, *A22p = A + IDX2C(n1,n1,lda);
    const int info1 = cholesky_factor( handle, uplo, A11p, n1, lda, sync );
    if ( info1 != 0 ) {
        return info1;
    }

    float const one = 1.f, minusOne = -1.f;

    if ( uplo == CUBLAS_FILL_MODE_UPPER )
    {
        float *A12p = A + IDX2C(0,n1,lda);

        cublasOperation_t trans = CUBLAS_OP_T;

        // Update and scale A12
        CUBLAS_CHECK( cublasStrsm( handle, CUBLAS_SIDE_LEFT, uplo, trans, CUBLAS_DIAG_NON_UNIT
                                 , n1, n2, &one, A11p, lda, A12p, lda
                                 )
                    );
        if ( sync ) CUDA_CHECK( cudaStreamSynchronize( stream ) );

        CUBLAS_CHECK( cublasSsyrk( handle, uplo, trans, n2, n1, &minusOne, A12p, lda, &one, A22p, lda ) );
        if ( sync ) CUDA_CHECK( cudaStreamSynchronize( stream ) );
    }
    else
    {
        float *A21 = A + IDX2C(n1,0,lda);

        CUBLAS_CHECK( cublasStrsm( handle, CUBLAS_SIDE_RIGHT, uplo, CUBLAS_OP_T, CUBLAS_DIAG_NON_UNIT
                                 , n2, n1, &one, A11p, lda, A21, lda
                                 )
                    );
        if ( sync ) CUDA_CHECK( cudaStreamSynchronize( stream ) );

        CUBLAS_CHECK( cublasSsyrk( handle, uplo, CUBLAS_OP_N, n2, n1, &minusOne, A21, lda, &one, A22p, lda ) );
        if ( sync ) CUDA_CHECK( cudaStreamSynchronize( stream ) );
    }

    const int info2 = cholesky_factor( handle, uplo, A22p, n2, lda, sync );

    if ( info2 != 0 ) {
      return info2 + n1;
    }

    return 0;
}

int cholesky_solve( cublasHandle_t handle, cublasFillMode_t uplo, int n, float * A, float * b, float * x, bool sync )
{
    cudaStream_t stream;
    cublasGetStream( handle, &stream );

    CUBLAS_CHECK( cublasScopy( handle, n, b, 1, x, 1 ) );
    if ( sync ) CUDA_CHECK( cudaStreamSynchronize( stream ) );

    if ( uplo == CUBLAS_FILL_MODE_UPPER )
    {
        CUBLAS_CHECK( cublasStrsv( handle, uplo, CUBLAS_OP_T, CUBLAS_DIAG_NON_UNIT, n, A, n, x, 1 ) );
        if ( sync ) CUDA_CHECK( cudaStreamSynchronize( stream ) );

        CUBLAS_CHECK( cublasStrsv( handle, uplo, CUBLAS_OP_N, CUBLAS_DIAG_NON_UNIT, n, A, n, x, 1 ) );
        if ( sync ) CUDA_CHECK( cudaStreamSynchronize( stream ) );
    }
    else
    {
        CUBLAS_CHECK( cublasStrsv( handle, uplo, CUBLAS_OP_N, CUBLAS_DIAG_NON_UNIT, n, A, n, x, 1 ) );
        if ( sync ) CUDA_CHECK( cudaStreamSynchronize( stream ) );

        CUBLAS_CHECK( cublasStrsv( handle, uplo, CUBLAS_OP_T, CUBLAS_DIAG_NON_UNIT, n, A, n, x, 1 ) );
        if ( sync ) CUDA_CHECK( cudaStreamSynchronize( stream ) );
    }

    return 0;
}

int main( int argc, char *argv[] )
{
  int N = 4;

  if ( argc == 2 )
  {
      N = atoi(argv[1]);
  }

#ifdef _VERBOSE
  std::cout << "Matrix size = " << N << std::endl;
#endif

  std::vector<data_type> A_vec( N * N ), A_save_vec( N * N ), b_vec( N ), b_save_vec( N ), x_vec( N ), b_inverted_vec( N );

  for ( int i = 0; i < N; ++i )
  {
    A_vec[i*N+i] = N;
    for ( int j = i + 1; j < N; ++j )
    {
      A_vec[i*N + j] = random_value();
      A_vec[j*N + i] = A_vec[i*N + j];
    }
  }
  for ( int i = 0; i < N; ++i )
  {
    b_vec[i] = random_value();
  }

  matrix_type A     (      A_vec.data(), N, N );
  matrix_type A_save( A_save_vec.data(), N, N );
  vector_type b         (          b_vec.data(), N );
  vector_type b_save    (     b_save_vec.data(), N );
  vector_type x         (          x_vec.data(), N );
  vector_type b_inverted( b_inverted_vec.data(), N );

  std::copy_n( A_vec.data(), A_vec.size(), A_save_vec.data() );
  std::copy_n( b_vec.data(), b_vec.size(), b_save_vec.data() );

#ifdef _VERBOSE
  print_matrix( N, N, A_vec.data(), N );
#endif

  cudaStream_t   stream  = NULL;
  cublasHandle_t cublasH = NULL;
  cublasFillMode_t uplo  = CUBLAS_FILL_MODE_LOWER;
  int lda = N;

  data_type * d_A = nullptr;
  data_type * d_b = nullptr;
  data_type * d_x = nullptr;

  cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking);
  cublasCreate(&cublasH);
  cublasSetStream(cublasH, stream);

  cudaMalloc( reinterpret_cast<void **>( &d_A ), sizeof(data_type) * A_vec.size() );
  cudaMalloc( reinterpret_cast<void **>( &d_b ), sizeof(data_type) * b_vec.size() );
  cudaMalloc( reinterpret_cast<void **>( &d_x ), sizeof(data_type) * x_vec.size() );

#ifndef _NO_SYNC
  bool syncFlag = true;
#else
  bool syncFlag = false;
#endif

  Timer timer;

  for ( int iRepeat = 0; iRepeat < N_REPEAT; ++iRepeat )
  {
    std::copy_n( A_save_vec.data(), A_vec.size(), A_vec.data() );
    std::copy_n( b_save_vec.data(), b_vec.size(), b_vec.data() );

    cudaMemcpyAsync( d_A, A_vec.data(), sizeof(data_type) * A_vec.size(), cudaMemcpyHostToDevice, stream );
    cudaMemcpyAsync( d_b, b_vec.data(), sizeof(data_type) * b_vec.size(), cudaMemcpyHostToDevice, stream );

    cudaStreamSynchronize( stream );
    timer.start();

    cholesky_factor( cublasH, uplo, d_A, N, lda, syncFlag );

    cholesky_solve ( cublasH, uplo, N, d_A, d_b, d_x, syncFlag );

    cudaStreamSynchronize( stream );
    double t = timer.end();
    std::cout << "Timing = " << t << std::endl;
  }

  cudaMemcpyAsync( x_vec.data(), d_x, sizeof(data_type) * x_vec.size(), cudaMemcpyDeviceToHost, stream );

  cudaStreamDestroy( stream );

  la::matrix_vector_product( std::execution::seq, A_save, x, b_inverted );

#ifdef _VERBOSE
  print_matrix( N, 1, b_vec.data(), N );
  print_matrix( N, 1, b_inverted_vec.data(), N );
#endif

  data_type maxDiff = 0.f;
  for ( int i = 0; i < N; ++i )
  {
    maxDiff = std::max( maxDiff, std::abs( b_vec[i] - b_inverted_vec[i] ) );
  }

  std::cout << "Max Diff = " << maxDiff << std::endl;

  cudaFree( d_A );
  cudaFree( d_b );
  cudaFree( d_x );

  cublasDestroy( cublasH );
}
