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

// A simple CG solver using cuBLAS

#include <vector>
#include <cmath>
#include <algorithm>
#include <iostream>
#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 = double;
using vector_type = std::mdspan<data_type, std::dims<1, size_type>>;
using matrix_type = std::mdspan<data_type, std::dims<2, size_type>>;

void cublas_cg_solver(cublasHandle_t handle, int n,
                      const data_type *A, const data_type *b, data_type *x,
                      data_type *r, data_type *p, data_type *Ap,
                      bool sync, data_type tol = 1e-10)
{
  cudaStream_t stream;
  cublasGetStream(handle, &stream);

  const data_type one = 1.0;
  const data_type neg_one = -1.0;
  const data_type zero = 0.0;

  // r = b - A*x
  CUBLAS_CHECK(cublasDcopy(handle, n, b, 1, r, 1));
  if (sync) CUDA_CHECK(cudaStreamSynchronize(stream));

  CUBLAS_CHECK(cublasDgemv(handle, CUBLAS_OP_N, n, n, &neg_one, A, n, x, 1, &one, r, 1));
  if (sync) CUDA_CHECK(cudaStreamSynchronize(stream));

  // p = r
  CUBLAS_CHECK(cublasDcopy(handle, n, r, 1, p, 1));
  if (sync) CUDA_CHECK(cudaStreamSynchronize(stream));

  data_type rs_old, rs_new;
  CUBLAS_CHECK(cublasDdot(handle, n, r, 1, r, 1, &rs_old));

  for (int i = 0; i < n; ++i) {
    // Ap = A * p
    CUBLAS_CHECK(cublasDgemv(handle, CUBLAS_OP_N, n, n, &one, A, n, p, 1, &zero, Ap, 1));
    if (sync) CUDA_CHECK(cudaStreamSynchronize(stream));

    // alpha = rs_old / (p^T * Ap)
    data_type pAp;
    CUBLAS_CHECK(cublasDdot(handle, n, p, 1, Ap, 1, &pAp));
    data_type alpha = rs_old / pAp;

    // x = x + alpha * p
    CUBLAS_CHECK(cublasDaxpy(handle, n, &alpha, p, 1, x, 1));
    if (sync) CUDA_CHECK(cudaStreamSynchronize(stream));

    // r = r - alpha * Ap
    data_type neg_alpha = -alpha;
    CUBLAS_CHECK(cublasDaxpy(handle, n, &neg_alpha, Ap, 1, r, 1));
    if (sync) CUDA_CHECK(cudaStreamSynchronize(stream));

    // Check convergence
    CUBLAS_CHECK(cublasDdot(handle, n, r, 1, r, 1, &rs_new));
    if (std::sqrt(rs_new) < tol) {
#ifdef _VERBOSE
      std::cout << "Converged in " << i + 1 << " iterations.\n";
#endif
      break;
    }

    // p = r + (rs_new / rs_old) * p
    data_type beta = rs_new / rs_old;
    CUBLAS_CHECK(cublasDscal(handle, n, &beta, p, 1));
    if (sync) CUDA_CHECK(cudaStreamSynchronize(stream));
    CUBLAS_CHECK(cublasDaxpy(handle, n, &one, r, 1, p, 1));
    if (sync) CUDA_CHECK(cudaStreamSynchronize(stream));

    rs_old = rs_new;
  }
}

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), b_vec(N), x_vec(N), b_check_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<data_type>();
      A_vec[j * N + i] = A_vec[i * N + j];
    }
  }
  for (int i = 0; i < N; ++i) {
    b_vec[i] = random_value<data_type>();
  }

  matrix_type A(A_vec.data(), N, N);
  vector_type b(b_vec.data(), N);
  vector_type x(x_vec.data(), N);
  vector_type b_check(b_check_vec.data(), N);

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

  cudaStream_t   stream  = NULL;
  cublasHandle_t cublasH = NULL;

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

  data_type *d_A = nullptr, *d_b = nullptr, *d_x = nullptr;
  data_type *d_r = nullptr, *d_p = nullptr, *d_Ap = nullptr;

  cudaMalloc(reinterpret_cast<void **>(&d_A),  sizeof(data_type) * N * N);
  cudaMalloc(reinterpret_cast<void **>(&d_b),  sizeof(data_type) * N);
  cudaMalloc(reinterpret_cast<void **>(&d_x),  sizeof(data_type) * N);
  cudaMalloc(reinterpret_cast<void **>(&d_r),  sizeof(data_type) * N);
  cudaMalloc(reinterpret_cast<void **>(&d_p),  sizeof(data_type) * N);
  cudaMalloc(reinterpret_cast<void **>(&d_Ap), sizeof(data_type) * N);

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

  // A and b are not modified by CG, copy to device once
  cudaMemcpyAsync(d_A, A_vec.data(), sizeof(data_type) * N * N, cudaMemcpyHostToDevice, stream);
  cudaMemcpyAsync(d_b, b_vec.data(), sizeof(data_type) * N,     cudaMemcpyHostToDevice, stream);

  Timer timer;

  for (int iRepeat = 0; iRepeat < N_REPEAT; ++iRepeat) {
    cudaMemsetAsync(d_x, 0, sizeof(data_type) * N, stream);

    cudaStreamSynchronize(stream);
    timer.start();

    cublas_cg_solver(cublasH, N, d_A, d_b, d_x, d_r, d_p, d_Ap, syncFlag);

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

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

  // Verify: compute A*x on host and compare with b
  la::matrix_vector_product(std::execution::seq, A, x, b_check);

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

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

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

  cudaFree(d_A);
  cudaFree(d_b);
  cudaFree(d_x);
  cudaFree(d_r);
  cudaFree(d_p);
  cudaFree(d_Ap);

  cudaStreamDestroy(stream);
  cublasDestroy(cublasH);
}
