/*
 * 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 C++26 std::linalg

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

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

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 cg_solver(matrix_type A, vector_type b, vector_type x, data_type tol = 1e-10) {
  const auto n = b.extent(0);

  std::vector<data_type> r_vec(n), p_vec(n), Ap_vec(n);
  vector_type r(r_vec.data(), n);
  vector_type p(p_vec.data(), n);
  vector_type Ap(Ap_vec.data(), n);

  // r = b - A*x
  la::matrix_vector_product(exec, A, x, r);
  la::add( exec, b, la::scaled( -1., r ), r );

  // p = r
  la::copy(exec, r, p);

  data_type rs_old = la::dot(exec, r, r);

  for (size_type i = 0; i < (size_type)n; ++i) {
    la::matrix_vector_product(exec, A, p, Ap);

    data_type alpha = rs_old / la::dot(exec, p, Ap);

    // x = x + alpha * p
    la::add(exec, la::scaled(alpha, p), x, x);

    // r = r - alpha * Ap
    la::add(exec, la::scaled(-alpha, Ap), r, r);

    data_type rs_new = la::dot(exec, r, r);
    if (std::sqrt(rs_new) < tol) {
#ifdef _VERBOSE
      std::cout << "Converged in " << i + 1 << " iterations.\n";
#endif
      return;
    }

    // p = r + (rs_new / rs_old) * p
    la::add(exec, r, la::scaled(rs_new / rs_old, p), p);
    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>();
  }

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

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

  Timer timer;

  for (int iRepeat = 0; iRepeat < N_REPEAT; ++iRepeat) {
    std::fill(x_vec.begin(), x_vec.end(), 0.0);

    timer.start();

    cg_solver(A, b, x);

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

  // Verify: compute A*x and compare with b
  la::matrix_vector_product(std::execution::par, 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;
}
