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

// Based on examples in the stdBLAS proposal, https://wg21.link/p1673, section 17

#include <vector>
#include <utility>
#include <cmath>
#include <algorithm>
#include <iostream>
#include <linalg>
#include "../utils/stdblas_utils.h"

// For GPU cuBLAS back-ends, use execution policy `no_sync(std::execution::par)` to skip the
// synchronization after each cuBLAS API call to match up with how a cuBLAS implementation would
// have been written and gain additional performance.
#if defined __NVCOMPILER_STDPAR_GPU and defined _NO_SYNC
#define exec no_sync(std::execution::par)
#else
#define exec std::execution::par
#endif

#define N_REPEAT 5

namespace la = std::linalg;

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

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

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 ); }

// Work array for cholesky_factor
std::vector<data_type> A22update_vec;

template<class inout_matrix_t,
         class Triangle>
int cholesky_factor(inout_matrix_t A, Triangle t)
{
  using element_type = typename inout_matrix_t::element_type;
  constexpr element_type ZERO {};
  constexpr element_type ONE (1.0);
  const auto n = A.extent(0);

  if (n == 0) {
    return 0;
  }
  else if (n == 1) {
#ifndef __NVCOMPILER_STDPAR_GPU
    if (A[0,0] <= ZERO || std::isnan(A[0,0])) {
      return 1;
    }
    A[0,0] = sqrt(A[0,0]);
#else
    sqrt_kernel<<<1,1>>>(A.data_handle());
#endif
    return 0;
  }
  else {
    // Partition A into [A11, A12,
    //                   A21, A22],
    // where A21 is the transpose of A12.
    const size_type n1 = n / 2;
    const size_type n2 = n - n1;
    auto A11 = submdspan(A, std::pair{0, n1}, std::pair{0, n1});
    auto A22 = submdspan(A, std::pair{n1, n}, std::pair{n1, n});
    matrix_type A22update(A22update_vec.data(), n2, n2);

    // Factor A11
    const int info1 = cholesky_factor(A11, t);
    if (info1 != 0) {
      return info1;
    }

    if constexpr (std::is_same_v<Triangle, la::upper_triangle_t>) {
      // Update and scale A12
      auto A12 = submdspan(A, std::pair{0, n1}, std::pair{n1, n});
      // Unlike a BLAS/cuBLAS call, here the `triangle` parameter tells the
      // algorithm which part of `transposed(A11)` (rather than A11) to use
      la::triangular_matrix_matrix_left_solve( exec
                                             , la::transposed(A11)
                                             , la::lower_triangle
                                             , la::explicit_diagonal
                                             , A12
                                             , A12
                                             );
      // A22 = A22 - A12^T * A12
      la::symmetric_matrix_rank_k_update(exec, -ONE, la::transposed(A12), A22update, t);
    }
    else {
      //
      // Compute the Cholesky factorization A = L * L^T
      //
      // Update and scale A21
      auto A21 = submdspan(A, std::pair{n1, n}, std::pair{0, n1});
      la::triangular_matrix_matrix_right_solve( exec
                                              , la::transposed(A11)
                                              , la::upper_triangle
                                              , la::explicit_diagonal
                                              , A21
                                              , A21
                                              );
      // A22 = A22 - A21 * A21^T
      la::symmetric_matrix_rank_k_update(exec, -ONE, A21, A22update, t);
    }
    la::add(exec, A22, A22update, A22);

    // Factor A22
    const int info2 = cholesky_factor(A22, t);
    if (info2 != 0) {
      return info2 + n1;
    }
  }

  return 0;
}

template<class in_matrix_t,
         class Triangle,
         class in_vector_t,
         class out_vector_t>
void cholesky_solve(
  in_matrix_t A,
  Triangle t,
  in_vector_t b,
  out_vector_t x)
{
  if constexpr (std::is_same_v<Triangle, la::upper_triangle_t>) {
    // Solve Ax=b where A = U^T U
    //
    // Solve U^T c = b, using x to store c.
    la::triangular_matrix_vector_solve(exec, transposed(A), la::lower_triangle, la::explicit_diagonal, b, x);

    // Solve U x = c, overwriting x with result.
    // Use std::execution::par to do a synchronization after the last call of the algorithm
    la::triangular_matrix_vector_solve(std::execution::par, A, la::upper_triangle, la::explicit_diagonal, x, x);
  }
  else {
    // Solve Ax=b where A = L L^T
    //
    // Solve L c = b, using x to store c.
    la::triangular_matrix_vector_solve(exec, A, t, la::explicit_diagonal, b, x);

    // Solve L^T x = c, overwriting x with result.
    // Use std::execution::par to do a synchronization after the last call of the algorithm
    la::triangular_matrix_vector_solve(std::execution::par, la::transposed(A), la::upper_triangle, la::explicit_diagonal, x, x);
  }
}

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 );

  // Allocate work array up-front
  // For the GPU back-end, the device allocation and a H2D copy happen upon the first touch of the data on the device, and no
  // more data movement happens in the iterations since the work array is only used on the device.
  A22update_vec.resize( N * 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();
  }

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

  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 );

  la::copy( std::execution::par, A, A_save );
  la::copy( std::execution::par, b, b_save );

  auto uplo = la::lower_triangle_t{};

  Timer timer;

  for ( int iRepeat = 0; iRepeat < N_REPEAT; ++iRepeat )
  {
    la::copy( std::execution::par, A_save, A );
    la::copy( std::execution::par, b_save, b );

    timer.start();

    cholesky_factor( A, uplo );

    cholesky_solve ( A, uplo, b, x );

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

  la::matrix_vector_product( std::execution::par, 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;
}
