#include <iostream>
#include <vector>
#include <cmath>
#include <mdspan>
#include <linalg>
#include <iomanip>
#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>>;

#ifdef __NVCOMPILER_STDPAR_GPU
__global__ void householder_u1(const data_type *x,
                               data_type x_norm, data_type *u1_out)
{
  if (threadIdx.x == 0 && blockIdx.x == 0) {
    data_type x0 = x[0];
    data_type alpha = (x0 > 0) ? -x_norm : x_norm;
    *u1_out = x0 - alpha;
  }
}

__global__ void set_value_kernel(data_type *p, data_type val)
{
  if (threadIdx.x == 0 && blockIdx.x == 0) {
    *p = val;
  }
}
//
data_type *u1_ptr = nullptr;
#endif

// Work arrays for qr_factor, allocated once to avoid per-call H2D transfers
std::vector<data_type> v_vec, w_vec, vwt_vec;

/**
 * QR Decomposition using Householder Reflections.
 * Transforms A into R (upper triangular) and b into Q^T * b in-place.
 */
void qr_factor(auto A, auto b) {
  const size_t rows = A.extent(0);
  const size_t cols = A.extent(1);

  for (size_t k = 0; k < cols && k < rows - 1; ++k) {
    size_t m_sub = rows - k;
    size_t n_sub = cols - k;

    // 1. Extract sub-column x
    auto x = std::submdspan(A, std::pair{k, rows}, k);

    // 2. Compute Householder vector v and tau
    data_type x_norm = la::vector_two_norm(exec, x);
    if (x_norm == 0) continue;

    auto v = std::mdspan(v_vec.data(), m_sub);

#ifndef __NVCOMPILER_STDPAR_GPU
    data_type alpha = (x[0] > 0) ? -x_norm : x_norm;
    data_type u1 = x[0] - alpha;
#else
    householder_u1<<<1, 1>>>(x.data_handle(), x_norm, u1_ptr);
    cudaDeviceSynchronize();
    data_type u1 = *u1_ptr;
    data_type alpha = (u1 > 0) ? -x_norm : x_norm;
#endif

    data_type tau = -u1 / alpha;
    la::copy(exec, x, v);
    if (u1 != 0) {
      auto v_tail = std::submdspan(v, std::pair{1, v.extent(0)});
      la::scale(exec, 1.0 / u1, v_tail);
    }

#ifndef __NVCOMPILER_STDPAR_GPU
    v[0] = 1.0;
#else
    set_value_kernel<<<1, 1>>>(v.data_handle(), 1.0);
#endif

    // 3. Apply H = I - tau*v*v^T to A[k:rows, k:cols]
    auto A_sub = std::submdspan(A, std::pair{k, rows}, std::pair{k, cols});

    // w = A_sub^T * v
    auto w = std::mdspan(w_vec.data(), n_sub);
    la::matrix_vector_product(exec, la::transposed(A_sub), v, w);

    // A_sub = A_sub - tau * v * w^T
    auto v_col = std::mdspan(v_vec.data(), m_sub, 1);
    auto w_row = std::mdspan(w_vec.data(), 1, n_sub);
    auto vwt = std::mdspan(vwt_vec.data(), m_sub, n_sub);
    la::matrix_product(exec, la::scaled(-tau, v_col), w_row, vwt);
    la::add(exec, vwt, A_sub, A_sub);
    // TODO: replace above with matrix_rank_1_update when the back-ends are implemented
    //la::matrix_rank_1_update(exec, la::scaled(-tau, v), w, A_sub, A_sub);

    // 4. Apply H to b: b_sub = b_sub - tau * v * (v^T * b_sub)
    auto b_sub = std::submdspan(b, std::pair{k, rows});
    data_type v_dot_b = la::dot(exec, v, b_sub);
    la::add(exec, la::scaled(-tau * v_dot_b, v), b_sub, b_sub);
  }
}

void qr_solve(auto A, auto b, auto x) {
  la::triangular_matrix_vector_solve(std::execution::par, A,
    la::upper_triangle, la::explicit_diagonal, b, 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);

  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] = random_value<data_type>();
    }
  }
  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);
  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 );

  // Allocate work arrays for qr_factor once upfront
#ifdef __NVCOMPILER_STDPAR_GPU
  cudaMallocHost(&u1_ptr, sizeof(data_type));
#endif
  v_vec.resize(N);
  w_vec.resize(N);
  vwt_vec.resize(N * N);

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

    qr_factor(A, b);

    qr_solve(A, b, x);

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

#ifdef __NVCOMPILER_STDPAR_GPU
  cudaFreeHost(u1_ptr);
#endif

  // Verify: A_save * x vs b_save
  la::matrix_vector_product(std::execution::par, A_save, x, b_inverted);

#ifdef _VERBOSE
  print_matrix(N, 1, b_save_vec.data(), N);
  print_matrix(N, 1, b_inverted_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_inverted_vec[i]));
  }

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