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

// QR solver using cuBLAS Householder reflections

#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>, std::layout_left>;

void qr_factor(cublasHandle_t handle, int n,
                      data_type *A, int lda, data_type *b,
                      data_type *v, data_type *w,
                      bool sync)
{
  cudaStream_t stream;
  cublasGetStream(handle, &stream);

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

  for (int k = 0; k < n - 1; ++k) {
    int m_sub = n - k;
    data_type *A_kk = A + IDX2C(k, k, lda);
    data_type *b_k  = b + k;

    // 1. x_norm = ||A[k:n, k]||
    data_type x_norm;
    CUBLAS_CHECK(cublasDnrm2(handle, m_sub, A_kk, 1, &x_norm));
    if (x_norm == 0.0) continue;

    // 2. Read x[0] = A[k,k] from device
    data_type x0;
    cudaMemcpyAsync(&x0, A_kk, sizeof(data_type), cudaMemcpyDeviceToHost, stream);
    CUDA_CHECK(cudaStreamSynchronize(stream));

    data_type alpha = (x0 > 0) ? -x_norm : x_norm;
    data_type u1 = x0 - alpha;
    data_type tau = -u1 / alpha;

    // 3. Build Householder vector: v = x / u1, v[0] = 1
    CUBLAS_CHECK(cublasDcopy(handle, m_sub, A_kk, 1, v, 1));
    if (sync) CUDA_CHECK(cudaStreamSynchronize(stream));

    data_type inv_u1 = 1.0 / u1;
    CUBLAS_CHECK(cublasDscal(handle, m_sub, &inv_u1, v, 1));
    if (sync) CUDA_CHECK(cudaStreamSynchronize(stream));

    data_type v0 = 1.0;
    cudaMemcpyAsync(v, &v0, sizeof(data_type), cudaMemcpyHostToDevice, stream);
    if (sync) CUDA_CHECK(cudaStreamSynchronize(stream));

    // 4. w = A[k:n, k:n]^T * v  (gemv with transpose)
    CUBLAS_CHECK(cublasDgemv(handle, CUBLAS_OP_T, m_sub, m_sub,
                             &one, A_kk, lda, v, 1, &zero, w, 1));
    if (sync) CUDA_CHECK(cudaStreamSynchronize(stream));

    // 5. A[k:n, k:n] -= tau * v * w^T  (rank-1 update via dger)
    data_type neg_tau = -tau;
    CUBLAS_CHECK(cublasDger(handle, m_sub, m_sub, &neg_tau, v, 1, w, 1, A_kk, lda));
    if (sync) CUDA_CHECK(cudaStreamSynchronize(stream));

    // 6. b[k:n] -= tau * (v^T * b[k:n]) * v
    data_type v_dot_b;
    CUBLAS_CHECK(cublasDdot(handle, m_sub, v, 1, b_k, 1, &v_dot_b));

    data_type scale_b = -tau * v_dot_b;
    CUBLAS_CHECK(cublasDaxpy(handle, m_sub, &scale_b, v, 1, b_k, 1));
    if (sync) CUDA_CHECK(cudaStreamSynchronize(stream));
  }
}

void qr_solve(cublasHandle_t handle, int n,
                     data_type *A, int lda,
                     data_type *b, data_type *x,
                     bool sync)
{
  cudaStream_t stream;
  cublasGetStream(handle, &stream);

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

  CUBLAS_CHECK(cublasDtrsv(handle, CUBLAS_FILL_MODE_UPPER, CUBLAS_OP_N,
                           CUBLAS_DIAG_NON_UNIT, n, A, lda, x, 1));
  if (sync) CUDA_CHECK(cudaStreamSynchronize(stream));
}

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_check_vec(N);

  // Column-major diagonally dominant matrix for cuBLAS
  for (int j = 0; j < N; ++j) {
    for (int i = 0; i < N; ++i) {
      A_vec[i + j * N] = random_value<data_type>();
    }
    A_vec[j + j * N] = N;
  }
  for (int i = 0; i < N; ++i) {
    b_vec[i] = random_value<data_type>();
  }

  matrix_type A_save(A_save_vec.data(), N, N);
  vector_type x(x_vec.data(), N);
  vector_type b_check(b_check_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;

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

  data_type *d_A = nullptr, *d_b = nullptr, *d_x = nullptr;
  data_type *d_v = nullptr, *d_w = 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_v), sizeof(data_type) * N);
  cudaMalloc(reinterpret_cast<void **>(&d_w), sizeof(data_type) * N);

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

  Timer timer;

  for (int iRepeat = 0; iRepeat < N_REPEAT; ++iRepeat) {
    // QR modifies both A and b, re-copy each iteration
    cudaMemcpyAsync(d_A, A_save_vec.data(), sizeof(data_type) * N * N, cudaMemcpyHostToDevice, stream);
    cudaMemcpyAsync(d_b, b_save_vec.data(), sizeof(data_type) * N,     cudaMemcpyHostToDevice, stream);

    cudaStreamSynchronize(stream);
    timer.start();

    qr_factor(cublasH, N, d_A, N, d_b, d_v, d_w, syncFlag);

    qr_solve(cublasH, N, d_A, N, 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) * N, cudaMemcpyDeviceToHost, stream);
  cudaStreamSynchronize(stream);

  // Verify: A_save * x vs b_save on host
  la::matrix_vector_product(std::execution::seq, A_save, x, b_check);

#ifdef _VERBOSE
  //print_matrix(N, 1, b_save_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_save_vec[i] - b_check_vec[i]));
  }

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

  cudaFree(d_A);
  cudaFree(d_b);
  cudaFree(d_x);
  cudaFree(d_v);
  cudaFree(d_w);

  cudaStreamDestroy(stream);
  cublasDestroy(cublasH);
}
